Python: Fix Model ID attribute not showing up in invoke_agent span (#2061)

* Best effort to surface the model id to invoke agent span

* Fix tests

* Fix tests

Author: TaoChenOSU

python/packages/core/agent_framework/_agents.py

@@ -642,6 +642,7 @@ def __init__(
             max_tokens: The maximum number of tokens to generate.
             metadata: Additional metadata to include in the request.
             model_id: The model_id to use for the agent.
+                This overrides the model_id set in the chat client if it contains one.
             presence_penalty: The presence penalty to use.
             response_format: The format of the response.
             seed: The random seed to use.
@@ -690,7 +691,7 @@ def __init__(
         self._local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
         agent_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
         self.chat_options = ChatOptions(
-            model_id=model_id,
+            model_id=model_id or (str(chat_client.model_id) if hasattr(chat_client, "model_id") else None),
             allow_multiple_tool_calls=allow_multiple_tool_calls,
             conversation_id=conversation_id,
             frequency_penalty=frequency_penalty,

python/packages/core/agent_framework/observability.py

@@ -846,6 +846,7 @@ async def trace_get_response(
                 kwargs.get("model_id")
                 or (chat_options.model_id if (chat_options := kwargs.get("chat_options")) else None)
                 or getattr(self, "model_id", None)
+                or "unknown"
             )
             service_url = str(
                 service_url_func()
@@ -933,6 +934,7 @@ async def trace_get_streaming_response(
                 kwargs.get("model_id")
                 or (chat_options.model_id if (chat_options := kwargs.get("chat_options")) else None)
                 or getattr(self, "model_id", None)
+                or "unknown"
             )
             service_url = str(
                 service_url_func()
@@ -1324,7 +1326,10 @@ def _get_span(
     attributes: dict[str, Any],
     span_name_attribute: str,
 ) -> Generator["trace.Span", Any, Any]:
-    """Start a span for a agent run."""
+    """Start a span for a agent run.
+
+    Note: `attributes` must contain the `span_name_attribute` key.
+    """
     span = get_tracer().start_span(f"{attributes[OtelAttr.OPERATION]} {attributes[span_name_attribute]}")
     span.set_attributes(attributes)
     with trace.use_span(
@@ -1353,7 +1358,8 @@ def _get_span_attributes(**kwargs: Any) -> dict[str, Any]:
         attributes[SpanAttributes.LLM_SYSTEM] = system_name
     if provider_name := kwargs.get("provider_name"):
         attributes[OtelAttr.PROVIDER_NAME] = provider_name
-    attributes[SpanAttributes.LLM_REQUEST_MODEL] = kwargs.get("model", "unknown")
+    if model_id := kwargs.get("model", chat_options.model_id):
+        attributes[SpanAttributes.LLM_REQUEST_MODEL] = model_id
     if service_url := kwargs.get("service_url"):
         attributes[OtelAttr.ADDRESS] = service_url
     if conversation_id := kwargs.get("conversation_id", chat_options.conversation_id):

python/packages/core/tests/core/test_observability.py

@@ -279,6 +279,45 @@ async def test_chat_client_streaming_observability(
         assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
 
 
+async def test_chat_client_without_model_id_observability(mock_chat_client, span_exporter: InMemorySpanExporter):
+    """Test telemetry shouldn't fail when the model_id is not provided for unknown reason."""
+    client = use_observability(mock_chat_client)()
+    messages = [ChatMessage(role=Role.USER, text="Test")]
+    span_exporter.clear()
+    response = await client.get_response(messages=messages)
+
+    assert response is not None
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+
+    assert span.name == "chat unknown"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+
+
+async def test_chat_client_streaming_without_model_id_observability(
+    mock_chat_client, span_exporter: InMemorySpanExporter
+):
+    """Test streaming telemetry shouldn't fail when the model_id is not provided for unknown reason."""
+    client = use_observability(mock_chat_client)()
+    messages = [ChatMessage(role=Role.USER, text="Test")]
+    span_exporter.clear()
+    # Collect all yielded updates
+    updates = []
+    async for update in client.get_streaming_response(messages=messages):
+        updates.append(update)
+
+    # Verify we got the expected updates, this shouldn't be dependent on otel
+    assert len(updates) == 2
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "chat unknown"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+
+
 def test_prepend_user_agent_with_none_value():
     """Test prepend user agent with None value in headers."""
     headers = {"User-Agent": None}
@@ -368,6 +407,7 @@ def __init__(self):
             self.name = "test_agent"
             self.display_name = "Test Agent"
             self.description = "Test agent description"
+            self.chat_options = ChatOptions(model_id="TestModel")
 
         async def run(self, messages=None, *, thread=None, **kwargs):
             return AgentRunResponse(
@@ -405,7 +445,7 @@ async def test_agent_instrumentation_enabled(
     assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
     assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
     assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
-    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "TestModel"
     assert span.attributes[OtelAttr.INPUT_TOKENS] == 15
     assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25
     if enable_sensitive_data:
@@ -433,7 +473,7 @@ async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
     assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
     assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
     assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
-    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "TestModel"
     if enable_sensitive_data:
         assert span.attributes.get(OtelAttr.OUTPUT_MESSAGES) is not None  # Streaming, so no usage yet