← Docs

Adding a New Agent

Step-by-step: subclass `BaseAgent`, add to `AGENT_REGISTRY`, write an `AgentSpec` in `_catalog.py`, add fixtures to `mock_claude._FIXTURES`, optionally seed it into a workflow.

Decide which pattern you're using

Three patterns. Pick before you write code — it determines what reference template to copy.

  • •studio_consuming — agent consumes Studio (asks for a generated asset, waits for approval, transcodes, publishes). Reference: ~/projects/tiktok-army/tiktok_army/agents/content_producer.py.
  • •llm_pure — agent only calls Claude (no Studio, no asset). Reference: ~/projects/tiktok-army/tiktok_army/agents/comment_triage.py.
  • •data_monitoring — agent computes scores from provider data. May call Claude lightly for summaries. Reference: ~/projects/tiktok-army/tiktok_army/agents/account_health.py (skeleton with significant scaffolding).

If you're not sure, start with llm_pure. It's the simplest shape.

Step 1: Create the agent file

Create ~/projects/tiktok-army/tiktok_army/agents/<your_name>.py. Skeleton (llm_pure pattern):

"""Your Agent — one-line purpose."""

from __future__ import annotations

import json
import logging
from typing import Any
from uuid import UUID

from sqlalchemy import text

from tiktok_army.agents.base import AgentContext, AgentResult, BaseAgent
from tiktok_army.lib.claude import call_claude_cached
from tiktok_army.lib.db import session_for_workspace

logger = logging.getLogger(__name__)


_SYSTEM_PROMPT = """You are <agent role> for the brand "{brand_name}".

<Stable, voice-aware system prompt. KEEP DETERMINISTIC — no timestamps, no
per-request data, no UUIDs. Cache requires byte-identical system prompts.>

OUTPUT FORMAT (strict JSON, no markdown fence):
{{
  "field_one": "...",
  "field_two": 0
}}
"""


class YourAgent(BaseAgent):
    name = "your_agent"

    async def _execute(self, ctx: AgentContext) -> AgentResult:
        handle = ctx.input.get("handle")
        if not handle:
            raise ValueError("your_agent requires a handle")

        # 1. Load context (brand profile, posts, whatever's needed).
        brand_profile = await self._load_brand(ctx.workspace_id, ctx.brand_id)

        # 2. Build cached system prompt (per-brand only — deterministic).
        system_prompt = _SYSTEM_PROMPT.format(brand_name=brand_profile.get("name", "Unknown"))

        # 3. Call Claude. Per-request data goes in user_message, NEVER in system.
        result = await call_claude_cached(
            system_stable=system_prompt,
            user_message=f"Handle: @{handle}\n<per-request context here>",
            model="claude-haiku-4-5-20251001",
            agent_name=self.name,
            max_tokens=1024,
            temperature=0.3,
        )

        # 4. Parse the JSON response.
        try:
            parsed = json.loads(result.text.strip())
        except json.JSONDecodeError as exc:
            raise ValueError(f"agent returned non-JSON: {result.text[:200]}") from exc

        # 5. (Optional) persist to DB.
        # await self._save(ctx.workspace_id, parsed)

        return AgentResult(
            output=parsed,
            cost_usd=result.cost_usd,
            model_used=result.model,
        )

    async def _load_brand(self, workspace_id: UUID, brand_id: UUID | None) -> dict[str, Any]:
        if brand_id is None:
            return {"name": "Unknown"}
        async with session_for_workspace(workspace_id) as session:
            row = (
                await session.execute(
                    text("SELECT id, name FROM brands WHERE id = :id"),
                    {"id": brand_id},
                )
            ).mappings().first()
        return dict(row) if row else {"name": "Unknown"}

Critical things in this skeleton:

  • •The _SYSTEM_PROMPT is a module-level constant. Anything that varies per call goes in user_message. The cache requires the system prompt to be byte-identical across calls.
  • •BaseAgent does the lifecycle. You only override _execute.
  • •DB ops use session_for_workspace(workspace_id) so RLS is set.
  • •Errors raise; BaseAgent.run() catches and persists them.
  • •Return AgentResult with output, cost_usd, model_used.

For a studio_consuming agent, copy from content_producer.py and adapt — the flow is more involved (request_generation → wait_for_approval → fetch_asset_to_local → to_tiktok_vertical → tiktok_publisher.publish → INSERT post → publish event). For data_monitoring, copy from account_health.py — start with provider reads, do compute in Python, optionally call Claude at the end.

Step 2: Register in `AGENT_REGISTRY`

Edit ~/projects/tiktok-army/tiktok_army/agents/__init__.py. Two changes:

  1. Import: from tiktok_army.agents.your_agent import YourAgent.
  2. Add to the registry dict (alphabetized by current convention):
AGENT_REGISTRY: dict[str, type[BaseAgent]] = {
    AccountHealthAgent.name: AccountHealthAgent,
    ...
    YourAgent.name: YourAgent,
    ...
}
  1. Add to __all__.

Without this step, the orchestrator will fail with unknown agent: your_agent when a workflow tries to invoke it.

Step 3: Write an `AgentSpec` in `_catalog.py`

Edit ~/projects/tiktok-army/tiktok_army/agents/_catalog.py. Add an entry to the SPECS dict:

"your_agent": AgentSpec(
    name="your_agent",
    display_name="Your Agent",
    purpose="One-line description of what this does.",
    pattern="llm_pure",  # or studio_consuming / data_monitoring
    typical_model="claude-haiku-4-5-20251001",
    inputs=[_HANDLE_FIELD],  # handle is the standard input
    options=[
        AgentField(
            name="some_option",
            kind="integer",
            label="Some option",
            help_text="What this option does.",
            default=10,
            min_value=1,
            max_value=100,
        ),
    ],
    outputs=[
        AgentField(name="field_one", kind="string", label="Field one", help_text=""),
        AgentField(name="field_two", kind="integer", label="Field two 0-100", help_text=""),
    ],
    human_touchpoints=[],  # add HumanTouchpoint(...) entries if there's an approval gate
    workflows=[],  # populated when you add it to a workflow
),

This drives:

  • •The /agents/catalog page in the dashboard (rendered from SPECS).
  • •The brief intake form (auto-generates input controls from inputs + options where human_editable=True).
  • •The workflow editor's validation.

If you skip this step, the agent will run from a workflow but won't show up in the catalog or any UI form. The system isn't introspecting it at runtime to validate — but every other entry-point assumes the spec exists.

Step 4: Add fixtures to `mock_claude._FIXTURES`

Edit ~/projects/tiktok-army/tiktok_army/lib/mock_claude.py. Add an entry to _FIXTURES:

"your_agent": [
    _MockResponse(
        text=json.dumps(
            {
                "field_one": "example value",
                "field_two": 42,
            },
            indent=2,
        )
    ),
    # Add more fixtures if you want to cycle through variants based on prompt hash.
],

Without this, the mock will return _DEFAULT_FIXTURE (a generic {"status": "ok", "note": "mock response — no fixture for this agent yet"}), which is fine for development but breaks any downstream agent that expects specific fields in your output.

A fixture's text should match the JSON shape your _execute parses. The mock dispatches by agent_name (set via the agent_name kwarg you pass to call_claude_cached, which defaults to the contextvar set by BaseAgent.run()).

If your agent might be invoked from a workflow's synthesis step, you may also want to add to _SYNTHESIS_FIXTURES keyed by workflow slug — but only when you're seeding a new workflow, not just a new agent.

Step 5: Optionally add to a workflow

If your agent should be part of one of the seeded workflows (Profile Audit, Campaign Launch, Post-Launch Loop), edit ~/projects/tiktok-army/tiktok_army/orchestrator/definitions.py and add a WorkflowStepDef:

WorkflowStepDef(
    key="your_step",
    agent_name="your_agent",
    label="Run your agent",
    depends_on=["audit_health"],  # or whatever upstream key feeds it
    input_map={"handle": "brief.handle"},
    options={"some_option": 20},  # optional per-step override
),

If your step's output is consumed by a downstream step, that downstream step references it via input_map:

WorkflowStepDef(
    key="downstream_step",
    agent_name="something",
    depends_on=["your_step"],
    input_map={"thing_we_need": "your_step.output.field_one"},
),

If your step's output should be reviewed by a human before something downstream runs, insert an APPROVAL_NODE step between them with target_output_field set to whatever output field the human is reviewing. See Workflow Contract for details.

Step 6: Write a test

Tests live in ~/projects/tiktok-army/tests/. Existing patterns to copy:

  • •tests/test_account_health.py — full agent test with mock data + assertions on output structure.
  • •tests/test_audience_mapper.py — same shape.

A minimal test:

import pytest
from uuid import uuid4
from tiktok_army.agents import YourAgent
from tiktok_army.models import AgentTriggerType


@pytest.mark.asyncio  # not strictly needed — pytest.ini sets asyncio_mode=auto
async def test_your_agent_happy_path():
    agent = YourAgent()
    result = await agent.run(
        workspace_id=uuid4(),
        brand_id=None,
        trigger_type=AgentTriggerType.MANUAL,
        input_data={"handle": "lakucosmetics"},
    )
    assert result.output["field_one"]
    assert isinstance(result.output["field_two"], int)
    assert result.cost_usd > 0

Run with uv run pytest tests/test_your_agent.py. Mock mode is the default in tests (the config falls back to placeholder secrets when CLAUDE_MODE=mock), so no API key is needed. See Testing for the full picture, including which tests today actually hit DB/Pub/Sub and how to skip those.

Checklist

  • •[ ] tiktok_army/agents/your_agent.py exists with a BaseAgent subclass and _execute.
  • •[ ] name class attribute matches the slug used in _FIXTURES, SPECS, and AGENT_REGISTRY.
  • •[ ] tiktok_army/agents/__init__.py imports the class and adds it to AGENT_REGISTRY.
  • •[ ] tiktok_army/agents/_catalog.py:SPECS has an AgentSpec entry.
  • •[ ] tiktok_army/lib/mock_claude.py:_FIXTURES has at least one fixture.
  • •[ ] If part of a workflow, tiktok_army/orchestrator/definitions.py has a WorkflowStepDef.
  • •[ ] A test in tests/ exercises the happy path.
  • •[ ] uv run ruff check . passes.
  • •[ ] uv run mypy tiktok_army passes (strict mode — typing must be complete).

Common gotchas

  • •System prompt mutates per call. Most common cache-killer. If your _SYSTEM_PROMPT includes anything that varies between requests (a date, a request ID, a counter), the cache misses every time. Move that data to user_message.
  • •Forgot RLS scope. Direct DB access without session_for_workspace returns nothing in dev (RLS denies) or — worse — leaks rows in real mode. Always use the helper.
  • •Forgot to bump cost_usd. If your agent makes multiple Claude calls, sum their cost_usd. The framework doesn't add them automatically.
  • •Forgot the spec. Agent works in tests but doesn't show up in the dashboard catalog. Always add the AgentSpec.
  • •Mock fixture has wrong JSON shape. Tests pass (the mock returns whatever you wrote), but a downstream agent in a real workflow chokes when the field it expected is missing. Make sure the fixture matches _catalog.py:outputs.