fix(openai-agents): Patch models functions following library refactor #5449
+134
−135
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -66,141 +66,128 @@ def _inject_trace_propagation_headers( | |||||
| headers[key] = value | ||||||
|
|
||||||
|
|
||||||
| def _get_model( | ||||||
| original_get_model: "Callable[..., agents.Model]", | ||||||
| agent: "agents.Agent", | ||||||
| run_config: "agents.RunConfig", | ||||||
| ) -> "agents.Model": | ||||||
| """ | ||||||
| Responsible for | ||||||
| - creating and managing AI client spans. | ||||||
| - adding trace propagation headers to tools with type HostedMCPTool. | ||||||
| - setting the response model on agent invocation spans. | ||||||
| """ | ||||||
| # copy the model to double patching its methods. We use copy on purpose here (instead of deepcopy) | ||||||
| # 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 | ||||||
|
|
||||||
| # Wrap _fetch_response if it exists (for OpenAI models) to capture 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: | ||||||
|
Member
There was a problem hiding this comment. Nitpick: I think an alternative approach that you could take here in order to make this conditional more concise (by removing the
Suggested change
This would also make things consistent with what you have on line 119 within the
Contributor
Author
There was a problem hiding this comment. We need to rethink this response model business because we probably shouldn't be setting response models on agent spans either. This was put in before I knew as much about agents, but there are libraries that let you vary the request model during agent execution as well (in particular Claude Code wrapped by the Claude Agent SDK). Created an issue to track #5458 (and the plan would be to clean up as part of that issue). |
||||||
| agent._sentry_response_model = str(response.model) | ||||||
| return response | ||||||
|
|
||||||
| model._fetch_response = wrapped_fetch_response | ||||||
|
|
||||||
| original_get_response = model.get_response | ||||||
|
|
||||||
| @wraps(original_get_response) | ||||||
| async def wrapped_get_response(*args: "Any", **kwargs: "Any") -> "Any": | ||||||
| mcp_tools = kwargs.get("tools") | ||||||
| hosted_tools = [] | ||||||
| if mcp_tools is not None: | ||||||
| hosted_tools = [ | ||||||
| tool for tool in mcp_tools if isinstance(tool, HostedMCPTool) | ||||||
| ] | ||||||
|
|
||||||
| with ai_client_span(agent, kwargs) 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") | ||||||
|
|
||||||
| _set_response_model_on_agent_span(agent, response_model) | ||||||
| update_ai_client_span(span, result, response_model, agent) | ||||||
|
|
||||||
| return result | ||||||
|
|
||||||
| model.get_response = wrapped_get_response | ||||||
|
|
||||||
| # Also wrap stream_response for streaming support | ||||||
| if hasattr(model, "stream_response"): | ||||||
| original_stream_response = model.stream_response | ||||||
|
|
||||||
| @wraps(original_stream_response) | ||||||
| async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any": | ||||||
| # Uses explicit try/finally instead of context manager to ensure cleanup | ||||||
| # even if the consumer abandons the stream (GeneratorExit). | ||||||
|
Comment on lines
+136
to
+137
Member
There was a problem hiding this comment. This comment looks like it needs to be updated or removed as there's no try/finally below and the code uses a context manager 😅
Contributor
Author
There was a problem hiding this comment. you're completely right, and it doesn't look like we catch I'll keep this commit atomic since we're moving the contents of |
||||||
| span_kwargs = dict(kwargs) | ||||||
| if len(args) > 0: | ||||||
| span_kwargs["system_instructions"] = args[0] | ||||||
| if len(args) > 1: | ||||||
| span_kwargs["input"] = args[1] | ||||||
|
|
||||||
| hosted_tools = [] | ||||||
| if len(args) > 3: | ||||||
| mcp_tools = args[3] | ||||||
|
|
||||||
| if mcp_tools is not None: | ||||||
| hosted_tools = [ | ||||||
| tool for tool in mcp_tools if isinstance(tool, HostedMCPTool) | ||||||
| ] | ||||||
|
|
||||||
| with ai_client_span(agent, span_kwargs) as span: | ||||||
| for hosted_tool in hosted_tools: | ||||||
| _inject_trace_propagation_headers(hosted_tool, span=span) | ||||||
|
|
||||||
| span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) | ||||||
|
|
||||||
| streaming_response = None | ||||||
| ttft_recorded = False | ||||||
| # Capture start time locally to avoid race conditions with concurrent requests | ||||||
| start_time = time.perf_counter() | ||||||
|
|
||||||
| async for event in original_stream_response(*args, **kwargs): | ||||||
| # Detect first content token (text delta event) | ||||||
| if not ttft_recorded and hasattr(event, "delta"): | ||||||
| ttft = time.perf_counter() - start_time | ||||||
| span.set_data( | ||||||
| SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft | ||||||
| ) | ||||||
| ttft_recorded = True | ||||||
|
|
||||||
| # Capture the full response from ResponseCompletedEvent | ||||||
| if hasattr(event, "response"): | ||||||
| streaming_response = event.response | ||||||
| yield event | ||||||
|
|
||||||
| # Update span with response data (usage, output, model) | ||||||
| if streaming_response: | ||||||
| response_model = ( | ||||||
| str(streaming_response.model) | ||||||
| if hasattr(streaming_response, "model") | ||||||
| and streaming_response.model | ||||||
|
Comment on lines
+182
to
+183
Member
There was a problem hiding this comment. Also a nitpick but similar to the comment I left above |
||||||
| else None | ||||||
| ) | ||||||
| _set_response_model_on_agent_span(agent, response_model) | ||||||
| update_ai_client_span( | ||||||
| span, streaming_response, response_model, agent | ||||||
| ) | ||||||
|
|
||||||
| model.stream_response = wrapped_stream_response | ||||||
|
|
||||||
| return model | ||||||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Model patch applied to wrong module
Medium Severity
In the v0.8+ branch, the wrapper is built from
turn_preparation.get_modelbut assigned toagents.run_internal.run_loop.get_model. This leavesturn_preparation.get_modelunpatched, so the new model path can bypass_get_modelinstrumentation inpatches/models.py.