@@ -74,6 +74,24 @@ def test_agent():
7474 )
7575
7676
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+
7795@pytest .mark .asyncio
7896async def test_agent_invocation_span (
7997 sentry_init , capture_events , test_agent , mock_model_response
@@ -128,6 +146,42 @@ async def test_agent_invocation_span(
128146 assert ai_client_span ["data" ]["gen_ai.request.top_p" ] == 1.0
129147
130148
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+
131185def test_agent_invocation_span_sync (
132186 sentry_init , capture_events , test_agent , mock_model_response
133187):
0 commit comments