Observability for agent runs

import time
import uuid
from contextlib import contextmanager

TRACE = []
_stack = []
_t0 = time.perf_counter()


@contextmanager
def span(name, **attributes):
    span_id = str(uuid.uuid4())[:8]
    start = time.perf_counter()
    depth = len(_stack)
    _stack.append(span_id)
    try:
        yield
    finally:
        _stack.pop()
        duration_ms = (time.perf_counter() - start) * 1000
        TRACE.append({"name": name, "start_ms": (start - _t0) * 1000,
                       "duration_ms": duration_ms, "depth": depth, **attributes})


def print_trace():
    TRACE.sort(key=lambda r: r["start_ms"])
    for row in TRACE:
        attrs = {k: v for k, v in row.items()
                 if k not in ("name", "start_ms", "duration_ms", "depth")}
        indent = "  " * row["depth"]
        label = f"{indent}{row['name']}"
        print(f"{label:<32}{attrs}")


def calculator_tool(expression):
    return eval(expression, {"__builtins__": {}})


def search_tool(query):
    facts = {"who invented the calculator": "Blaise Pascal built an early mechanical calculator in 1642."}
    return facts.get(query.lower(), "no result found")


def fake_agent_policy(step, scratchpad):
    """
    Stands in for an LLM deciding the next action given the scratchpad so
    far. A real agent asks a model "what should I do next?" and parses its
    reply; this fixed, deterministic script plays that role so the whole
    example runs with no API key.
    """
    plan = [
        ("calculator", "240 * 0.15"),
        ("search", "who invented the calculator"),
        ("final_answer", None),
    ]
    return plan[step]


def run_agent(max_steps=5):
    scratchpad = []
    with span("agent_run", max_steps=max_steps):
        for step in range(max_steps):
            with span("think", step=step):
                tool, arg = fake_agent_policy(step, scratchpad)

            if tool == "final_answer":
                with span("final_answer"):
                    answer = (f"15% of 240 is {scratchpad[0]}. "
                              f"The calculator was invented by: {scratchpad[1]}")
                return answer, step + 1

            with span("tool_call", tool=tool, arg=str(arg)):
                if tool == "calculator":
                    result = calculator_tool(arg)
                elif tool == "search":
                    result = search_tool(arg)
                scratchpad.append(result)

    return None, max_steps  # ran out of steps without finishing


answer, steps_used = run_agent(max_steps=5)
print(f"answer: {answer}")
print(f"steps used: {steps_used}\n")
print_trace()


def buggy_agent_policy(step, scratchpad):
    """A policy with a bug: it keeps re-searching and never reaches
    final_answer, standing in for a real agent stuck re-planning forever."""
    return ("search", "who invented the calculator")


def run_buggy_agent(max_steps=4):
    TRACE.clear()
    scratchpad = []
    with span("agent_run", max_steps=max_steps):
        for step in range(max_steps):
            with span("think", step=step):
                tool, arg = buggy_agent_policy(step, scratchpad)
            with span("tool_call", tool=tool, arg=str(arg)):
                scratchpad.append(search_tool(arg))
    return None, max_steps


print("\n--- a second run: a policy with a bug ---\n")
answer2, steps_used2 = run_buggy_agent(max_steps=4)
hit_cap = steps_used2 >= 4
print(f"answer: {answer2}")
print(f"steps used: {steps_used2}   hit the step cap without finishing: {hit_cap}\n")
print_trace()