Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion python/packages/core/agent_framework/_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down
10 changes: 8 additions & 2 deletions python/packages/core/agent_framework/observability.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand Down
44 changes: 42 additions & 2 deletions python/packages/core/tests/core/test_observability.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down