diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 9cecdc57fc..80c610e35e 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -124,6 +124,9 @@ ) INNER_RESPONSE_ID_CAPTURED_FIELD: Final[str] = "response_id" INNER_USAGE_CAPTURED_FIELD: Final[str] = "usage" +INNER_CAPTURED_RESPONSE_ID: Final[contextvars.ContextVar[str | None]] = contextvars.ContextVar( + "inner_captured_response_id", default=None +) # Tracks accumulated token usage from all inner chat completion spans within an agent invoke. INNER_ACCUMULATED_USAGE: Final[contextvars.ContextVar[UsageDetails | None]] = contextvars.ContextVar( @@ -2290,6 +2293,14 @@ def __init__( self.token_usage_histogram = _get_token_usage_histogram() self.duration_histogram = _get_duration_histogram() + def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]: + """Return provider-specific attributes emitted on agent spans.""" + return {} + + def _should_capture_agent_response_id(self) -> bool: + """Return whether the agent span must retain an inner response ID.""" + return False + def _trace_agent_invocation( self, *, @@ -2333,6 +2344,7 @@ def _trace_agent_invocation( all_options=dict(merged_options), **merged_client_kwargs, ) + attributes.update(self._get_additional_otel_agent_attributes()) if stream: # Do NOT set the inner-telemetry context vars here: this synchronous run() body executes @@ -2344,6 +2356,7 @@ def _trace_agent_invocation( # below), so set and reset both happen in the consumer's context. inner_response_telemetry_captured_fields: set[str] = set() inner_response_telemetry_captured_fields_token: contextvars.Token[set[str] | None] | None = None + inner_captured_response_id_token: contextvars.Token[str | None] | None = None inner_accumulated_usage_token: contextvars.Token[UsageDetails | None] | None = None # Agent Framework's agents run in-process (the actual network call happens on a nested # chat span), so invoke_agent spans use the default INTERNAL kind. @@ -2411,10 +2424,14 @@ async def _finalize_stream() -> None: response_attributes = _get_response_attributes( attributes, response, - capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD - not in inner_response_telemetry_captured_fields, + capture_response_id=( + self._should_capture_agent_response_id() + or 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, ) + if self._should_capture_agent_response_id(): + _apply_captured_response_id(response_attributes) _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields) _capture_response(span=span, attributes=response_attributes, duration=duration) if ( @@ -2436,6 +2453,8 @@ async def _finalize_stream() -> None: # pull-context factory below set the tokens in — so the reset is cross-context safe. if inner_response_telemetry_captured_fields_token is not None: INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token) + if inner_captured_response_id_token is not None: + INNER_CAPTURED_RESPONSE_ID.reset(inner_captured_response_id_token) if inner_accumulated_usage_token is not None: INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token) _close_span() @@ -2447,11 +2466,14 @@ def _inner_telemetry_pull_context() -> contextlib.AbstractContextManager[Any]: # avoiding the cross-context Token reset failure. Setting happens before the # underlying iterator is pulled, so inner chat completion spans created during the # pull can still accumulate usage / mark captured fields. - nonlocal inner_response_telemetry_captured_fields_token, inner_accumulated_usage_token + nonlocal inner_response_telemetry_captured_fields_token + nonlocal inner_captured_response_id_token + nonlocal inner_accumulated_usage_token if inner_response_telemetry_captured_fields_token is None: inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set( inner_response_telemetry_captured_fields ) + inner_captured_response_id_token = INNER_CAPTURED_RESPONSE_ID.set(None) inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({}) return _activate_span(span) @@ -2483,6 +2505,7 @@ async def _run() -> AgentResponse[Any]: inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set( inner_response_telemetry_captured_fields ) + inner_captured_response_id_token = INNER_CAPTURED_RESPONSE_ID.set(None) inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({}) try: with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span: @@ -2500,12 +2523,16 @@ async def _run() -> AgentResponse[Any]: response_attributes = _get_response_attributes( attributes, response, - capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD - not in inner_response_telemetry_captured_fields, + capture_response_id=( + self._should_capture_agent_response_id() + or 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 ), ) + if self._should_capture_agent_response_id(): + _apply_captured_response_id(response_attributes) _apply_accumulated_usage( response_attributes, inner_response_telemetry_captured_fields, @@ -2531,6 +2558,7 @@ async def _run() -> AgentResponse[Any]: raise finally: INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token) + INNER_CAPTURED_RESPONSE_ID.reset(inner_captured_response_id_token) INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token) return _run() @@ -3439,6 +3467,7 @@ def _mark_inner_response_telemetry_captured( return if response.response_id: captured_fields.add(INNER_RESPONSE_ID_CAPTURED_FIELD) + INNER_CAPTURED_RESPONSE_ID.set(response.response_id) if response.usage_details: captured_fields.add(INNER_USAGE_CAPTURED_FIELD) accumulated = INNER_ACCUMULATED_USAGE.get() @@ -3448,6 +3477,12 @@ def _mark_inner_response_telemetry_captured( INNER_ACCUMULATED_USAGE.set(add_usage_details(accumulated, response.usage_details)) +def _apply_captured_response_id(attributes: dict[str, Any]) -> None: + """Apply the inner chat response ID to an agent span when the provider requires it.""" + if response_id := INNER_CAPTURED_RESPONSE_ID.get(): + attributes.setdefault(OtelAttr.RESPONSE_ID, response_id) + + 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: diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index b2cd1603ad..684f4b7e62 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -4847,10 +4847,13 @@ async def _get() -> ChatResponse: @pytest.mark.parametrize("stream", [False, True]) -async def test_agent_and_chat_spans_do_not_duplicate_response_telemetry( - span_exporter: InMemorySpanExporter, stream: bool +@pytest.mark.parametrize("retain_agent_response_id", [False, True]) +async def test_agent_provider_hooks_control_response_telemetry( + span_exporter: InMemorySpanExporter, + stream: bool, + retain_agent_response_id: bool, ): - """The inner chat span owns response-id; usage is aggregated on the agent span.""" + """Provider hooks can add root attributes and retain an inner response ID when required.""" class NestedTelemetryChatClient(ChatTelemetryLayer, BaseChatClient[Any]): def service_url(self): @@ -4890,7 +4893,14 @@ async def _get() -> ChatResponse: return _get() - agent = Agent( + class ProviderTelemetryAgent(Agent): + def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]: + return {"test.provider.attribute": "provider-value"} + + def _should_capture_agent_response_id(self) -> bool: + return retain_agent_response_id + + agent = ProviderTelemetryAgent( client=NestedTelemetryChatClient(), # ty: ignore[invalid-argument-type] id="nested_agent_id", name="nested_agent", @@ -4921,7 +4931,11 @@ async def _get() -> ChatResponse: assert chat_span.attributes[OtelAttr.INPUT_TOKENS] == 11 # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] assert chat_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22 # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] - assert OtelAttr.RESPONSE_ID not in agent_span.attributes # type: ignore[operator] # pyrefly: ignore[not-iterable] # ty: ignore[unsupported-operator] + assert agent_span.attributes["test.provider.attribute"] == "provider-value" # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] + if retain_agent_response_id: + assert agent_span.attributes[OtelAttr.RESPONSE_ID] == "nested_resp_123" # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] + else: + assert OtelAttr.RESPONSE_ID not in agent_span.attributes # type: ignore[operator] # pyrefly: ignore[not-iterable] # ty: ignore[unsupported-operator] # The agent span carries the aggregated usage from all inner chat completions assert agent_span.attributes[OtelAttr.INPUT_TOKENS] == 11 # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] assert agent_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22 # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index cd44b46088..a989421844 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -36,6 +36,7 @@ from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient from azure.ai.projects.aio import AIProjectClient +from azure.ai.projects.models import ConnectionType from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -95,6 +96,7 @@ class FoundryAgentSettings(TypedDict, total=False): FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY = "foundry_hosted_agent_session_id" +_FOUNDRY_PROJECT_ARM_ID_ATTRIBUTE = "microsoft.foundry.project.id" class FoundryAgentOptions(OpenAIChatOptions, total=False): @@ -750,6 +752,7 @@ def __init__( client_kwargs["function_invocation_configuration"] = function_invocation_configuration client = actual_client_type(**client_kwargs) + self._foundry_project_arm_id: str | None = None super().__init__( client=client, # type: ignore[arg-type] @@ -841,6 +844,22 @@ def _update_session_from_chat_response_update( if session is not None and isinstance(agent_session_id, str) and agent_session_id: session.state[FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY] = agent_session_id + async def _get_foundry_project_arm_id(self) -> str: + """Get the Foundry project ARM ID from its Application Insights connection.""" + client = cast(RawFoundryAgentChatClient, self.client) + # AIProjectClient does not expose the project ARM ID directly. Derive it from the + # project-scoped connection until https://github.com/Azure/azure-sdk-for-python/issues/48825 is addressed. + connections = client.project_client.connections.list(connection_type=ConnectionType.APPLICATION_INSIGHTS) + async for connection in connections: + connection_suffix = f"/connections/{connection.name}" + if not connection.id.lower().endswith(connection_suffix.lower()): + raise ValueError( + f"The Foundry Application Insights connection ID has an unexpected format: {connection.id!r}." + ) + return connection.id[: -len(connection_suffix)] + + raise ValueError("The Foundry project does not have an Application Insights connection.") + async def configure_azure_monitor( self, enable_sensitive_data: bool = False, @@ -858,6 +877,8 @@ async def configure_azure_monitor( Raises: ImportError: If azure-monitor-opentelemetry-exporter is not installed. + ValueError: If the Application Insights connection does not contain the expected + project-scoped ARM resource ID. """ from agent_framework.observability import ( OBSERVABILITY_SETTINGS, @@ -897,6 +918,8 @@ async def configure_azure_monitor( "Install it with: pip install azure-monitor-opentelemetry" ) from exc + self._foundry_project_arm_id = await self._get_foundry_project_arm_id() + if "resource" not in kwargs: kwargs["resource"] = create_resource() @@ -951,6 +974,18 @@ class FoundryAgent( # type: ignore[misc] ) """ + @override + def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]: + """Return Foundry attributes required to discover the agent trace.""" + if self._foundry_project_arm_id: + return {_FOUNDRY_PROJECT_ARM_ID_ATTRIBUTE: self._foundry_project_arm_id} + return {} + + @override + def _should_capture_agent_response_id(self) -> bool: + """Keep the response ID on the client agent span for Foundry trace discovery.""" + return True + def __init__( self, *, diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index a4c3b8ca25..b71de78072 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -947,6 +947,43 @@ def test_raw_foundry_agent_init_creates_client() -> None: assert agent.client is not None assert cast(Any, agent.client).agent_name == "test-agent" + assert agent.name == "test-agent" + + +async def test_get_foundry_project_arm_id_from_application_insights_connection() -> None: + """Test that project attribution uses the public project-scoped connection ID.""" + + project_arm_id = ( + "/subscriptions/test-sub/resourceGroups/test-rg/providers/" + "Microsoft.CognitiveServices/accounts/test-account/projects/test-project" + ) + project_client = MagicMock() + + async def connections(): + yield SimpleNamespace(id=f"{project_arm_id}/connections/appinsights", name="appinsights") + + project_client.connections.list.return_value = connections() + agent = RawFoundryAgent(project_client=project_client, agent_name="test-agent") + + assert await agent._get_foundry_project_arm_id() == project_arm_id + project_client.connections.list.assert_called_once_with( + connection_type=projects_models.ConnectionType.APPLICATION_INSIGHTS + ) + + +async def test_get_foundry_project_arm_id_rejects_unexpected_connection_id() -> None: + """Test that malformed connection metadata does not silently disable portal attribution.""" + + project_client = MagicMock() + + async def connections(): + yield SimpleNamespace(id="appinsights", name="appinsights") + + project_client.connections.list.return_value = connections() + agent = RawFoundryAgent(project_client=project_client, agent_name="test-agent") + + with pytest.raises(ValueError, match="unexpected format"): + await agent._get_foundry_project_arm_id() def test_raw_foundry_agent_init_passes_default_headers_to_client() -> None: @@ -1243,6 +1280,8 @@ def test_foundry_agent_init() -> None: assert agent.client is not None assert cast(Any, agent.client).agent_name == "test-agent" + assert agent.name == "test-agent" + assert agent._should_capture_agent_response_id() def test_foundry_agent_init_with_middleware() -> None: @@ -1278,6 +1317,10 @@ async def test_foundry_agent_configure_azure_monitor() -> None: mock_views = MagicMock(return_value=[]) mock_resource = MagicMock() mock_enable = MagicMock() + project_arm_id = ( + "/subscriptions/test-sub/resourceGroups/test-rg/providers/" + "Microsoft.CognitiveServices/accounts/test-account/projects/test-project" + ) with ( patch.dict( @@ -1287,15 +1330,24 @@ async def test_foundry_agent_configure_azure_monitor() -> None: patch("agent_framework.observability.create_metric_views", mock_views), patch("agent_framework.observability.create_resource", return_value=mock_resource), patch("agent_framework.observability.enable_instrumentation", mock_enable), + patch( + "agent_framework_foundry._agent.RawFoundryAgent._get_foundry_project_arm_id", + new_callable=AsyncMock, + return_value=project_arm_id, + ) as mock_get_project_arm_id, ): await agent.configure_azure_monitor(enable_sensitive_data=True) mock_project.telemetry.get_application_insights_connection_string.assert_called_once() + mock_get_project_arm_id.assert_awaited_once_with() call_kwargs = mock_configure.call_args.kwargs assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint" assert call_kwargs["views"] == [] assert call_kwargs["resource"] is mock_resource mock_enable.assert_called_once_with(enable_sensitive_data=True) + assert agent._get_additional_otel_agent_attributes() == { + "microsoft.foundry.project.id": project_arm_id, + } async def test_foundry_agent_configure_azure_monitor_resource_not_found() -> None: