diff --git a/src/lightspeed_agentic/audit.py b/src/lightspeed_agentic/audit.py index fcc5a9b..eb006a0 100644 --- a/src/lightspeed_agentic/audit.py +++ b/src/lightspeed_agentic/audit.py @@ -7,6 +7,9 @@ from datetime import UTC, datetime from typing import Any +from opentelemetry.context import Context +from opentelemetry.trace import StatusCode + from lightspeed_agentic.tracing import get_tracer from lightspeed_agentic.types import ProviderEvent @@ -35,10 +38,20 @@ def __init__( self._text_buffer: list[str] = [] self._thinking_buffer: list[str] = [] self._last_tool_name = "unknown" + self._last_tool_input = "" self._tool_span: Any = None self._tracer = get_tracer() + self._parent_context: Context | None = None self._emit("audit.agent.started", model=model, provider=provider) + def set_parent_context(self, ctx: Context) -> None: + """Set the parent span context for tool spans. + + Must be called from within the agent.run span block so that + tool spans are correctly parented under agent.run. + """ + self._parent_context = ctx + def process_event(self, event: ProviderEvent) -> None: match event.type: case "text_delta": @@ -52,7 +65,17 @@ def process_event(self, event: ProviderEvent) -> None: if self._tool_span is not None: self._tool_span.end() self._last_tool_name = event.name or "unknown" - self._tool_span = self._tracer.start_span(f"tool.{self._last_tool_name}") + self._last_tool_input = event.input or "" + # Create tool span with explicit parent context so it's correctly + # parented under agent.run even in async code + self._tool_span = self._tracer.start_span( + f"tool.{self._last_tool_name}", + context=self._parent_context, + attributes={ + "tool.name": self._last_tool_name, + "tool.input": self._last_tool_input[:300], + }, + ) self._emit( "audit.agent.tool.call", tool_name=self._last_tool_name, tool_input=event.input ) @@ -64,6 +87,8 @@ def process_event(self, event: ProviderEvent) -> None: success=True, ) if self._tool_span is not None: + self._tool_span.set_attribute("tool.output", event.output[:500]) + self._tool_span.set_status(StatusCode.OK) self._tool_span.end() self._tool_span = None case "result": diff --git a/src/lightspeed_agentic/routes/query.py b/src/lightspeed_agentic/routes/query.py index f907722..2a08afd 100644 --- a/src/lightspeed_agentic/routes/query.py +++ b/src/lightspeed_agentic/routes/query.py @@ -12,6 +12,7 @@ from typing import Any from fastapi import APIRouter, Request +from opentelemetry import context as otel_context from lightspeed_agentic.audit import AuditLogger, derive_phase from lightspeed_agentic.logging import EventLogger @@ -108,6 +109,9 @@ async def run() -> None: context=trace_ctx, attributes={"model": model, "provider": provider.name}, ): + # Pass the current context to AuditLogger so tool spans + # are correctly parented under agent.run + audit_logger.set_parent_context(otel_context.get_current()) result = provider.query( ProviderQueryOptions( prompt=prompt,