@@ -74,6 +74,24 @@ def test_agent():
74
74
)
75
75
76
76
77
+ @pytest .fixture
78
+ def test_agent_custom_model ():
79
+ """Create a real Agent instance for testing."""
80
+ return Agent (
81
+ name = "test_agent_custom_model" ,
82
+ instructions = "You are a helpful test assistant." ,
83
+ # the model could be agents.OpenAIChatCompletionsModel()
84
+ model = MagicMock (model = "my-custom-model" ),
85
+ model_settings = ModelSettings (
86
+ max_tokens = 100 ,
87
+ temperature = 0.7 ,
88
+ top_p = 1.0 ,
89
+ presence_penalty = 0.0 ,
90
+ frequency_penalty = 0.0 ,
91
+ ),
92
+ )
93
+
94
+
77
95
@pytest .mark .asyncio
78
96
async def test_agent_invocation_span (
79
97
sentry_init , capture_events , test_agent , mock_model_response
@@ -128,6 +146,42 @@ async def test_agent_invocation_span(
128
146
assert ai_client_span ["data" ]["gen_ai.request.top_p" ] == 1.0
129
147
130
148
149
+ @pytest .mark .asyncio
150
+ async def test_client_span_custom_model (
151
+ sentry_init , capture_events , test_agent_custom_model , mock_model_response
152
+ ):
153
+ """
154
+ Test that the integration uses the correct model name if a custom model is used.
155
+ """
156
+
157
+ with patch .dict (os .environ , {"OPENAI_API_KEY" : "test-key" }):
158
+ with patch (
159
+ "agents.models.openai_responses.OpenAIResponsesModel.get_response"
160
+ ) as mock_get_response :
161
+ mock_get_response .return_value = mock_model_response
162
+
163
+ sentry_init (
164
+ integrations = [OpenAIAgentsIntegration ()],
165
+ traces_sample_rate = 1.0 ,
166
+ )
167
+
168
+ events = capture_events ()
169
+
170
+ result = await agents .Runner .run (
171
+ test_agent_custom_model , "Test input" , run_config = test_run_config
172
+ )
173
+
174
+ assert result is not None
175
+ assert result .final_output == "Hello, how can I help you?"
176
+
177
+ (transaction ,) = events
178
+ spans = transaction ["spans" ]
179
+ _ , ai_client_span = spans
180
+
181
+ assert ai_client_span ["description" ] == "chat my-custom-model"
182
+ assert ai_client_span ["data" ]["gen_ai.request.model" ] == "my-custom-model"
183
+
184
+
131
185
def test_agent_invocation_span_sync (
132
186
sentry_init , capture_events , test_agent , mock_model_response
133
187
):
0 commit comments