Tracing

from exo.trace import traced

@traced(name="research_task")
async def research(topic: str) -> str:
    """A traced function -- automatically creates a span."""
    # ... do work ...
    return f"Results for {topic}"

# When called, this creates an OpenTelemetry span named "research_task"
result = await research("quantum computing")

from exo.trace import traced

# Async function
@traced(name="fetch_data", attributes={"source": "api"})
async def fetch_data(url: str) -> str:
    return await http_get(url)

# Sync function
@traced(name="process")
def process_data(data: str) -> str:
    return data.upper()

# Async generator (streaming)
@traced(name="stream_results")
async def stream_results(query: str):
    for chunk in chunks:
        yield chunk

@traced(
    name="my_operation",          # span name (defaults to function name)
    attributes={"key": "value"},  # static span attributes
    extract_args=True,            # include function args as attributes
)

from exo.trace.decorator import span_sync, span_async

# Sync context manager
with span_sync("my_sync_operation", attributes={"step": 1}):
    do_sync_work()

# Async context manager
async with span_async("my_async_operation"):
    await do_async_work()

from exo.trace import TraceConfig, TraceBackend

config = TraceConfig(
    service_name="my-agent-app",
    backend=TraceBackend.OTLP,
    endpoint="http://localhost:4317",
    sample_rate=1.0,
    enabled=True,
)

from exo.trace.instrumentation import record_agent_run, record_tool_step, Timer

# Time an operation
timer = Timer()
timer.start()
# ... do work ...
elapsed = timer.stop()

# Record an agent run span
record_agent_run(
    agent_name="researcher",
    step=1,
    model="gpt-4o",
    prompt_tokens=500,
    output_tokens=150,
    elapsed_ms=elapsed * 1000,
)

# Record a tool execution span
record_tool_step(
    tool_name="search",
    status="success",
    elapsed_ms=120.5,
)

from exo.trace.instrumentation import build_agent_attributes, build_tool_attributes

agent_attrs = build_agent_attributes(
    name="researcher",
    step=1,
    model="gpt-4o",
)
# {"exo.agent.name": "researcher", "exo.agent.step": 1, ...}

tool_attrs = build_tool_attributes(
    name="search",
    status="success",
)
# {"exo.tool.name": "search", "exo.tool.status": "success"}

from exo.trace import PromptLogger, compute_token_breakdown

# Compute per-role token breakdown
breakdown = compute_token_breakdown(messages)
print(f"System: {breakdown.system_tokens}")
print(f"User: {breakdown.user_tokens}")
print(f"Assistant: {breakdown.assistant_tokens}")
print(f"Tool: {breakdown.tool_tokens}")

# Logger tracks execution entries
logger = PromptLogger()
logger.log(
    messages=messages,
    response=response_text,
    model="gpt-4o",
    tokens=breakdown,
)

@dataclass(frozen=True)
class TokenBreakdown:
    system_tokens: int
    user_tokens: int
    assistant_tokens: int
    tool_tokens: int
    total_tokens: int

from exo.trace.propagation import BaggagePropagator, DictCarrier

propagator = BaggagePropagator()

# Inject baggage into outgoing request headers
carrier = DictCarrier()
propagator.inject(carrier)

# Extract baggage from incoming request headers
propagator.extract(incoming_headers)

from exo.trace.propagation import set_baggage, get_baggage, clear_baggage

# Set baggage values
set_baggage("user_id", "u-123")
set_baggage("session_id", "s-456")

# Read baggage
user_id = get_baggage("user_id")  # "u-123"

# Clear all baggage
clear_baggage()

from exo.trace.propagation import SpanConsumer, register_consumer

class LogConsumer(SpanConsumer):
    """Log all spans to a file."""

    def consume(self, span_data: dict) -> None:
        with open("spans.log", "a") as f:
            f.write(f"{span_data}\n")

# Register the consumer
register_consumer(LogConsumer())

@traced(name="agent_pipeline", extract_args=True)
async def run_pipeline(input: str) -> str:
    # Each sub-step creates a child span
    context = await prepare_context(input)
    response = await call_llm(context)
    result = await execute_tools(response)
    return result

@traced(name="prepare_context")
async def prepare_context(input: str) -> dict:
    ...

@traced(name="call_llm")
async def call_llm(context: dict) -> str:
    ...

@traced(name="execute_tools")
async def execute_tools(response: str) -> str:
    ...

from exo.trace.instrumentation import Timer

async def timed_operation():
    timer = Timer()
    timer.start()

    result = await expensive_operation()

    elapsed = timer.stop()
    record_agent_run(
        agent_name="worker",
        step=1,
        elapsed_ms=elapsed * 1000,
        prompt_tokens=0,
        output_tokens=0,
    )
    return result

from exo.trace.propagation import BaggagePropagator, DictCarrier

# Sender: inject context into request
carrier = DictCarrier()
propagator = BaggagePropagator()
propagator.inject(carrier)
# Send carrier.items as HTTP headers

# Receiver: extract context from request
propagator.extract(request.headers)
# Spans created here are linked to the sender's trace