From bffe7e2148103a9b19382c915e1bf5c0657dd95b Mon Sep 17 00:00:00 2001 From: Ansh Rohilla Date: Mon, 19 Oct 2026 17:32:52 +0530 Subject: [PATCH] feat(pydantic_ai): Set request parameters Add gen_ai.request.seed and gen_ai.request.top_k to pydantic_ai model settings recording, with fallback support for extra_body. Update test_model_settings to verify all standard parameters, and add test cases for zero values, extra_body fallback, and omitted settings across streaming and static trace modes. Fixes #8015 --- sentry_sdk/integrations/pydantic_ai/utils.py | 19 +- .../pydantic_ai/test_pydantic_ai.py | 207 ++++++++++++++++++ 2 files changed, 222 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index c847c143b5..110bbd27d7 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -183,21 +183,32 @@ def _set_model_data( "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, + "seed": SPANDATA.GEN_AI_REQUEST_SEED, + "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K, } # ModelSettings is a TypedDict (dict at runtime), so use dict access if isinstance(settings, dict): + extra_body = settings.get("extra_body") for setting_name, spandata_key in settings_map.items(): value = settings.get(setting_name) + if value is None and isinstance(extra_body, dict): + value = extra_body.get(setting_name) if value is not None: set_on_span(spandata_key, value) else: # Fallback for object-style settings + extra_body = getattr(settings, "extra_body", None) for setting_name, spandata_key in settings_map.items(): - if hasattr(settings, setting_name): - value = getattr(settings, setting_name) - if value is not None: - set_on_span(spandata_key, value) + value = ( + getattr(settings, setting_name) + if hasattr(settings, setting_name) + else None + ) + if value is None and isinstance(extra_body, dict): + value = extra_body.get(setting_name) + if value is not None: + set_on_span(spandata_key, value) def _set_available_tools( diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index e100e62671..64ae1f6c80 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -58,6 +58,10 @@ def inner(): temperature=0.7, max_tokens=100, top_p=0.9, + frequency_penalty=0.5, + presence_penalty=0.5, + seed=42, + top_k=50, ), ) @@ -1192,6 +1196,10 @@ async def test_model_settings( assert chat_span["attributes"].get("gen_ai.request.temperature") == 0.7 assert chat_span["attributes"].get("gen_ai.request.max_tokens") == 100 assert chat_span["attributes"].get("gen_ai.request.top_p") == 0.9 + assert chat_span["attributes"].get("gen_ai.request.frequency_penalty") == 0.5 + assert chat_span["attributes"].get("gen_ai.request.presence_penalty") == 0.5 + assert chat_span["attributes"].get("gen_ai.request.seed") == 42 + assert chat_span["attributes"].get("gen_ai.request.top_k") == 50 else: events = capture_events() @@ -1209,6 +1217,205 @@ async def test_model_settings( assert chat_span["data"].get("gen_ai.request.temperature") == 0.7 assert chat_span["data"].get("gen_ai.request.max_tokens") == 100 assert chat_span["data"].get("gen_ai.request.top_p") == 0.9 + assert chat_span["data"].get("gen_ai.request.frequency_penalty") == 0.5 + assert chat_span["data"].get("gen_ai.request.presence_penalty") == 0.5 + assert chat_span["data"].get("gen_ai.request.seed") == 42 + assert chat_span["data"].get("gen_ai.request.top_k") == 50 + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_model_settings_zero_values( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + """ + Test that zero values for model settings are correctly captured. + """ + from pydantic_ai import ModelSettings + + agent = Agent( + "test", + name="test_agent_zero_settings", + model_settings=ModelSettings( + temperature=0.0, + max_tokens=100, + top_p=0.0, + frequency_penalty=0.0, + presence_penalty=0.0, + seed=0, + top_k=0, + ), + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + await agent.run("Test input") + + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert chat_span["attributes"].get("gen_ai.request.temperature") == 0.0 + assert chat_span["attributes"].get("gen_ai.request.max_tokens") == 100 + assert chat_span["attributes"].get("gen_ai.request.top_p") == 0.0 + assert chat_span["attributes"].get("gen_ai.request.frequency_penalty") == 0.0 + assert chat_span["attributes"].get("gen_ai.request.presence_penalty") == 0.0 + assert chat_span["attributes"].get("gen_ai.request.seed") == 0 + assert chat_span["attributes"].get("gen_ai.request.top_k") == 0 + else: + events = capture_events() + + await agent.run("Test input") + + (transaction,) = events + spans = transaction["spans"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert chat_span["data"].get("gen_ai.request.temperature") == 0.0 + assert chat_span["data"].get("gen_ai.request.max_tokens") == 100 + assert chat_span["data"].get("gen_ai.request.top_p") == 0.0 + assert chat_span["data"].get("gen_ai.request.frequency_penalty") == 0.0 + assert chat_span["data"].get("gen_ai.request.presence_penalty") == 0.0 + assert chat_span["data"].get("gen_ai.request.seed") == 0 + assert chat_span["data"].get("gen_ai.request.top_k") == 0 + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_model_settings_fallback_extra_body( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + """ + Test that parameters in extra_body are resolved as fallbacks. + """ + agent = Agent( + "test", + name="test_agent_extra_body", + model_settings={ + "temperature": 0.7, + "extra_body": { + "seed": 999, + "top_k": 25, + }, + }, + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + await agent.run("Test input") + + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert chat_span["attributes"].get("gen_ai.request.temperature") == 0.7 + assert chat_span["attributes"].get("gen_ai.request.seed") == 999 + assert chat_span["attributes"].get("gen_ai.request.top_k") == 25 + else: + events = capture_events() + + await agent.run("Test input") + + (transaction,) = events + spans = transaction["spans"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert chat_span["data"].get("gen_ai.request.temperature") == 0.7 + assert chat_span["data"].get("gen_ai.request.seed") == 999 + assert chat_span["data"].get("gen_ai.request.top_k") == 25 + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_model_settings_omitted( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + """ + Test that omitted model settings are not set on the span. + """ + agent = Agent( + "test", + name="test_agent_omitted", + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + await agent.run("Test input") + + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert "gen_ai.request.seed" not in chat_span["attributes"] + assert "gen_ai.request.top_k" not in chat_span["attributes"] + else: + events = capture_events() + + await agent.run("Test input") + + (transaction,) = events + spans = transaction["spans"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) == 1 + chat_span = chat_spans[0] + + assert "gen_ai.request.seed" not in chat_span["data"] + assert "gen_ai.request.top_k" not in chat_span["data"] @pytest.mark.parametrize("span_streaming", [True, False])