Skip to content

Commit 1df6c9a

Browse files
authored
Fix custom model name (#4569)
The `model` parameter can be a string or an object of the model itself. If it is a model, get the models name from the `.model` attribute.
1 parent 5709b25 commit 1df6c9a

File tree

3 files changed

+58
-2
lines changed

3 files changed

+58
-2
lines changed

sentry_sdk/integrations/openai_agents/spans/ai_client.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,10 @@
1919
def ai_client_span(agent, get_response_kwargs):
2020
# type: (Agent, dict[str, Any]) -> sentry_sdk.tracing.Span
2121
# TODO-anton: implement other types of operations. Now "chat" is hardcoded.
22+
model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
2223
span = sentry_sdk.start_span(
2324
op=OP.GEN_AI_CHAT,
24-
description=f"chat {agent.model}",
25+
description=f"chat {model_name}",
2526
origin=SPAN_ORIGIN,
2627
)
2728
# TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on

sentry_sdk/integrations/openai_agents/utils.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,8 @@ def _set_agent_data(span, agent):
5353
)
5454

5555
if agent.model:
56-
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, agent.model)
56+
model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
57+
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
5758

5859
if agent.model_settings.presence_penalty:
5960
span.set_data(

tests/integrations/openai_agents/test_openai_agents.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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
7896
async 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+
131185
def test_agent_invocation_span_sync(
132186
sentry_init, capture_events, test_agent, mock_model_response
133187
):

0 commit comments

Comments
 (0)