"""agent.py: the same tool-using agent from M19, now INSTRUMENTED with the tracer.

It is the hand-rolled ReAct loop (ask the model, run a tool if asked, feed the result back,
repeat), except every model call and every tool call is recorded as a span in a Trace. The
function returns BOTH the answer and the trace, so you can see exactly what happened.

Run (venv active, key in .env):
    python agent.py
"""

import os
from dotenv import load_dotenv
import anthropic
from tracer import Trace

load_dotenv()
MODEL = "claude-opus-4-8"

TOOLS = [{
    "name": "multiply",
    "description": "Multiply two integers and return the product.",
    "input_schema": {
        "type": "object",
        "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
        "required": ["a", "b"],
    },
}]


def multiply(a, b):
    return a * b


def run(task, client=None, trace=None):
    """Run the agent. Returns (final_answer_text, trace)."""
    client = client or anthropic.Anthropic()
    trace = trace or Trace(label=f"task: {task[:40]}")
    messages = [{"role": "user", "content": task}]

    while True:
        span = trace.record("model", MODEL, {"messages": len(messages)})
        try:
            resp = client.messages.create(model=MODEL, max_tokens=1024, tools=TOOLS, messages=messages)
        except Exception as e:                              # record failures too: that IS observability
            span.finish(output=str(e), status="error")
            raise
        tokens = getattr(getattr(resp, "usage", None), "output_tokens", 0) or 0
        span.finish(output=resp.stop_reason, tokens=tokens)
        messages.append({"role": "assistant", "content": resp.content})

        if resp.stop_reason == "tool_use":
            results = []
            for block in resp.content:
                if getattr(block, "type", None) == "tool_use":
                    tspan = trace.record("tool", block.name, block.input)
                    try:
                        out = multiply(**block.input)
                        tspan.finish(output=out)
                    except Exception as e:
                        tspan.finish(output=str(e), status="error")
                        out = f"error: {e}"
                    results.append({"type": "tool_result", "tool_use_id": block.id, "content": str(out)})
            messages.append({"role": "user", "content": results})
            continue

        answer = "".join(b.text for b in resp.content if getattr(b, "type", None) == "text")
        return answer, trace


if __name__ == "__main__":
    answer, trace = run("What is 23 times 17? Use the multiply tool.")
    trace.print_tree()
    print("\nANSWER:", answer)