runtime/agent.py

"""agent.py - minimal tool-calling agent loop over the orchestrator.

Demonstrates: model router -> decide tool -> execute (Exa/Firecrawl/Mapbox/etc.)
-> feed result back -> final answer. Pluggable; production agents should use
LangGraph/Agno/CrewAI with checkpointing + tracing.
"""
from __future__ import annotations
import json
from typing import Callable

from orchestrator import Orchestrator

# Tool registry: name -> callable. Wire more from registry/tools.yaml.
TOOLS: dict[str, Callable] = {}


def register_tool(name: str):
    def deco(fn):
        TOOLS[name] = fn
        return fn
    return deco


@register_tool("exa.search")
def _exa_search(query: str, num_results: int = 5):
    from clients.python import exa
    return exa.search(query, num_results=num_results)


@register_tool("firecrawl.scrape")
def _fc_scrape(url: str):
    from clients.python import firecrawl
    return firecrawl.scrape(url)


@register_tool("mapbox.geocode")
def _mb_geocode(query: str):
    from clients.python import mapbox
    return mapbox.geocode(query)


SYSTEM = (
    "You are a toolbox agent. You have tools: "
    + ", ".join(TOOLS)
    + ". Use them when needed, then answer."
)


def run(prompt: str, max_steps: int = 4, orch: Orchestrator | None = None) -> str:
    orch = orch or Orchestrator()
    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}]
    for _ in range(max_steps):
        resp = orch.chat(messages)
        content = (resp.get("choices") or [{}])[0].get("message", {})
        tool_calls = content.get("tool_calls") or []
        if not tool_calls:
            return content.get("content", json.dumps(resp))
        for call in tool_calls:
            fn = call.get("function", {})
            name, args = fn.get("name"), json.loads(fn.get("arguments") or "{}")
            result = TOOLS.get(name, lambda **_: {"error": "unknown_tool"})(**args)
            messages.append({
                "role": "tool",
                "tool_call_id": call.get("id"),
                "content": json.dumps(result)[:4000],
            })
        messages.append(content)
    return messages[-1].get("content", "no_answer")


if __name__ == "__main__":
    import sys
    print(run(sys.argv[1] if len(sys.argv) > 1 else "What is the latest headline about NVIDIA?"))