diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index a0cbd6a1a0..7462181036 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -14,6 +14,7 @@ Commonly used exports: from __future__ import annotations import contextlib +import contextvars import json import logging import os @@ -65,6 +66,7 @@ if TYPE_CHECKING: # pragma: no cover GeneratedEmbeddings, Message, ResponseStream, + UsageDetails, ) ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) @@ -93,6 +95,18 @@ ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") logger = logging.getLogger("agent_framework") +INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS: Final[contextvars.ContextVar[set[str] | None]] = contextvars.ContextVar( + "inner_response_telemetry_captured_fields", default=None +) +INNER_RESPONSE_ID_CAPTURED_FIELD: Final[str] = "response_id" +INNER_USAGE_CAPTURED_FIELD: Final[str] = "usage" + +# Tracks accumulated token usage from all inner chat completion spans within an agent invoke. +INNER_ACCUMULATED_USAGE: Final[contextvars.ContextVar[UsageDetails | None]] = contextvars.ContextVar( + "inner_accumulated_usage", default=None +) + + OTEL_METRICS: Final[str] = "__otel_metrics__" TOKEN_USAGE_BUCKET_BOUNDARIES: Final[tuple[float, ...]] = ( 1, @@ -1314,6 +1328,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): operation_duration_histogram=getattr(self, "duration_histogram", None), duration=duration, ) + _mark_inner_response_telemetry_captured(response) if ( OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and isinstance(response, ChatResponse) @@ -1373,6 +1388,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): operation_duration_histogram=getattr(self, "duration_histogram", None), duration=duration, ) + _mark_inner_response_telemetry_captured(response) if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages: finish_reason = cast( "FinishReason | None", @@ -1527,8 +1543,6 @@ class AgentTelemetryLayer: super().run, # type: ignore[misc] ) provider_name = str(self.otel_provider_name) - capture_usage = bool(getattr(self, "_otel_capture_usage", True)) - if not OBSERVABILITY_SETTINGS.ENABLED: return super_run( # type: ignore[no-any-return] messages=messages, @@ -1557,23 +1571,34 @@ class AgentTelemetryLayer: **merged_client_kwargs, ) + inner_response_telemetry_captured_fields: set[str] = set() + inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set( + inner_response_telemetry_captured_fields + ) + inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({}) + if stream: - run_result: object = super_run( - messages=messages, - stream=True, - session=session, - compaction_strategy=compaction_strategy, - tokenizer=tokenizer, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, - **kwargs, - ) - if isinstance(run_result, ResponseStream): - result_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = run_result # pyright: ignore[reportUnknownVariableType] - elif isinstance(run_result, Awaitable): - result_stream = ResponseStream.from_awaitable(run_result) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - else: - raise RuntimeError("Streaming telemetry requires a ResponseStream result.") + try: + run_result: object = super_run( + messages=messages, + stream=True, + session=session, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + function_invocation_kwargs=function_invocation_kwargs, + client_kwargs=client_kwargs, + **kwargs, + ) + if isinstance(run_result, ResponseStream): + result_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = run_result # pyright: ignore[reportUnknownVariableType] + elif isinstance(run_result, Awaitable): + result_stream = ResponseStream.from_awaitable(run_result) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + else: + raise RuntimeError("Streaming telemetry requires a ResponseStream result.") + except Exception: + INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token) + INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token) + raise # Create span directly without trace.use_span() context attachment. # Streaming spans are closed asynchronously in cleanup hooks, which run @@ -1613,8 +1638,11 @@ class AgentTelemetryLayer: response_attributes = _get_response_attributes( attributes, response, - capture_usage=capture_usage, + capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD + not in inner_response_telemetry_captured_fields, + capture_usage=INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields, ) + _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields) _capture_response(span=span, attributes=response_attributes, duration=duration) if ( OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED @@ -1630,6 +1658,8 @@ class AgentTelemetryLayer: except Exception as exception: capture_exception(span=span, exception=exception, timestamp=time_ns()) finally: + INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token) + INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token) _close_span() # Register a weak reference callback to close the span if stream is garbage collected @@ -1641,41 +1671,52 @@ class AgentTelemetryLayer: return wrapped_stream async def _run() -> AgentResponse: - with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span: - if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages: - _capture_messages( - span=span, - provider_name=provider_name, - messages=messages, - system_instructions=_get_instructions_from_options(merged_options), - ) - start_time_stamp = perf_counter() - try: - response: AgentResponse[Any] = await super_run( - messages=messages, - stream=False, - session=session, - compaction_strategy=compaction_strategy, - tokenizer=tokenizer, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, - **kwargs, - ) - except Exception as exception: - capture_exception(span=span, exception=exception, timestamp=time_ns()) - raise - duration = perf_counter() - start_time_stamp - if response: - response_attributes = _get_response_attributes(attributes, response, capture_usage=capture_usage) - _capture_response(span=span, attributes=response_attributes, duration=duration) - if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages: + try: + with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span: + if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages: _capture_messages( span=span, provider_name=provider_name, - messages=response.messages, - output=True, + messages=messages, + system_instructions=_get_instructions_from_options(merged_options), ) - return response # type: ignore[return-value,no-any-return] + start_time_stamp = perf_counter() + try: + response: AgentResponse[Any] = await super_run( + messages=messages, + stream=False, + session=session, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + function_invocation_kwargs=function_invocation_kwargs, + client_kwargs=client_kwargs, + **kwargs, + ) + except Exception as exception: + capture_exception(span=span, exception=exception, timestamp=time_ns()) + raise + duration = perf_counter() - start_time_stamp + if response: + response_attributes = _get_response_attributes( + attributes, + response, + capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD + not in inner_response_telemetry_captured_fields, + capture_usage=INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields, + ) + _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields) + _capture_response(span=span, attributes=response_attributes, duration=duration) + if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages: + _capture_messages( + span=span, + provider_name=provider_name, + messages=response.messages, + output=True, + ) + return response # type: ignore[return-value,no-any-return] + finally: + INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token) + INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token) return _run() @@ -1931,14 +1972,46 @@ def _to_otel_part(content: Content) -> dict[str, Any] | None: return None +def _mark_inner_response_telemetry_captured(response: ChatResponse | AgentResponse) -> None: + """Record when an inner chat telemetry span already captured response metadata.""" + captured_fields = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.get() + if captured_fields is None: + return + if response.response_id: + captured_fields.add(INNER_RESPONSE_ID_CAPTURED_FIELD) + if response.usage_details: + captured_fields.add(INNER_USAGE_CAPTURED_FIELD) + accumulated = INNER_ACCUMULATED_USAGE.get() + if accumulated is not None: + from ._types import add_usage_details + + INNER_ACCUMULATED_USAGE.set(add_usage_details(accumulated, response.usage_details)) + + +def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[str]) -> None: + """Apply accumulated usage from inner chat spans to the invoke_agent span attributes.""" + if INNER_USAGE_CAPTURED_FIELD not in captured_fields: + return + accumulated = INNER_ACCUMULATED_USAGE.get() + if not accumulated: + return + input_tokens = accumulated.get("input_token_count") + if input_tokens: + attributes[OtelAttr.INPUT_TOKENS] = input_tokens + output_tokens = accumulated.get("output_token_count") + if output_tokens: + attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens + + def _get_response_attributes( attributes: dict[str, Any], response: ChatResponse | AgentResponse, *, + capture_response_id: bool = True, capture_usage: bool = True, ) -> dict[str, Any]: """Get the response attributes from a response.""" - if response.response_id: + if capture_response_id and response.response_id: attributes[OtelAttr.RESPONSE_ID] = response.response_id finish_reason = getattr(response, "finish_reason", None) if not finish_reason: diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 7982985b94..5152712b8c 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -11,6 +11,7 @@ from opentelemetry.trace import StatusCode from agent_framework import ( AGENT_FRAMEWORK_USER_AGENT, + Agent, AgentResponse, BaseChatClient, ChatResponse, @@ -473,10 +474,10 @@ def mock_chat_agent(): @pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True) -async def test_agent_instrumentation_enabled( +async def test_agent_span_captures_response_telemetry_without_inner_chat_span( mock_chat_agent: SupportsAgentRun, span_exporter: InMemorySpanExporter, enable_sensitive_data ): - """Test that when agent diagnostics are enabled, telemetry is applied.""" + """Agent spans should retain response telemetry when no inner chat span owns it.""" agent = mock_chat_agent() @@ -492,6 +493,7 @@ async def test_agent_instrumentation_enabled( assert span.attributes[OtelAttr.AGENT_NAME] == "test_agent" assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description" assert span.attributes[OtelAttr.REQUEST_MODEL] == "TestModel" + assert span.attributes[OtelAttr.RESPONSE_ID] == "test_response_id" assert span.attributes[OtelAttr.INPUT_TOKENS] == 15 assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25 if enable_sensitive_data: @@ -1700,6 +1702,24 @@ def test_get_response_attributes_capture_usage_false(): assert OtelAttr.OUTPUT_TOKENS not in result +def test_get_response_attributes_capture_response_id_false(): + """Test _get_response_attributes skips response_id when capture_response_id is False.""" + from unittest.mock import Mock + + from agent_framework.observability import OtelAttr, _get_response_attributes + + response = Mock() + response.response_id = "resp_123" + response.finish_reason = None + response.raw_representation = None + response.usage_details = None + + attrs = {} + result = _get_response_attributes(attrs, response, capture_response_id=False) + + assert OtelAttr.RESPONSE_ID not in result + + # region Test _get_exporters_from_env @@ -2530,6 +2550,82 @@ async def test_layer_ordering_span_sequence_with_function_calling(span_exporter: assert sorted_spans[2].name.startswith("chat"), f"Third span should be 'chat', got '{sorted_spans[2].name}'" +@pytest.mark.parametrize("stream", [False, True]) +async def test_agent_and_chat_spans_do_not_duplicate_response_telemetry( + span_exporter: InMemorySpanExporter, stream: bool +): + """The inner chat span owns response-id; usage is aggregated on the agent span.""" + + class NestedTelemetryChatClient(ChatTelemetryLayer, BaseChatClient[Any]): + def service_url(self): + return "https://test.example.com" + + def _inner_get_response( + self, *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + if stream: + + async def _stream() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(contents=[Content.from_text("Nested")], role="assistant") + yield ChatResponseUpdate(contents=[Content.from_text(" response")], role="assistant") + + def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: + return ChatResponse( + messages=[Message(role="assistant", text="Nested response")], + response_id="nested_resp_123", + usage_details=UsageDetails(input_token_count=11, output_token_count=22), + finish_reason="stop", + ) + + return ResponseStream(_stream(), finalizer=_finalize) + + async def _get() -> ChatResponse: + return ChatResponse( + messages=[Message(role="assistant", text="Nested response")], + response_id="nested_resp_123", + usage_details=UsageDetails(input_token_count=11, output_token_count=22), + finish_reason="stop", + ) + + return _get() + + agent = Agent( + client=NestedTelemetryChatClient(), + id="nested_agent_id", + name="nested_agent", + description="Nested telemetry agent", + default_options={"model_id": "NestedModel"}, + ) + + span_exporter.clear() + + if stream: + result_stream = agent.run("Test message", stream=True) + async for _ in result_stream: + pass + response = await result_stream.get_final_response() + else: + response = await agent.run("Test message") + + assert response is not None + + spans = span_exporter.get_finished_spans() + assert len(spans) == 2 + + span_by_operation = {span.attributes[OtelAttr.OPERATION.value]: span for span in spans} + agent_span = span_by_operation[OtelAttr.AGENT_INVOKE_OPERATION] + chat_span = span_by_operation[OtelAttr.CHAT_COMPLETION_OPERATION] + + assert chat_span.attributes[OtelAttr.RESPONSE_ID] == "nested_resp_123" + assert chat_span.attributes[OtelAttr.INPUT_TOKENS] == 11 + assert chat_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22 + + assert OtelAttr.RESPONSE_ID not in agent_span.attributes + # The agent span carries the aggregated usage from all inner chat completions + assert agent_span.attributes[OtelAttr.INPUT_TOKENS] == 11 + assert agent_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22 + + # region Test non-ASCII character handling in JSON serialization