Decorate LLM and tool calls

Learn how to decorate LLM and tool calls with Agent Control.

Add @control() to any function to enforce server-managed safety controls on its inputs and outputs. This guide walks you through:

• Setting up your environment.
• Creating two agent controls—block-ssn-output and block-dangerous-sql— to block social security numbers and dangerous SQL queries, respectively.
• Decorating an LLM call that asks “What is the capital of France?” and “DROP TABLE users”.
• Returning the answer to the former and blocking the potentially dangerous SQL injection.

Prerequisites

• Python 3.12+, uv, Docker

1. Start the server

   git clone https://github.com/agentcontrol/agent-control.git
   cd agent-control
   make sync
   
   cd server && docker-compose up -d && cd ..
   make server-alembic-upgrade
   make server-run                    # leave running — use a new terminal below
   
2. Create your project

   mkdir my-agent && cd my-agent
   uv init
   uv add agent-control-sdk anthropic  # swap anthropic for openai, etc.
   
3. Create setup_controls.py

   cd my-agent
   touch setup_controls.py
   

   This script registers your agent, creates two controls, and associates them directly to the agent. Copy and paste the following into setup_controls.py.

   """One-time setup: create an agent and two controls, then attach them to the agent."""
   import asyncio
   from agent_control import AgentControlClient, controls
   
   AGENT_NAME = "my-agent"
   SERVER_URL = "http://localhost:8000"
   
   
   async def main():
       async with AgentControlClient(base_url=SERVER_URL) as client:
           # 1. Register the agent
           resp = await client.http_client.post(
               "/api/v1/agents/initAgent",
               json={
                   "agent": {
                       "agent_name": AGENT_NAME,
                       "agent_description": "Demo agent",
                   },
                   "steps": [],
               },
           )
           resp.raise_for_status()
   
           # 2. Create controls
           ssn = await controls.create_control(client, "block-ssn-output", data={
               "enabled": True,
               "execution": "server",
               "scope": {"step_types": ["llm"], "stages": ["post"]},
               "condition": {
                   "selector": {"path": "output"},
                   "evaluator": {
                       "name": "regex",
                       "config": {"pattern": r"\b\d{3}-\d{2}-\d{4}\b"},
                   },
               },
               "action": {"decision": "deny"},
           })
   
           sql = await controls.create_control(client, "block-dangerous-sql", data={
               "enabled": True,
               "execution": "server",
               "scope": {"step_types": ["llm"], "stages": ["pre"]},
               "condition": {
                   "selector": {"path": "input"},
                   "evaluator": {
                       "name": "list",
                       "config": {
                           "values": ["DROP", "DELETE", "TRUNCATE"],
                           "logic": "any",
                           "match_on": "match",
                           "match_mode": "contains",
                           "case_sensitive": False,
                       },
                   },
               },
               "action": {"decision": "deny"},
           })
   
           # 3. Associate controls directly to the agent
           await client.http_client.post(f"/api/v1/agents/{AGENT_NAME}/controls/{ssn['control_id']}")
           await client.http_client.post(f"/api/v1/agents/{AGENT_NAME}/controls/{sql['control_id']}")
   
           print(f"Done — agent '{AGENT_NAME}' ready with 2 controls")
   
   
   asyncio.run(main())
   

   Agent and control names must be unique. If you get a 409 conflict, pick new names or reset the database.

   uv run python setup_controls.py
   
4. Create main.py

   Three things to add to a normal LLM script: import, init(), @control().

   import asyncio
   import anthropic
   import agent_control
   from agent_control import control, ControlViolationError
   
   AGENT_NAME = "my-agent"
   
   # Initialize once at startup
   
   agent_control.init(
       agent_name=AGENT_NAME,
       agent_description="Demo agent",
       server_url="http://localhost:8000",
   )
   
   client = anthropic.Anthropic()          # uses ANTHROPIC_API_KEY env var
   
   @control()
   async def chat(message: str) -> str:
       response = client.messages.create(
           model="claude-sonnet-4-20250514",
           max_tokens=1024,
           messages=[{"role": "user", "content": message}],
       )
       return response.content[0].text
   
   async def main():
       for prompt in ["What is the capital of France?", "DROP TABLE users"]:
           try:
               result = await chat(prompt)
               print(f"✅ {prompt} → {result[:60]}")
           except ControlViolationError as e:
               print(f"🚫 {prompt} → blocked by {e.control_name}")
   
   asyncio.run(main())
   
5. Run it

   export ANTHROPIC_API_KEY="your-key"
   uv run python main.py
   

   ✅ What is the capital of France? → The capital of France is Paris.
   🚫 DROP TABLE users → blocked by block-dangerous-sql
   

---

How it works

1. Pre-stage — before the function runs, the decorator sends its input to the server. Controls scoped to "pre" evaluate it. If denied, the function never executes.
2. Execution — the LLM call runs normally.
3. Post-stage — after the function returns, the decorator sends the output to the server. Controls scoped to "post" evaluate it. If denied, the output is blocked.

Decorate tool calls the same way

@control()
async def execute_query(query: str) -> str:
    return await db.run(query)

Any LLM SDK works

@control() wraps the function, not a specific provider:

@control()
async def openai_chat(message: str) -> str:
    return openai_client.chat.completions.create(...).choices[0].message.content

Key points

• Works on both async and sync functions.
• Controls live on the server — update them without redeploying your agent.
• Fail-safe: if the server is unreachable, the call is blocked, not silently allowed.

---

Troubleshooting

409 Conflict — name already exists

Agent and control names are unique. Re-running setup_controls.py against a database that already has those names will return a 409 Conflict.

Option A — pick new names. Change the name strings in setup_controls.py (and update AGENT_NAME in main.py to match).

Option B — reset the database. From the repo root, stop the server, wipe the Docker volume, and re-run migrations:

# stop the running server (Ctrl-C), then:

cd server
docker compose down -v          # removes the postgres volume
docker compose up -d            # recreates a fresh database
make alembic-upgrade            # re-applies migrations
cd ..
make server-run                 # restart the server

Then re-run setup_controls.py.

422 Unprocessable Entity on initAgent

The /initAgent payload must include agent_name inside the agent object. Double-check your setup_controls.py sends it:

"agent": {
    "agent_name": AGENT_NAME,
}