Middleware

Intercept and modify model and tool calls with before/after hooks and wrap decorators

Middleware lets your agent intercept every LLM call and every tool call — log them, modify inputs, retry on failure, or block execution entirely.

from praisonaiagents import Agent
from praisonaiagents.hooks import before_model

@before_model
def log_request(event):
    print(f"Model call: {len(event.messages)} messages")

agent = Agent(name="Guarded", instructions="Answer helpfully.")
agent.start("Explain middleware in one paragraph.")

The user adds hooks or wrap decorators; middleware runs on every model and tool call in the agent loop.

Quick Start

1. Add a logging middleware

   from praisonaiagents import Agent
   from praisonaiagents.hooks import before_model, after_model
   
   @before_model
   def log_request(request):
       print(f"Calling model: {request.model}")
       return request
   
   @after_model
   def log_response(response):
       print(f"Model replied: {response.content[:80]}")
       return response
   
   agent = Agent(
       name="Assistant",
       instructions="Answer questions helpfully.",
       hooks=[log_request, log_response]
   )
   
   result = agent.start("What is 2 + 2?")

2. Wrap calls for retry logic

   from praisonaiagents import Agent
   from praisonaiagents.hooks import wrap_model_call, wrap_tool_call
   
   @wrap_model_call
   def retry_model(request, call_next):
       for attempt in range(3):
           try:
               return call_next(request)
           except Exception as exc:
               if attempt == 2:
                   raise
       return call_next(request)
   
   @wrap_tool_call
   def retry_tool(request, call_next):
       for attempt in range(3):
           try:
               return call_next(request)
           except Exception as exc:
               if attempt == 2:
                   raise
   
   agent = Agent(
       name="RobustBot",
       instructions="Answer questions with retry protection.",
       hooks=[retry_model, retry_tool]
   )
   
   result = agent.start("Search for today's news.")

3. Async parity

   The same hooks gate async tool calls — write them once and they fire for both agent.start(...) and agent.astart(...).

   import asyncio
   from praisonaiagents import Agent
   from praisonaiagents.hooks import before_tool, ToolRequest
   
   @before_tool
   def audit(request: ToolRequest):
       print(f"[audit] {request.tool_name}({request.arguments})")
       return request
   
   agent = Agent(name="AsyncAudited", instructions="Answer with tools.", hooks=[audit])
   
   async def main():
       await agent.astart("Look up the weather in Paris.")
   
   asyncio.run(main())

   Under the hood the synchronous wrap_tool_call chain runs in a thread-pool worker and the actual async tool execution is scheduled back on the event loop — so you write hooks the same way for sync and async.

How It Works

Execution order for a model call:

1. All before_model hooks run in registration order
2. wrap_model_call middleware chain wraps the actual LLM call
3. The real LLM call executes
4. All after_model hooks run in reverse registration order

Under the hood:

• _chat_completion_with_retry / _achat_completion_with_retry build a ModelRequest (with InvocationContext), pass it through MiddlewareManager.execute_model_call, and unwrap the returned ModelResponse back to the agent’s expected return shape.
• The retry/backoff logic is the final handler of the middleware chain — so if a hook wraps call_next, it wraps the retry loop.
• before_model / wrap_model_call mutations to request.messages, request.temperature, and request.tools are threaded into the retry core.
• after_model / wrap_model_call edits to response.content win over the raw provider payload; unchanged content passes the original object through untouched.
• The async path runs the sync middleware chain in an executor and bridges back to the async retry core via asyncio.run_coroutine_threadsafe — hooks are written the same way for sync and async.

The same order applies for tool calls with before_tool, wrap_tool_call, and after_tool — including async tool calls:

Configuration Options

Pass middleware as a flat list to Agent(hooks=[...]). Mix decorator types freely.

agent = Agent(
    name="MyAgent",
    instructions="...",
    hooks=[log_request, retry_model, validate_tools]  # order matters
)

Data types available:

from praisonaiagents.hooks import (
    ModelRequest,    # messages, model, temperature, context, tools, extra
    ModelResponse,   # content, model, usage, context, tool_calls, extra
    ToolRequest,     # tool_name, arguments, context, extra
    ToolResponse,    # tool_name, result, error, context, extra
    InvocationContext,  # agent_id, run_id, session_id, tool_name, model_name, metadata
)

Common Patterns

Short-circuit the model with a cached response:

from praisonaiagents import Agent
from praisonaiagents.hooks import wrap_model_call, ModelResponse

_cache = {}

@wrap_model_call
def cache_model(request, call_next):
    key = (request.model, tuple((m["role"], m["content"]) for m in request.messages))
    if key in _cache:
        return _cache[key]
    response = call_next(request)
    _cache[key] = response
    return response

agent = Agent(name="Cached", instructions="Answer briefly.", hooks=[cache_model])
agent.start("What is 2+2?")   # hits the LLM
agent.start("What is 2+2?")   # served from cache

Inject a system prompt on every call:

from praisonaiagents.hooks import before_model

@before_model
def add_safety_prompt(request):
    safety = {"role": "system", "content": "Always be safe and helpful."}
    if safety not in request.messages:
        request.messages.insert(0, safety)
    return request

Block a specific tool:

from praisonaiagents.hooks import before_tool, ToolRequest

@before_tool
def block_delete(request: ToolRequest):
    if request.tool_name == "delete_file":
        raise PermissionError("delete_file is disabled in this environment.")
    return request

Cache expensive tool results:

import hashlib, json
from praisonaiagents.hooks import wrap_tool_call

_cache = {}

@wrap_tool_call
def cache_results(request, call_next):
    key = hashlib.md5(json.dumps(request.arguments, sort_keys=True).encode()).hexdigest()
    if key in _cache:
        return _cache[key]
    result = call_next(request)
    _cache[key] = result
    return result

Gate a tool through an external policy engine:

from praisonaiagents.hooks import wrap_tool_call, ToolResponse

@wrap_tool_call
def guard_shell(request, call_next):
    if request.tool_name != "run_shell":
        return call_next(request)
    if is_command_allowed(request.arguments.get("command", "")):
        return call_next(request)
    return ToolResponse(
        tool_name=request.tool_name,
        result="Blocked by policy.",
        error="Blocked by policy.",
        context=request.context,
    )

See HOL Guard Tool Policy for a complete fail-closed integration with an external security engine.

Best Practices

from praisonaiagents import Agent
from praisonaiagents.hooks import before_tool
from praisonaiagents.session.context import get_session_context, set_session_context

@before_tool
def audit(req):
    return req

async def whoami():
    "current user"
    return get_session_context().user_id

agent = Agent(name="t", instructions="x", tools=[whoami], hooks=[audit])
set_session_context(user_id="alice", platform="telegram")
# whoami() returns "alice" even with the hook attached.

Related