Record actual served model as response model for Azure OpenAI

This commit is contained in:
Tao Chen
2026-05-17 19:50:21 -07:00
Unverified
parent da308f5f1e
commit 3f65aa055f
4 changed files with 390 additions and 39 deletions
@@ -101,6 +101,12 @@ INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS: Final[contextvars.ContextVar[set[str]
INNER_RESPONSE_ID_CAPTURED_FIELD: Final[str] = "response_id"
INNER_USAGE_CAPTURED_FIELD: Final[str] = "usage"
# Response header set by Azure OpenAI naming the model that actually served the
# request (which can differ from the deployment alias the caller sent). Chat
# clients may surface this on ``ChatResponse.additional_properties`` so the
# telemetry layer can promote it to ``gen_ai.response.model``.
AZURE_OPENAI_SERVED_MODEL_HEADER: Final[str] = "x-ms-served-model"
# 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
@@ -2125,6 +2131,14 @@ def _get_response_attributes(
attributes[OtelAttr.FINISH_REASONS] = json.dumps([finish_reason])
if model := getattr(response, "model", None):
attributes[OtelAttr.RESPONSE_MODEL] = model
# If the underlying provider reports the actually served model via the
# ``x-ms-served-model`` response header (Azure OpenAI), prefer it over the
# model reported on the response body for the response model attribute.
additional_properties = getattr(response, "additional_properties", None)
if isinstance(additional_properties, Mapping):
candidate = cast("Mapping[str, Any]", additional_properties).get(AZURE_OPENAI_SERVED_MODEL_HEADER)
if isinstance(candidate, str) and candidate:
attributes[OtelAttr.RESPONSE_MODEL] = candidate
if capture_usage and (usage := response.usage_details):
input_tokens = usage.get("input_token_count")
if input_tokens:
@@ -1739,6 +1739,78 @@ def test_get_response_attributes_capture_response_id_false():
assert OtelAttr.RESPONSE_ID not in result
def test_get_response_attributes_served_model_overrides_response_model():
"""When the response carries the Azure ``x-ms-served-model`` header, it should override RESPONSE_MODEL."""
from unittest.mock import Mock
from agent_framework.observability import (
AZURE_OPENAI_SERVED_MODEL_HEADER,
OtelAttr,
_get_response_attributes,
)
response = Mock()
response.response_id = None
response.finish_reason = None
response.raw_representation = None
response.usage_details = None
response.model = "gpt-4"
response.additional_properties = {AZURE_OPENAI_SERVED_MODEL_HEADER: "gpt-4o-2024-08-06"}
attrs = {OtelAttr.REQUEST_MODEL: "my-deployment-alias"}
result = _get_response_attributes(attrs, response)
# REQUEST_MODEL is left untouched; RESPONSE_MODEL is overridden by the served-model header.
assert result[OtelAttr.REQUEST_MODEL] == "my-deployment-alias"
assert result[OtelAttr.RESPONSE_MODEL] == "gpt-4o-2024-08-06"
def test_get_response_attributes_no_served_model_keeps_response_model():
"""Without the served-model header RESPONSE_MODEL should reflect the response's reported model."""
from unittest.mock import Mock
from agent_framework.observability import OtelAttr, _get_response_attributes
response = Mock()
response.response_id = None
response.finish_reason = None
response.raw_representation = None
response.usage_details = None
response.model = "gpt-4"
response.additional_properties = {}
attrs = {OtelAttr.REQUEST_MODEL: "my-deployment-alias"}
result = _get_response_attributes(attrs, response)
assert result[OtelAttr.REQUEST_MODEL] == "my-deployment-alias"
assert result[OtelAttr.RESPONSE_MODEL] == "gpt-4"
def test_get_response_attributes_ignores_non_string_served_model():
"""A non-string / empty value in the served-model header should not override RESPONSE_MODEL."""
from unittest.mock import Mock
from agent_framework.observability import (
AZURE_OPENAI_SERVED_MODEL_HEADER,
OtelAttr,
_get_response_attributes,
)
response = Mock()
response.response_id = None
response.finish_reason = None
response.raw_representation = None
response.usage_details = None
response.model = "gpt-4"
response.additional_properties = {AZURE_OPENAI_SERVED_MODEL_HEADER: ""}
attrs = {OtelAttr.REQUEST_MODEL: "my-deployment-alias"}
result = _get_response_attributes(attrs, response)
assert result[OtelAttr.REQUEST_MODEL] == "my-deployment-alias"
assert result[OtelAttr.RESPONSE_MODEL] == "gpt-4"
# region Test _get_exporters_from_env
@@ -2481,6 +2553,28 @@ def test_capture_response(span_exporter: InMemorySpanExporter):
assert spans[0].attributes.get(OtelAttr.OUTPUT_TOKENS) == 50
def test_capture_response_does_not_update_span_name_with_request_model(span_exporter: InMemorySpanExporter):
"""_capture_response should not rename the span even when REQUEST_MODEL is set."""
from agent_framework.observability import OtelAttr, _capture_response, get_tracer
span_exporter.clear()
tracer = get_tracer()
attrs = {
OtelAttr.OPERATION: "chat",
OtelAttr.REQUEST_MODEL: "my-deployment-alias",
OtelAttr.RESPONSE_MODEL: "gpt-4o-2024-08-06",
}
with tracer.start_as_current_span("chat my-deployment-alias") as span:
_capture_response(span=span, attributes=attrs)
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
assert spans[0].name == "chat my-deployment-alias"
assert spans[0].attributes.get(OtelAttr.RESPONSE_MODEL) == "gpt-4o-2024-08-06"
async def test_layer_ordering_span_sequence_with_function_calling(span_exporter: InMemorySpanExporter):
"""Test that with correct layer ordering, spans appear in the expected sequence.