Hooks & Middleware¶

AgentHooks provides a lightweight callback system for injecting behavior before and after LLM calls and tool calls. No subclassing, no abstract base classes — just callback lists with decorator registration.

Quick start¶

from shipit_agent import Agent, AgentHooks
from shipit_agent.llms import OpenAIChatLLM

hooks = AgentHooks()

@hooks.on_before_llm
def log_llm_call(messages, tools):
    print(f"Calling LLM with {len(messages)} messages, {len(tools)} tools")

@hooks.on_after_llm
def track_tokens(response):
    usage = response.usage
    if usage:
        print(f"Tokens: {usage.get('total_tokens', 0)}")

@hooks.on_before_tool
def log_tool_start(name, arguments):
    print(f"Running {name}...")

@hooks.on_after_tool
def log_tool_end(name, result):
    print(f"{name} returned {len(result.output)} chars")

agent = Agent.with_builtins(
    llm=OpenAIChatLLM(model="gpt-4o-mini"),
    hooks=hooks,
)

result = agent.run("What is the weather in Tokyo?")

Hook types¶

Hook	Signature	When it fires
`before_llm`	`fn(messages: list, tools: list)`	Before each LLM completion call
`after_llm`	`fn(response: LLMResponse)`	After each LLM completion returns
`before_tool`	`fn(name: str, arguments: dict)`	Before a tool is executed
`after_tool`	`fn(name: str, result: ToolResult)`	After a tool returns (success or error)

Registration¶

Two ways to register hooks:

Decorator styleAppend style

hooks = AgentHooks()

@hooks.on_before_llm
def my_hook(messages, tools):
    ...

hooks = AgentHooks()
hooks.before_llm.append(lambda msgs, tools: print("calling LLM"))

Both are equivalent. The decorator returns the original function, so you can still call it directly.

Common patterns¶

Cost tracking¶

total_cost = {"tokens": 0}

@hooks.on_after_llm
def accumulate(response):
    total_cost["tokens"] += response.usage.get("total_tokens", 0)

agent.run("Do something complex")
print(f"Total tokens used: {total_cost['tokens']}")

Rate limiting¶

import time

last_call = {"time": 0.0}

@hooks.on_before_llm
def rate_limit(messages, tools):
    elapsed = time.time() - last_call["time"]
    if elapsed < 1.0:
        time.sleep(1.0 - elapsed)
    last_call["time"] = time.time()

Content filtering¶

@hooks.on_after_tool
def filter_pii(name, result):
    if "email" in result.output.lower():
        print(f"Warning: {name} output may contain PII")

Guardrails¶

BLOCKED_TOOLS = {"code_execution", "workspace_files"}

@hooks.on_before_tool
def block_dangerous_tools(name, arguments):
    if name in BLOCKED_TOOLS:
        raise PermissionError(f"Tool {name} is blocked by policy")

Via the profile builder¶

from shipit_agent import AgentProfileBuilder, AgentHooks

hooks = AgentHooks()
hooks.before_llm.append(my_logger)

profile = (
    AgentProfileBuilder("monitored-agent")
    .hooks(hooks)
    .build_profile()
)

Works with async too¶

AgentHooks works identically with AsyncAgentRuntime. The hook callbacks themselves are synchronous — the async runtime calls them inline between awaits.