diff --git a/sentry_sdk/integrations/openai_agents/patches/models.py b/sentry_sdk/integrations/openai_agents/patches/models.py index c2b030c42a..bf769ca518 100644 --- a/sentry_sdk/integrations/openai_agents/patches/models.py +++ b/sentry_sdk/integrations/openai_agents/patches/models.py @@ -12,12 +12,12 @@ add_sentry_baggage_to_headers, should_propagate_trace, ) -from sentry_sdk.utils import logger +from sentry_sdk.utils import capture_internal_exceptions, logger from ..spans import ai_client_span, update_ai_client_span if TYPE_CHECKING: - from typing import Any, Callable, Union + from typing import Any, Callable, Optional, Union from sentry_sdk.tracing import Span @@ -56,6 +56,39 @@ def _inject_trace_propagation_headers( headers[key] = value +class _ResponseModelRecordingStream: + """ + Proxies a provider chunk stream returned by `_fetch_response(stream=True)` and records + the model reported on each chunk. + + The default `OpenAIResponsesModel` is not affected: its terminal streaming event carries + the model that responded. This only applies to the opt-in Chat Completions code paths, + i.e. an explicit `OpenAIChatCompletionsModel` (commonly used with Azure OpenAI and + OpenAI-compatible providers), `set_default_openai_api("chat_completions")`, and + `LitellmModel`. For those, the Agents SDK synthesizes the streamed response with `model` + set to the requested model name, e.g. an Azure deployment name, so only the provider + chunks carry the model that actually responded. + """ + + def __init__(self, stream: "Any", record: "Callable[[str], None]") -> None: + self._sentry_stream = stream + self._sentry_record = record + + def __getattr__(self, name: str) -> "Any": + return getattr(self._sentry_stream, name) + + def __aiter__(self) -> "_ResponseModelRecordingStream": + return self + + async def __anext__(self) -> "Any": + chunk = await self._sentry_stream.__anext__() + with capture_internal_exceptions(): + chunk_model = getattr(chunk, "model", None) + if chunk_model: + self._sentry_record(str(chunk_model)) + return chunk + + def _get_model( original_get_model: "Callable[..., agents.Model]", agent: "agents.Agent", @@ -71,19 +104,49 @@ def _get_model( # because we only patch its direct methods, all underlying data can remain unchanged. model = copy.copy(original_get_model(agent, run_config)) - # Capture the request model name for spans (agent.model can be None when using defaults) - request_model_name = model.model if hasattr(model, "model") else str(model) - agent._sentry_request_model = request_model_name # type: ignore[attr-defined] + # The resolved model honors `RunConfig.model` and has provider prefixes such as "litellm/" + # stripped, so it is the source of truth for the request model (agent.model can be None + # when using defaults, or overridden by the run config). + resolved_model_name = getattr(model, "model", None) + request_model_name = ( + str(resolved_model_name) if resolved_model_name is not None else None + ) + agent._sentry_request_model = request_model_name or str(model) # type: ignore[attr-defined] + + # The model copy is created per turn, so this state is not shared between concurrent runs. + # It holds the model reported by the provider for the in-flight request. + response_model_state: "dict[str, Optional[str]]" = {"model": None} - # Wrap _fetch_response if it exists (for OpenAI models) to capture response model + def _record_response_model(response_model: str) -> None: + response_model_state["model"] = response_model + + def _pop_response_model() -> "Optional[str]": + response_model = response_model_state["model"] + response_model_state["model"] = None + return response_model + + # Wrap _fetch_response if it exists (for OpenAI and LiteLLM models) to capture the response model if hasattr(model, "_fetch_response"): original_fetch_response = model._fetch_response @wraps(original_fetch_response) async def wrapped_fetch_response(*args: "Any", **kwargs: "Any") -> "Any": response = await original_fetch_response(*args, **kwargs) - if hasattr(response, "model") and response.model: - agent._sentry_response_model = str(response.model) # type: ignore[attr-defined] + with capture_internal_exceptions(): + if hasattr(response, "model") and response.model: + _record_response_model(str(response.model)) + elif ( + isinstance(response, tuple) + and len(response) == 2 + and hasattr(type(response[1]), "__anext__") + ): + # Streamed Chat Completions return (synthesized Response, chunk stream). + return ( + response[0], + _ResponseModelRecordingStream( + response[1], _record_response_model + ), + ) return response model._fetch_response = wrapped_fetch_response @@ -99,16 +162,15 @@ async def wrapped_get_response(*args: "Any", **kwargs: "Any") -> "Any": tool for tool in mcp_tools if isinstance(tool, HostedMCPTool) ] - with ai_client_span(agent, kwargs) as span: + _pop_response_model() + with ai_client_span(agent, kwargs, request_model=request_model_name) as span: for hosted_tool in hosted_tools: _inject_trace_propagation_headers(hosted_tool, span=span) result = await original_get_response(*args, **kwargs) - # Get response model captured from _fetch_response and clean up - response_model = getattr(agent, "_sentry_response_model", None) - if response_model: - delattr(agent, "_sentry_response_model") + # Get response model captured from _fetch_response + response_model = _pop_response_model() update_ai_client_span(span, result, response_model, agent) @@ -137,7 +199,10 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any": tool for tool in mcp_tools if isinstance(tool, HostedMCPTool) ] - with ai_client_span(agent, span_kwargs) as span: + _pop_response_model() + with ai_client_span( + agent, span_kwargs, request_model=request_model_name + ) as span: for hosted_tool in hosted_tools: _inject_trace_propagation_headers(hosted_tool, span=span) @@ -167,7 +232,9 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any": # Update span with response data (usage, output, model) if streaming_response: - response_model = ( + # Prefer the model reported by provider chunks: the terminal response of + # streamed Chat Completions only echoes the requested model name. + response_model = _pop_response_model() or ( str(streaming_response.model) if hasattr(streaming_response, "model") and streaming_response.model diff --git a/sentry_sdk/integrations/openai_agents/spans/ai_client.py b/sentry_sdk/integrations/openai_agents/spans/ai_client.py index 697f0d6cf7..950cc3ed78 100644 --- a/sentry_sdk/integrations/openai_agents/spans/ai_client.py +++ b/sentry_sdk/integrations/openai_agents/spans/ai_client.py @@ -46,6 +46,7 @@ from ..consts import SPAN_ORIGIN from ..utils import ( + _get_agent_model_name, _set_agent_data, _set_input_data, _set_output_data, @@ -198,14 +199,12 @@ def _transform_tool_definitions(tools: "list[Tool]") -> "list[ToolDefinition]": def ai_client_span( - agent: "Agent", get_response_kwargs: "dict[str, Any]" + agent: "Agent", + get_response_kwargs: "dict[str, Any]", + request_model: "Optional[str]" = None, ) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": # TODO-anton: implement other types of operations. Now "chat" is hardcoded. - model_name = None - if agent.model: - model_name = agent.model.model if hasattr(agent.model, "model") else agent.model - elif hasattr(agent, "_sentry_request_model"): - model_name = agent._sentry_request_model + model_name = request_model or _get_agent_model_name(agent) client_options = sentry_sdk.get_client().options @@ -232,7 +231,7 @@ def ai_client_span( set_on_span = span.set_data - _set_agent_data(span, agent) + _set_agent_data(span, agent, request_model=request_model) if has_data_collection_enabled(client_options): if client_options["data_collection"]["gen_ai"]["inputs"]: diff --git a/sentry_sdk/integrations/openai_agents/utils.py b/sentry_sdk/integrations/openai_agents/utils.py index 01303e0e58..afbb9cd412 100644 --- a/sentry_sdk/integrations/openai_agents/utils.py +++ b/sentry_sdk/integrations/openai_agents/utils.py @@ -48,9 +48,23 @@ def _capture_exception(exc: "Any") -> None: sentry_sdk.capture_event(event, hint=hint) +def _get_agent_model_name(agent: "agents.Agent") -> "Optional[str]": + if isinstance(agent.model, Model) and hasattr(agent.model, "model"): + return agent.model.model + if isinstance(agent.model, str): + return agent.model + return getattr(agent, "_sentry_request_model", None) + + def _set_agent_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "agents.Agent" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + agent: "agents.Agent", + request_model: "Optional[str]" = None, ) -> None: + """ + `request_model` is the name of the model instance the runner resolved for this call. When + given, it takes precedence over `agent.model`, which `RunConfig.model` can override. + """ set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) @@ -64,14 +78,7 @@ def _set_agent_data( if agent.model_settings.max_tokens: set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, agent.model_settings.max_tokens) - model_name: "Optional[str]" = None - if isinstance(agent.model, Model) and hasattr(agent.model, "model"): - model_name = agent.model.model - elif isinstance(agent.model, str): - model_name = agent.model - elif hasattr(agent, "_sentry_request_model"): - model_name = agent._sentry_request_model - + model_name = request_model or _get_agent_model_name(agent) if model_name: set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index 8bc49a2c9f..92074b9bd7 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -7483,3 +7483,160 @@ async def test_runner_run_streamed_with_starting_agent_kwarg( (transaction,) = events assert transaction["transaction"] == "test_agent workflow" + + +def _get_chat_spans(items): + """Return (name, attributes) for chat spans, whether sent as span items or in a transaction.""" + chat_spans = [] + for item in items: + if item.type == "span": + if item.payload["attributes"].get("sentry.op") == OP.GEN_AI_CHAT: + chat_spans.append((item.payload["name"], item.payload["attributes"])) + elif item.type == "transaction": + for span in item.payload["spans"]: + if span["op"] == OP.GEN_AI_CHAT: + chat_spans.append((span["description"], span["data"])) + return chat_spans + + +def _mock_openai_client(handler): + return AsyncOpenAI( + api_key="z", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_ai_client_span_request_model_uses_run_config_model( + sentry_init, + capture_items, + test_agent, + stream_gen_ai_spans, + span_streaming, +): + """ + `RunConfig.model` takes precedence over `Agent.model`, so the chat span must report the + model that was actually requested. + """ + + def handler(request): + assert json.loads(request.content)["model"] == "gpt-4.1-mini" + return httpx.Response( + 200, + json=EXAMPLE_RESPONSE.model_copy( + update={"model": "gpt-4.1-mini-2025-04-14"} + ).model_dump(mode="json", exclude_none=True), + ) + + run_config = agents.RunConfig( + model=OpenAIResponsesModel( + model="gpt-4.1-mini", openai_client=_mock_openai_client(handler) + ), + tracing_disabled=True, + ) + + sentry_init( + integrations=[OpenAIAgentsIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + items = capture_items("span", "transaction") + + # test_agent.model is "gpt-4" + await agents.Runner.run(test_agent, "Test input", run_config=run_config) + sentry_sdk.flush() + + ((name, attributes),) = _get_chat_spans(items) + assert name == "chat gpt-4.1-mini" + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-4.1-mini" + assert attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-mini-2025-04-14" + + +@pytest.mark.skipif( + parse_version(OPENAI_AGENTS_VERSION) < (0, 7, 0), + reason="Chat Completions streaming in older openai-agents is incompatible with the tested openai versions", +) +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_ai_client_span_response_model_from_streamed_chat_completions( + sentry_init, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + """ + For streamed Chat Completions, the Agents SDK synthesizes the final response with the + requested model name (e.g. an Azure deployment name). The response model must come from + the provider chunks, since it is used to calculate model costs. + """ + + def chunk(**kwargs): + return { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 10000000, + "model": "gpt-4o-2024-08-06", + **kwargs, + } + + def handler(request): + chunks = [ + chunk( + choices=[ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hello"}, + "finish_reason": None, + } + ] + ), + chunk(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]), + chunk( + choices=[], + usage={ + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + }, + ), + ] + body = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + return httpx.Response( + 200, + content=(body + "data: [DONE]\n\n").encode("utf-8"), + headers={"Content-Type": "text/event-stream"}, + ) + + agent = Agent( + name="test_agent", + model=agents.OpenAIChatCompletionsModel( + model="my-azure-deployment", openai_client=_mock_openai_client(handler) + ), + model_settings=ModelSettings(include_usage=True), + ) + + sentry_init( + integrations=[OpenAIAgentsIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + items = capture_items("span", "transaction") + + result = agents.Runner.run_streamed(agent, "Test input", run_config=test_run_config) + async for _event in result.stream_events(): + pass + sentry_sdk.flush() + + ((name, attributes),) = _get_chat_spans(items) + assert name == "chat my-azure-deployment" + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == "my-azure-deployment" + assert attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4o-2024-08-06" + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20