agent.py

# agent.py — defines what the agent is and how to run it.
# No infrastructure here: no FastAPI, no MCP spawning, no HTTP.
# main.py handles all of that and calls into this module.

import json
from agents import Agent, Runner

# The Ollama model tag to use. Changing this is the only thing needed
# to swap models across the whole application.
MODEL = "gemma4:12b"

TRACE_WIDTH = 60


def create_agent(*mcp_servers):
    # Agent is a stateless config object — it describes the agent but doesn't
    # run anything. Runner.run() is what actually executes it.
    # mcp_servers tells the SDK which tool servers are available; it handles
    # discovery and tool-call routing automatically.
    return Agent(
        name="assistant",
        model=MODEL,
        instructions=(
            "You are a helpful assistant. "
            "Use the web_search tool whenever you need current information."
        ),
        mcp_servers=list(mcp_servers),
    )


def _fmt_content(content):
    """Extract readable text from a content field (string or list of blocks)."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict):
                parts.append(block.get("text", repr(block)))
            else:
                parts.append(repr(block))
        return "\n".join(parts)
    return repr(content)


def _print_trace(input, result):
    """Print the full conversation trace in a human-readable format."""
    sep = "=" * TRACE_WIDTH
    thin = "-" * TRACE_WIDTH

    print(f"\n{sep}")
    print("  AGENT TRACE")
    print(sep)

    # --- Input ---
    print("\n  INPUT")
    print(thin)
    if isinstance(input, str):
        print(f"  user: {input}")
    else:
        for msg in input:
            role = msg.get("role", "?")
            print(f"  {role}: {msg.get('content', '')}")

    # --- Steps ---
    for item in result.new_items:
        name = type(item).__name__
        raw = item.raw_item
        data = raw.model_dump() if hasattr(raw, "model_dump") else {}

        print(f"\n  {thin}")

        if name == "MessageOutputItem":
            print("  ASSISTANT")
            print(thin)
            print(f"  {_fmt_content(data.get('content', ''))}")

        elif name == "ToolCallItem":
            print(f"  TOOL CALL  →  {data.get('name', '?')}")
            print(thin)
            try:
                args = json.loads(data.get("arguments", "{}"))
                for k, v in args.items():
                    print(f"  {k}: {v}")
            except json.JSONDecodeError:
                print(f"  {data.get('arguments', '')}")

        elif name == "ToolCallOutputItem":
            print("  TOOL OUTPUT")
            print(thin)
            # Output may be on the item itself or nested in raw
            output = getattr(item, "output", None) or str(data)
            if len(output) > 600:
                print(f"  {output[:600]}")
                print(f"  ... [{len(output) - 600} chars truncated]")
            else:
                print(f"  {output}")

        else:
            # Catch-all for any other item types (handoffs, guardrails, etc.)
            print(f"  {name.upper()}")
            print(thin)
            print(f"  {repr(raw)}")

    print(f"\n{sep}\n")


async def run(agent, input):
    # Runner.run() is the agentic loop: it sends the input to the model,
    # executes any tool calls the model requests, feeds results back, and
    # repeats until the model returns a plain text response.
    # result.final_output is that final text string.
    result = await Runner.run(agent, input)

    # Print the full trace so we can see exactly what the model did:
    # input, any tool calls made and their outputs, and the final reply.
    _print_trace(input, result)

    return result.final_output