Custom Agent Runtimes Basics

python -m venv .venv && source .venv/bin/activate
pip install openai  # OpenAI-compatible client works for many gateways
export OPENAI_API_KEY="sk-..."   # or gateway key
export OPENAI_BASE_URL="${OPENAI_BASE_URL:-}"  # optional gateway

from typing import Any
 
 
def get_order_status(order_id: str) -> dict[str, Any]:
    """Return a demo order status payload."""
    demo = {
        "A-100": {"status": "shipped", "eta_days": 2},
        "A-200": {"status": "processing", "eta_days": 5},
    }
    return demo.get(order_id, {"status": "not_found", "order_id": order_id})
 
 
def add(a: float, b: float) -> float:
    """Add two numbers for deterministic demos."""
    return a + b

from typing import Any, Callable
 
TOOLS_SPEC: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "get_order_status",
            "description": "Look up order status by id.",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string"},
                },
                "required": ["order_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "Add two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "number"},
                    "b": {"type": "number"},
                },
                "required": ["a", "b"],
            },
        },
    },
]
 
DISPATCH: dict[str, Callable[..., Any]] = {
    "get_order_status": get_order_status,
    "add": add,
}
 
 
def run_tool(name: str, arguments: dict[str, Any]) -> Any:
    fn = DISPATCH.get(name)
    if fn is None:
        raise KeyError(f"unknown_tool:{name}")
    return fn(**arguments)

import json
import os
from openai import OpenAI
 
client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ.get("OPENAI_BASE_URL") or None,
)
 
MODEL = os.environ.get("AGENT_MODEL", "gpt-4.1-mini")  # verify at build
 
 
def chat_once(messages: list[dict]) -> object:
    return client.chat.completions.create(
        model=MODEL,
        messages=messages,
        tools=TOOLS_SPEC,
        tool_choice="auto",
    )

import json
from typing import Any
 
 
def run_agent(user_text: str, max_turns: int = 5) -> dict[str, Any]:
    messages: list[dict[str, Any]] = [
        {
            "role": "system",
            "content": (
                "You are a careful assistant. "
                "Use tools when needed. Prefer short final answers."
            ),
        },
        {"role": "user", "content": user_text},
    ]
 
    for turn in range(1, max_turns + 1):
        resp = chat_once(messages)
        choice = resp.choices[0]
        msg = choice.message
 
        # Record assistant message (including tool_calls if any)
        assistant_msg: dict[str, Any] = {
            "role": "assistant",
            "content": msg.content or "",
        }
        if msg.tool_calls:
            assistant_msg["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in msg.tool_calls
            ]
        messages.append(assistant_msg)
 
        if not msg.tool_calls:
            return {
                "ok": True,
                "stop_reason": "final_answer",
                "turns": turn,
                "output": msg.content or "",
                "messages": messages,
            }
 
        for tc in msg.tool_calls:
            name = tc.function.name
            try:
                args = json.loads(tc.function.arguments or "{}")
                result = run_tool(name, args)
                content = json.dumps(result)
            except Exception as exc:  # production: narrower + metrics
                content = json.dumps({"error": str(exc), "tool": name})
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": content,
                }
            )
 
    return {
        "ok": False,
        "stop_reason": "max_turns",
        "turns": max_turns,
        "output": None,
        "messages": messages,
    }
 
 
if __name__ == "__main__":
    print(run_agent("What is the status of order A-100?"))

from enum import Enum
 
 
class StopReason(str, Enum):
    FINAL_ANSWER = "final_answer"
    MAX_TURNS = "max_turns"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"
    TOOL_POLICY = "tool_policy"
    MODEL_ERROR = "model_error"
 
 
def summarize_run(result: dict) -> dict:
    return {
        "stop_reason": result["stop_reason"],
        "turns": result["turns"],
        "ok": result["ok"],
        "output_preview": (result.get("output") or "")[:200],
    }

import time
from typing import Any, Callable
 
 
def run_agent_with_deadline(
    user_text: str,
    max_turns: int = 5,
    deadline_s: float = 30.0,
    runner: Callable[..., dict[str, Any]] = run_agent,
) -> dict[str, Any]:
    started = time.monotonic()
    # Thin wrapper pattern: production often checks deadline inside each turn.
    result = runner(user_text, max_turns=max_turns)
    if time.monotonic() - started > deadline_s and result.get("ok"):
        result = {
            **result,
            "ok": False,
            "stop_reason": "timeout",
        }
    return result

import json
 
 
def test_get_order_status_dispatches():
    out = run_tool("get_order_status", {"order_id": "A-100"})
    assert out["status"] == "shipped"
 
 
def test_unknown_tool_raises():
    try:
        run_tool("drop_database", {})
        assert False, "expected KeyError"
    except KeyError as exc:
        assert "unknown_tool" in str(exc)
 
 
def test_tool_error_becomes_json_observation():
    # Simulate the loop's error packaging
    try:
        run_tool("add", {"a": 1})  # missing b
    except TypeError as exc:
        payload = json.dumps({"error": str(exc), "tool": "add"})
        assert "error" in json.loads(payload)

import time
from typing import Any
 
 
def run_agent_traced(user_text: str, max_turns: int = 5) -> dict[str, Any]:
    t0 = time.monotonic()
    events: list[dict[str, Any]] = []
    result = run_agent(user_text, max_turns=max_turns)
 
    for m in result.get("messages") or []:
        if m.get("role") == "assistant" and m.get("tool_calls"):
            for tc in m["tool_calls"]:
                events.append(
                    {
                        "type": "tool_call",
                        "name": tc["function"]["name"],
                        # redacted: do not log raw args in prod by default
                    }
                )
        if m.get("role") == "tool":
            events.append({"type": "tool_result", "tool_call_id": m.get("tool_call_id")})
 
    return {
        **summarize_run(result),
        "latency_ms": int((time.monotonic() - t0) * 1000),
        "events": events,
        "output": result.get("output"),
    }