diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index b15b27f802..489f323679 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -15,7 +15,7 @@ raise DidNotEnable("pydantic-ai not installed") if TYPE_CHECKING: - from typing import Any, Callable, Optional, Union + from typing import Any, Callable, Optional class _StreamingContextManagerWrapper: @@ -37,7 +37,7 @@ def __init__( self.model_settings = model_settings self.is_streaming = is_streaming self._isolation_scope: "Any" = None - self._span: "Optional[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]]" = None + self._span: "Optional[sentry_sdk.traces.StreamedSpan]" = None self._result: "Any" = None async def __aenter__(self) -> "Any": @@ -53,7 +53,6 @@ async def __aenter__(self) -> "Any": self.model_settings, self.is_streaming, ) - self._span.__enter__() # Push agent to contextvar stack after span is successfully created and entered # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 62fb294b14..d71aba33af 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -5,13 +5,9 @@ from sentry_sdk.ai.utils import ( normalize_message_roles, set_data_normalized, - truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing_utils import ( - has_span_streaming_enabled, -) from sentry_sdk.utils import safe_serialize from ..consts import SPAN_ORIGIN @@ -103,9 +99,7 @@ def _get_system_instructions( return permanent_instructions, current_instructions -def _set_input_messages( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any" -) -> None: +def _set_input_messages(span: "StreamedSpan", messages: "Any") -> None: """Set input messages data on a span.""" if not _should_send_inputs(): return @@ -115,24 +109,14 @@ def _set_input_messages( permanent_instructions, current_instructions = _get_system_instructions(messages) if len(permanent_instructions) > 0 or len(current_instructions) > 0: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), - ) - else: - span.set_data( - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), - ) + span.set_attribute( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps( + _transform_system_instructions( + permanent_instructions, current_instructions + ) + ), + ) try: formatted_messages = [] @@ -216,15 +200,11 @@ def _set_input_messages( if formatted_messages: normalized_messages = normalize_message_roles(formatted_messages) - client = sentry_sdk.get_client() - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages(normalized_messages, span, scope) - if not has_span_streaming_enabled(client.options) - else normalized_messages - ) set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + normalized_messages, + unpack=False, ) except Exception: # If we fail to format messages, just skip it @@ -232,7 +212,7 @@ def _set_input_messages( def _set_output_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + span: "StreamedSpan", response: "Optional[ModelResponse]", ) -> None: """Set output data on a span.""" @@ -241,10 +221,7 @@ def _set_output_data( if not response: return - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) # type: ignore[arg-type] + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) # type: ignore[arg-type] if not record_outputs: return @@ -280,7 +257,7 @@ def _set_output_data( parts.append(tool_part) if parts: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_OUTPUT_MESSAGES, json.dumps([{"role": "assistant", "parts": parts}]), ) @@ -292,7 +269,7 @@ def _set_output_data( def ai_client_span( messages: "Any", agent: "Any", model: "Any", model_settings: "Any" -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": +) -> "StreamedSpan": """Create a span for an AI client call (model request). Args: @@ -308,27 +285,15 @@ def ai_client_span( model_name = _get_model_name(model_obj) or "unknown" - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - span = sentry_sdk.traces.start_span( - name=f"chat {model_name}", - attributes={ - "sentry.op": OP.GEN_AI_CHAT, - "sentry.origin": SPAN_ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(), - }, - ) - else: - span = sentry_sdk.start_span( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=SPAN_ORIGIN, - ) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") - # Set streaming flag from contextvar - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) + span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(), + }, + ) _set_agent_data(span, agent) _set_model_data(span, model, model_settings) @@ -345,7 +310,7 @@ def ai_client_span( def update_ai_client_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + span: "StreamedSpan", model_response: "Optional[ModelResponse]", ) -> None: """Update the AI client span with response data.""" diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 03c96455f2..7a8974f072 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -3,14 +3,13 @@ import sentry_sdk from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import safe_serialize from ..consts import SPAN_ORIGIN from ..utils import _set_agent_data, _should_send_inputs, _should_send_outputs if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Optional from pydantic_ai._tool_manager import ToolDefinition # type: ignore @@ -20,7 +19,7 @@ def execute_tool_span( tool_args: "Any", agent: "Any", tool_definition: "Optional[ToolDefinition]" = None, -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": +) -> "StreamedSpan": """Create a span for tool execution. Args: @@ -29,33 +28,18 @@ def execute_tool_span( agent: The agent executing the tool tool_definition: The definition of the tool, if available """ - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - span = sentry_sdk.traces.start_span( - name=f"execute_tool {tool_name}", - attributes={ - "sentry.op": OP.GEN_AI_EXECUTE_TOOL, - "sentry.origin": SPAN_ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", - SPANDATA.GEN_AI_TOOL_NAME: tool_name, - }, - ) - - set_on_span = span.set_attribute - else: - span = sentry_sdk.start_span( - op=OP.GEN_AI_EXECUTE_TOOL, - name=f"execute_tool {tool_name}", - origin=SPAN_ORIGIN, - ) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") - span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) - - set_on_span = span.set_data + span = sentry_sdk.traces.start_span( + name=f"execute_tool {tool_name}", + attributes={ + "sentry.op": OP.GEN_AI_EXECUTE_TOOL, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", + SPANDATA.GEN_AI_TOOL_NAME: tool_name, + }, + ) if tool_definition is not None and hasattr(tool_definition, "description"): - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_definition.description, ) @@ -63,14 +47,12 @@ def execute_tool_span( _set_agent_data(span, agent) if _should_send_inputs() and tool_args is not None: - set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) + span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) return span -def update_execute_tool_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any" -) -> None: +def update_execute_tool_span(span: "StreamedSpan", result: "Any") -> None: """Update the execute tool span with the result.""" if not span: return @@ -78,7 +60,4 @@ def update_execute_tool_span( if not _should_send_outputs() or result is None: return - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) - else: - span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) + span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index 810731b7aa..8839bb8dc3 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -2,16 +2,11 @@ import sentry_sdk from sentry_sdk.ai.utils import ( - get_start_span_function, normalize_message_roles, set_data_normalized, - truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing_utils import ( - has_span_streaming_enabled, -) from ..consts import SPAN_ORIGIN from ..utils import ( @@ -27,7 +22,7 @@ ) if TYPE_CHECKING: - from typing import Any, Union + from typing import Any try: from pydantic_ai.messages import BinaryContent, ImageUrl @@ -42,31 +37,21 @@ def invoke_agent_span( model: "Any", model_settings: "Any", is_streaming: bool = False, -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": +) -> "StreamedSpan": """Create a span for invoking the agent.""" # Determine agent name for span name = "agent" if agent and getattr(agent, "name", None): name = agent.name - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - span = sentry_sdk.traces.start_span( - name=f"invoke_agent {name}", - attributes={ - "sentry.op": OP.GEN_AI_INVOKE_AGENT, - "sentry.origin": SPAN_ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", - }, - ) - else: - span = get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name=f"invoke_agent {name}", - origin=SPAN_ORIGIN, - ) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + span = sentry_sdk.traces.start_span( + name=f"invoke_agent {name}", + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + }, + ) _set_agent_data(span, agent) _set_model_data(span, model, model_settings) @@ -137,22 +122,18 @@ def invoke_agent_span( if messages: normalized_messages = normalize_message_roles(messages) - client = sentry_sdk.get_client() - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages(normalized_messages, span, scope) - if not has_span_streaming_enabled(client.options) - else normalized_messages - ) set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + normalized_messages, + unpack=False, ) return span def update_invoke_agent_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + span: "StreamedSpan", result: "Any", ) -> None: """Update and close the invoke agent span.""" @@ -173,12 +154,7 @@ def update_invoke_agent_span( try: response = result.response if hasattr(response, "model_name") and response.model_name: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name - ) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) except Exception: # If response access fails, continue without setting model name pass diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index c847c143b5..38d2660922 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -12,7 +12,7 @@ ) if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Optional # Store the current agent context in a contextvar for re-entrant safety @@ -90,9 +90,7 @@ def _should_send_outputs() -> bool: return _should_send_prompts_legacy() -def _set_agent_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" -) -> None: +def _set_agent_data(span: "StreamedSpan", agent: "Any") -> None: """Set agent-related data on a span. Args: @@ -106,10 +104,7 @@ def _set_agent_data( agent_obj = get_current_agent() if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) - else: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) def _get_model_name(model_obj: "Any") -> "Optional[str]": @@ -138,7 +133,7 @@ def _get_model_name(model_obj: "Any") -> "Optional[str]": def _set_model_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + span: "StreamedSpan", model: "Any", model_settings: "Any", ) -> None: @@ -157,19 +152,15 @@ def _set_model_data( if not model_obj and agent_obj and hasattr(agent_obj, "model"): model_obj = agent_obj.model - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - if model_obj: # Set system from model if hasattr(model_obj, "system"): - set_on_span(SPANDATA.GEN_AI_SYSTEM, model_obj.system) + span.set_attribute(SPANDATA.GEN_AI_SYSTEM, model_obj.system) # Set model name model_name = _get_model_name(model_obj) if model_name: - set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) # Extract model settings settings = model_settings @@ -190,19 +181,17 @@ def _set_model_data( for setting_name, spandata_key in settings_map.items(): value = settings.get(setting_name) if value is not None: - set_on_span(spandata_key, value) + span.set_attribute(spandata_key, value) else: # Fallback for object-style settings 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) + span.set_attribute(spandata_key, value) -def _set_available_tools( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" -) -> None: +def _set_available_tools(span: "StreamedSpan", agent: "Any") -> None: """Set available tools data on a span from an agent's function toolset. Args: @@ -237,14 +226,10 @@ def _set_available_tools( tools.append(tool_info) if tools: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) - else: - span.set_data( - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) + span.set_attribute( + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + except Exception: # If we can't extract tools, just skip it pass diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 73721dc893..3b4d22a6e4 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -78,14 +78,11 @@ def sync_event_loop(): asyncio.set_event_loop(None) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_run_async( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that the integration creates spans for async agent runs. @@ -94,104 +91,65 @@ async def test_agent_run_async( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - result = await test_agent.run( - ["Message demonstrating the absence of truncation.", "Test input"] - ) - - assert result is not None - assert result.output is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[1]["name"] == "invoke_agent test_agent" - assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" - - assert spans[1]["attributes"]["sentry.op"] == "gen_ai.invoke_agent" - - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - assert len(chat_spans) == 1 - - # Check chat span - chat_span = chat_spans[0] - assert "chat" in chat_span["name"] - assert chat_span["attributes"]["gen_ai.operation.name"] == "chat" - assert chat_span["attributes"]["gen_ai.response.streaming"] is False - assert json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) == [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Message demonstrating the absence of truncation.", - }, - { - "type": "text", - "text": "Test input", - }, - ], - } - ] - assert "gen_ai.usage.input_tokens" in chat_span["attributes"] - assert "gen_ai.usage.output_tokens" in chat_span["attributes"] - else: - events = capture_events() - - result = await test_agent.run("Test input") + result = await test_agent.run( + ["Message demonstrating the absence of truncation.", "Test input"] + ) - assert result is not None - assert result.output is not None + assert result is not None + assert result.output is not None - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify transaction (the transaction IS the invoke_agent span) - assert transaction["transaction"] == "invoke_agent test_agent" - assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" - # The transaction itself should have invoke_agent data - assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent" + assert spans[1]["attributes"]["sentry.op"] == "gen_ai.invoke_agent" - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 - # Check chat span - chat_span = chat_spans[0] - assert "chat" in chat_span["description"] - assert chat_span["data"]["gen_ai.operation.name"] == "chat" - assert chat_span["data"]["gen_ai.response.streaming"] is False - assert "gen_ai.request.messages" in chat_span["data"] - assert "gen_ai.usage.input_tokens" in chat_span["data"] - assert "gen_ai.usage.output_tokens" in chat_span["data"] + # Check chat span + chat_span = chat_spans[0] + assert "chat" in chat_span["name"] + assert chat_span["attributes"]["gen_ai.operation.name"] == "chat" + assert chat_span["attributes"]["gen_ai.response.streaming"] is False + assert json.loads(chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Test input", + }, + ], + } + ] + assert "gen_ai.usage.input_tokens" in chat_span["attributes"] + assert "gen_ai.usage.output_tokens" in chat_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_run_async_model_error( sentry_init, - capture_events, capture_items, - span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) def failing_model(messages, info): @@ -201,43 +159,25 @@ def failing_model(messages, info): FunctionModel(failing_model), name="test_agent", ) + items = capture_items("event", "span") - if span_streaming: - items = capture_items("event", "span") - - with pytest.raises(RuntimeError, match="model exploded"): - await agent.run("Test input") - - (error,) = (item.payload for item in items if item.type == "event") - assert error["level"] == "error" - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert len(spans) == 2 - - assert spans[0]["status"] == "error" - else: - events = capture_events() - - with pytest.raises(RuntimeError, match="model exploded"): - await agent.run("Test input") + with pytest.raises(RuntimeError, match="model exploded"): + await agent.run("Test input") - (error, transaction) = events - assert error["level"] == "error" + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" - spans = transaction["spans"] - assert len(spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert len(spans) == 2 - assert spans[0]["status"] == "internal_error" + assert spans[0]["status"] == "error" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_agent_run_sync( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, sync_event_loop, ): """ @@ -247,71 +187,43 @@ def test_agent_run_sync( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - result = test_agent.run_sync( - ["Message demonstrating the absence of truncation.", "Test input"] - ) - - assert result is not None - assert result.output is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[1]["name"] == "invoke_agent test_agent" - assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" - - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - assert len(chat_spans) == 1 - - # Verify streaming flag is False for sync - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False - else: - events = capture_events() - - result = test_agent.run_sync("Test input") + result = test_agent.run_sync( + ["Message demonstrating the absence of truncation.", "Test input"] + ) - assert result is not None - assert result.output is not None + assert result is not None + assert result.output is not None - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify transaction - assert transaction["transaction"] == "invoke_agent test_agent" - assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" - # Find span types - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 - # Verify streaming flag is False for sync - assert chat_spans[0]["data"]["gen_ai.response.streaming"] is False + # Verify streaming flag is False for sync + assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False -@pytest.mark.parametrize("span_streaming", [True, False]) def test_agent_run_sync_model_error( sentry_init, - capture_events, capture_items, - span_streaming, sync_event_loop, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) def failing_model(messages, info): @@ -321,44 +233,26 @@ def failing_model(messages, info): FunctionModel(failing_model), name="test_agent", ) + items = capture_items("event", "span") - if span_streaming: - items = capture_items("event", "span") + with pytest.raises(RuntimeError, match="model exploded"): + agent.run_sync("Test input") - with pytest.raises(RuntimeError, match="model exploded"): - agent.run_sync("Test input") + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" - (error,) = (item.payload for item in items if item.type == "event") - assert error["level"] == "error" + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert len(spans) == 2 - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert len(spans) == 2 + assert spans[0]["status"] == "error" - assert spans[0]["status"] == "error" - else: - events = capture_events() - - with pytest.raises(RuntimeError, match="model exploded"): - agent.run_sync("Test input") - - (error, transaction) = events - assert error["level"] == "error" - - spans = transaction["spans"] - assert len(spans) == 1 - - assert spans[0]["status"] == "internal_error" - -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_run_stream( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that the integration creates spans for streaming agent runs. @@ -367,96 +261,62 @@ async def test_agent_run_stream( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - async with test_agent.run_stream( - ["Message demonstrating the absence of truncation.", "Test input"] - ) as result: - # Consume the stream - async for _ in result.stream_output(): - pass - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[1]["name"] == "invoke_agent test_agent" - assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" + async with test_agent.run_stream( + ["Message demonstrating the absence of truncation.", "Test input"] + ) as result: + # Consume the stream + async for _ in result.stream_output(): + pass - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - assert len(chat_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify streaming flag is True for streaming - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True - assert json.loads( - chat_spans[0]["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) == [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Message demonstrating the absence of truncation.", - }, - { - "type": "text", - "text": "Test input", - }, - ], - } - ] - assert "gen_ai.usage.input_tokens" in chat_spans[0]["attributes"] - # Streaming responses should still have output data - assert ( - "gen_ai.response.text" in chat_spans[0]["attributes"] - or "gen_ai.response.model" in chat_spans[0]["attributes"] - ) - else: - events = capture_events() + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" - async with test_agent.run_stream("Test input") as result: - # Consume the stream - async for _ in result.stream_output(): - pass + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 - (transaction,) = events - spans = transaction["spans"] - - # Verify transaction - assert transaction["transaction"] == "invoke_agent test_agent" - assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" - - # Find chat spans - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 - - # Verify streaming flag is True for streaming - assert chat_spans[0]["data"]["gen_ai.response.streaming"] is True - assert "gen_ai.request.messages" in chat_spans[0]["data"] - assert "gen_ai.usage.input_tokens" in chat_spans[0]["data"] - # Streaming responses should still have output data - assert ( - "gen_ai.response.text" in chat_spans[0]["data"] - or "gen_ai.response.model" in chat_spans[0]["data"] - ) + # Verify streaming flag is True for streaming + assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True + assert json.loads( + chat_spans[0]["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Test input", + }, + ], + } + ] + assert "gen_ai.usage.input_tokens" in chat_spans[0]["attributes"] + # Streaming responses should still have output data + assert ( + "gen_ai.response.text" in chat_spans[0]["attributes"] + or "gen_ai.response.model" in chat_spans[0]["attributes"] + ) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_run_stream_events( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that run_stream_events creates spans (it uses run internally, so non-streaming). @@ -465,73 +325,44 @@ async def test_agent_run_stream_events( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Consume all events test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - if PYDANTIC_AI_VERSION > (2,): - async with test_agent.run_stream_events( - ["Message demonstrating the absence of truncation.", "Test input"] - ) as stream_events: - async for _ in stream_events: - pass - else: - async for _ in test_agent.run_stream_events( - ["Message demonstrating the absence of truncation.", "Test input"] - ): + if PYDANTIC_AI_VERSION > (2,): + async with test_agent.run_stream_events( + ["Message demonstrating the absence of truncation.", "Test input"] + ) as stream_events: + async for _ in stream_events: pass - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[-1]["name"] == "invoke_agent test_agent" - - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - assert len(chat_spans) == 1 - - # run_stream_events uses run() internally, so streaming should be False - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False else: - events = capture_events() - - if PYDANTIC_AI_VERSION > (2,): - async with test_agent.run_stream_events("Test input") as stream_events: - async for _ in stream_events: - pass - else: - async for _ in test_agent.run_stream_events("Test input"): - pass + async for _ in test_agent.run_stream_events( + ["Message demonstrating the absence of truncation.", "Test input"] + ): + pass - (transaction,) = events + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify transaction - assert transaction["transaction"] == "invoke_agent test_agent" + assert spans[-1]["name"] == "invoke_agent test_agent" - # Find chat spans - spans = transaction["spans"] - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 - # run_stream_events uses run() internally, so streaming should be False - assert chat_spans[0]["data"]["gen_ai.response.streaming"] is False + # run_stream_events uses run() internally, so streaming should be False + assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_with_tools( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that tool execution creates execute_tool spans. @@ -540,8 +371,7 @@ async def test_agent_with_tools( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -551,79 +381,44 @@ def add_numbers(a: int, b: int) -> int: """Add two numbers together.""" return a + b - if span_streaming: - items = capture_items("span") - - result = await test_agent.run("What is 5 + 3?") - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] - - # Should have tool spans - assert len(tool_spans) >= 1 + items = capture_items("span") - # Check tool span - tool_span = tool_spans[0] - assert "execute_tool" in tool_span["name"] - assert tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" - assert tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in tool_span["attributes"] - assert "gen_ai.tool.output" in tool_span["attributes"] - - # Check chat spans have available_tools - for chat_span in chat_spans: - assert "gen_ai.request.available_tools" in chat_span["attributes"] - available_tools_str = chat_span["attributes"][ - "gen_ai.request.available_tools" - ] - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str - else: - events = capture_events() - - result = await test_agent.run("What is 5 + 3?") + result = await test_agent.run("What is 5 + 3?") - assert result is not None + assert result is not None - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + # Find child span types (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # Should have tool spans - assert len(tool_spans) >= 1 + # Should have tool spans + assert len(tool_spans) >= 1 - # Check tool span - tool_span = tool_spans[0] - assert "execute_tool" in tool_span["description"] - assert tool_span["data"]["gen_ai.operation.name"] == "execute_tool" - assert tool_span["data"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in tool_span["data"] - assert "gen_ai.tool.output" in tool_span["data"] + # Check tool span + tool_span = tool_spans[0] + assert "execute_tool" in tool_span["name"] + assert tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" + assert tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in tool_span["attributes"] + assert "gen_ai.tool.output" in tool_span["attributes"] - # Check chat spans have available_tools - for chat_span in chat_spans: - assert "gen_ai.request.available_tools" in chat_span["data"] - available_tools_str = chat_span["data"]["gen_ai.request.available_tools"] - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str + # Check chat spans have available_tools + for chat_span in chat_spans: + assert "gen_ai.request.available_tools" in chat_span["attributes"] + available_tools_str = chat_span["attributes"]["gen_ai.request.available_tools"] + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "handled_tool_call_exceptions", [False, True], @@ -631,11 +426,9 @@ def add_numbers(a: int, b: int) -> int: @pytest.mark.asyncio async def test_agent_with_tool_model_retry( sentry_init, - capture_events, capture_items, get_test_agent, handled_tool_call_exceptions, - span_streaming, ): """ Test that a handled exception is captured when a tool raises ModelRetry. @@ -648,8 +441,7 @@ async def test_agent_with_tool_model_retry( ], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) retries = 0 @@ -665,106 +457,57 @@ def add_numbers(a: int, b: int) -> float: raise ModelRetry(message="Try again with the same arguments.") return a + b - if span_streaming: - items = capture_items("event", "span") - - result = await test_agent.run("What is 5 + 3?") - - assert result is not None - - if handled_tool_call_exceptions: - (error,) = (item.payload for item in items if item.type == "event") - assert error["level"] == "error" - assert error["exception"]["values"][0]["mechanism"]["handled"] + items = capture_items("event", "span") - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] + result = await test_agent.run("What is 5 + 3?") - # Should have tool spans - assert len(tool_spans) >= 1 + assert result is not None - # Check tool spans - model_retry_tool_span = tool_spans[0] - assert "execute_tool" in model_retry_tool_span["name"] - assert ( - model_retry_tool_span["attributes"]["gen_ai.operation.name"] - == "execute_tool" - ) - assert model_retry_tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in model_retry_tool_span["attributes"] + if handled_tool_call_exceptions: + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" + assert error["exception"]["values"][0]["mechanism"]["handled"] - tool_span = tool_spans[1] - assert "execute_tool" in tool_span["name"] - assert tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" - assert tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in tool_span["attributes"] - assert "gen_ai.tool.output" in tool_span["attributes"] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + # Find child span types (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # Check chat spans have available_tools - for chat_span in chat_spans: - assert "gen_ai.request.available_tools" in chat_span["attributes"] - available_tools_str = chat_span["attributes"][ - "gen_ai.request.available_tools" - ] + # Should have tool spans + assert len(tool_spans) >= 1 - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str - else: - events = capture_events() + # Check tool spans + model_retry_tool_span = tool_spans[0] + assert "execute_tool" in model_retry_tool_span["name"] + assert ( + model_retry_tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" + ) + assert model_retry_tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in model_retry_tool_span["attributes"] - result = await test_agent.run("What is 5 + 3?") + tool_span = tool_spans[1] + assert "execute_tool" in tool_span["name"] + assert tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" + assert tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in tool_span["attributes"] + assert "gen_ai.tool.output" in tool_span["attributes"] - assert result is not None + # Check chat spans have available_tools + for chat_span in chat_spans: + assert "gen_ai.request.available_tools" in chat_span["attributes"] + available_tools_str = chat_span["attributes"]["gen_ai.request.available_tools"] - if handled_tool_call_exceptions: - (error, transaction) = events - else: - (transaction,) = events - spans = transaction["spans"] - - if handled_tool_call_exceptions: - assert error["level"] == "error" - assert error["exception"]["values"][0]["mechanism"]["handled"] - - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] - - # Should have tool spans - assert len(tool_spans) >= 1 - - # Check tool spans - model_retry_tool_span = tool_spans[0] - assert "execute_tool" in model_retry_tool_span["description"] - assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool" - assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in model_retry_tool_span["data"] - - tool_span = tool_spans[1] - assert "execute_tool" in tool_span["description"] - assert tool_span["data"]["gen_ai.operation.name"] == "execute_tool" - assert tool_span["data"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in tool_span["data"] - assert "gen_ai.tool.output" in tool_span["data"] - - # Check chat spans have available_tools - for chat_span in chat_spans: - assert "gen_ai.request.available_tools" in chat_span["data"] - available_tools_str = chat_span["data"]["gen_ai.request.available_tools"] - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "handled_tool_call_exceptions", [False, True], @@ -772,11 +515,9 @@ def add_numbers(a: int, b: int) -> float: @pytest.mark.asyncio async def test_agent_with_tool_validation_error( sentry_init, - capture_events, capture_items, get_test_agent, handled_tool_call_exceptions, - span_streaming, ): """ Test that a handled exception is captured when a tool has unsatisfiable constraints. @@ -789,8 +530,7 @@ async def test_agent_with_tool_validation_error( ], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -800,107 +540,59 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: """Add two numbers together.""" return a + b - if span_streaming: - items = capture_items("event", "span") - - result = None - with pytest.raises(UnexpectedModelBehavior): - result = await test_agent.run("What is 5 + 3?") - - assert result is None + items = capture_items("event", "span") - if handled_tool_call_exceptions: - ( - error, - model_behaviour_error, - ) = (item.payload for item in items if item.type == "event") - - assert error["level"] == "error" - assert error["exception"]["values"][0]["mechanism"]["handled"] + result = None + with pytest.raises(UnexpectedModelBehavior): + result = await test_agent.run("What is 5 + 3?") - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] + assert result is None - # Should have tool spans - assert len(tool_spans) >= 1 + if handled_tool_call_exceptions: + ( + error, + model_behaviour_error, + ) = (item.payload for item in items if item.type == "event") - # Check tool spans - model_retry_tool_span = tool_spans[0] - assert "execute_tool" in model_retry_tool_span["name"] - assert ( - model_retry_tool_span["attributes"]["gen_ai.operation.name"] - == "execute_tool" - ) - assert model_retry_tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in model_retry_tool_span["attributes"] + assert error["level"] == "error" + assert error["exception"]["values"][0]["mechanism"]["handled"] - # Check chat spans have available_tools - assert "gen_ai.request.available_tools" in chat_spans[0]["attributes"] - available_tools_str = chat_spans[0]["attributes"][ - "gen_ai.request.available_tools" - ] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str - else: - events = capture_events() + # Should have tool spans + assert len(tool_spans) >= 1 - result = None - with pytest.raises(UnexpectedModelBehavior): - result = await test_agent.run("What is 5 + 3?") + # Check tool spans + model_retry_tool_span = tool_spans[0] + assert "execute_tool" in model_retry_tool_span["name"] + assert ( + model_retry_tool_span["attributes"]["gen_ai.operation.name"] == "execute_tool" + ) + assert model_retry_tool_span["attributes"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in model_retry_tool_span["attributes"] - assert result is None + # Check chat spans have available_tools + assert "gen_ai.request.available_tools" in chat_spans[0]["attributes"] + available_tools_str = chat_spans[0]["attributes"]["gen_ai.request.available_tools"] - if handled_tool_call_exceptions: - (error, model_behaviour_error, transaction) = events - else: - ( - model_behaviour_error, - transaction, - ) = events - spans = transaction["spans"] - - if handled_tool_call_exceptions: - assert error["level"] == "error" - assert error["exception"]["values"][0]["mechanism"]["handled"] - - # Find child span types (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] - - # Should have tool spans - assert len(tool_spans) >= 1 - - # Check tool spans - model_retry_tool_span = tool_spans[0] - assert "execute_tool" in model_retry_tool_span["description"] - assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool" - assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers" - assert "gen_ai.tool.input" in model_retry_tool_span["data"] - - # Check chat spans have available_tools - assert "gen_ai.request.available_tools" in chat_spans[0]["data"] - available_tools_str = chat_spans[0]["data"]["gen_ai.request.available_tools"] - # Available tools is serialized as a string - assert "add_numbers" in available_tools_str + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_with_tools_streaming( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that tool execution works correctly with streaming. @@ -909,8 +601,7 @@ async def test_agent_with_tools_streaming( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -920,72 +611,43 @@ def multiply(a: int, b: int) -> int: """Multiply two numbers.""" return a * b - if span_streaming: - items = capture_items("span") - - async with test_agent.run_stream("What is 7 times 8?") as result: - async for _ in result.stream_output(): - pass - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find span types - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] - - # Should have tool spans - assert len(tool_spans) >= 1 - - # Verify streaming flag is True - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True + items = capture_items("span") - # Check tool span - tool_span = tool_spans[0] - assert tool_span["attributes"]["gen_ai.tool.name"] == "multiply" - assert "gen_ai.tool.input" in tool_span["attributes"] - assert "gen_ai.tool.output" in tool_span["attributes"] - else: - events = capture_events() - - async with test_agent.run_stream("What is 7 times 8?") as result: - async for _ in result.stream_output(): - pass + async with test_agent.run_stream("What is 7 times 8?") as result: + async for _ in result.stream_output(): + pass - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find span types - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + # Find span types + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # Should have tool spans - assert len(tool_spans) >= 1 + # Should have tool spans + assert len(tool_spans) >= 1 - # Verify streaming flag is True - assert chat_spans[0]["data"]["gen_ai.response.streaming"] is True + # Verify streaming flag is True + assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True - # Check tool span - tool_span = tool_spans[0] - assert tool_span["data"]["gen_ai.tool.name"] == "multiply" - assert "gen_ai.tool.input" in tool_span["data"] - assert "gen_ai.tool.output" in tool_span["data"] + # Check tool span + tool_span = tool_spans[0] + assert tool_span["attributes"]["gen_ai.tool.name"] == "multiply" + assert "gen_ai.tool.input" in tool_span["attributes"] + assert "gen_ai.tool.output" in tool_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_model_settings( sentry_init, - capture_events, capture_items, get_test_agent_with_settings, - span_streaming, ): """ Test that model settings are captured in spans. @@ -993,51 +655,30 @@ async def test_model_settings( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent_with_settings = get_test_agent_with_settings() + items = capture_items("span") - if span_streaming: - items = capture_items("span") + await test_agent_with_settings.run("Test input") - await test_agent_with_settings.run("Test input") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find chat span - 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] - # Check that model settings are captured - 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 - else: - events = capture_events() + sentry_sdk.flush() + spans = [item.payload for item in items] - await test_agent_with_settings.run("Test input") - - (transaction,) = events - spans = transaction["spans"] - - # Find chat span - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + # Find chat span + 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] - # Check that model settings are captured - 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 + chat_span = chat_spans[0] + # Check that model settings are captured + 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 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -1050,11 +691,9 @@ async def test_model_settings( ) async def test_system_prompt_attribute( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, - span_streaming, ): """ Test that system prompts are included as the first message. @@ -1069,71 +708,41 @@ async def test_system_prompt_attribute( integrations=[PydanticAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run("Hello") + await agent.run("Hello") - sentry_sdk.flush() - spans = [item.payload for item in items] + sentry_sdk.flush() + spans = [item.payload for item in items] - # The transaction IS the invoke_agent span, check for messages in chat spans instead - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - assert len(chat_spans) == 1 + # The transaction IS the invoke_agent span, check for messages in chat spans instead + 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] + chat_span = chat_spans[0] - if send_default_pii and include_prompts: - system_instructions = chat_span["attributes"][ - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS - ] - assert json.loads(system_instructions) == [ - { - "type": "text", - "content": "You are a helpful assistant specialized in testing.", - } - ] - else: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["attributes"] + if send_default_pii and include_prompts: + system_instructions = chat_span["attributes"][ + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS + ] + assert json.loads(system_instructions) == [ + { + "type": "text", + "content": "You are a helpful assistant specialized in testing.", + } + ] else: - events = capture_events() - - await agent.run("Hello") - - (transaction,) = events - spans = transaction["spans"] - - # The transaction IS the invoke_agent span, check for messages in chat spans instead - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["attributes"] - chat_span = chat_spans[0] - - if send_default_pii and include_prompts: - system_instructions = chat_span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] - assert json.loads(system_instructions) == [ - { - "type": "text", - "content": "You are a helpful assistant specialized in testing.", - } - ] - else: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["data"] - -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_error_handling( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test error handling in agent execution. @@ -1148,45 +757,25 @@ async def test_error_handling( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # Simple run that should succeed - await agent.run("Hello") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[1]["is_segment"] is True - assert spans[1]["status"] != "error" # Could be None or some other status - else: - events = capture_events() - - # Simple run that should succeed - await agent.run("Hello") + # Simple run that should succeed + await agent.run("Hello") - # At minimum, we should have a transaction - assert len(events) == 1 - transaction = [e for e in events if e.get("type") == "transaction"][0] + sentry_sdk.flush() + spans = [item.payload for item in items] - assert transaction["transaction"] == "invoke_agent test_error" - # Transaction should complete successfully (status key may not exist if no error) - trace_status = transaction["contexts"]["trace"].get("status") - assert trace_status != "error" # Could be None or some other status + assert spans[1]["is_segment"] is True + assert spans[1]["status"] != "error" # Could be None or some other status -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_without_pii( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that PII is not captured when send_default_pii is False. @@ -1195,54 +784,32 @@ async def test_without_pii( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=False, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - test_agent = get_test_agent() - await test_agent.run("Sensitive input") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - - # Verify that messages and response text are not captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["attributes"] - assert "gen_ai.response.text" not in span["attributes"] - else: - events = capture_events() - - test_agent = get_test_agent() - await test_agent.run("Sensitive input") + test_agent = get_test_agent() + await test_agent.run("Sensitive input") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + # Find child spans (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # Verify that messages and response text are not captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["data"] - assert "gen_ai.response.text" not in span["data"] + # Verify that messages and response text are not captured + for span in chat_spans: + assert "gen_ai.request.messages" not in span["attributes"] + assert "gen_ai.response.text" not in span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_without_pii_tools( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that tool input/output are not captured when send_default_pii is False. @@ -1251,8 +818,7 @@ async def test_without_pii_tools( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=False, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -1262,50 +828,31 @@ def sensitive_tool(data: str) -> str: """A tool with sensitive data.""" return f"Processed: {data}" - if span_streaming: - items = capture_items("span") + items = capture_items("span") - await test_agent.run("Use sensitive tool with private data") + await test_agent.run("Use sensitive tool with private data") - sentry_sdk.flush() - spans = [item.payload for item in items] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find tool spans - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] - - # If tool was executed, verify input/output are not captured - for tool_span in tool_spans: - assert "gen_ai.tool.input" not in tool_span["attributes"] - assert "gen_ai.tool.output" not in tool_span["attributes"] - else: - events = capture_events() - - await test_agent.run("Use sensitive tool with private data") - - (transaction,) = events - spans = transaction["spans"] - - # Find tool spans - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + # Find tool spans + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # If tool was executed, verify input/output are not captured - for tool_span in tool_spans: - assert "gen_ai.tool.input" not in tool_span["data"] - assert "gen_ai.tool.output" not in tool_span["data"] + # If tool was executed, verify input/output are not captured + for tool_span in tool_spans: + assert "gen_ai.tool.input" not in tool_span["attributes"] + assert "gen_ai.tool.output" not in tool_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_multiple_agents_concurrent( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that multiple agents can run concurrently without interfering. @@ -1313,8 +860,7 @@ async def test_multiple_agents_concurrent( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -1322,44 +868,25 @@ async def test_multiple_agents_concurrent( async def run_agent(input_text): return await test_agent.run(input_text) - if span_streaming: - items = capture_items("span") - - # Run 3 agents concurrently - results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)]) - - assert len(results) == 3 - - sentry_sdk.flush() - spans = [item.payload for item in items] - for span in spans: - if span["is_segment"] is False: - continue - assert span["name"] == "invoke_agent test_agent" - else: - events = capture_events() + items = capture_items("span") - # Run 3 agents concurrently - results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)]) + # Run 3 agents concurrently + results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)]) - assert len(results) == 3 - assert len(events) == 3 + assert len(results) == 3 - # Verify each transaction is separate - for i, transaction in enumerate(events): - assert transaction["type"] == "transaction" - assert transaction["transaction"] == "invoke_agent test_agent" - # Each should have its own spans - assert len(transaction["spans"]) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + if span["is_segment"] is False: + continue + assert span["name"] == "invoke_agent test_agent" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_message_history( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that full conversation history is captured in chat spans. @@ -1373,8 +900,7 @@ async def test_message_history( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Second message with history @@ -1389,59 +915,32 @@ async def test_message_history( model_name="test", ), ] + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # First message - await agent.run("Hello, I'm Alice") - - await agent.run("What is my name?", message_history=history) - - 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" - ] - - if chat_spans: - chat_span = chat_spans[0] - if "gen_ai.request.messages" in chat_span["attributes"]: - messages_data = chat_span["attributes"]["gen_ai.request.messages"] - # Should have multiple messages including history - assert len(messages_data) > 1 - else: - events = capture_events() - - # First message - await agent.run("Hello, I'm Alice") + # First message + await agent.run("Hello, I'm Alice") - await agent.run("What is my name?", message_history=history) + await agent.run("What is my name?", message_history=history) - # We should have 2 transactions - assert len(events) == 2 - - # Check the second transaction has the full history - second_transaction = events[1] - spans = second_transaction["spans"] - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + 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" + ] - if chat_spans: - chat_span = chat_spans[0] - if "gen_ai.request.messages" in chat_span["data"]: - messages_data = chat_span["data"]["gen_ai.request.messages"] - # Should have multiple messages including history - assert len(messages_data) > 1 + if chat_spans: + chat_span = chat_spans[0] + if "gen_ai.request.messages" in chat_span["attributes"]: + messages_data = chat_span["attributes"]["gen_ai.request.messages"] + # Should have multiple messages including history + assert len(messages_data) > 1 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_gen_ai_system( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that gen_ai.system is set from the model. @@ -1449,56 +948,34 @@ async def test_gen_ai_system( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await test_agent.run("Test input") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find chat span - 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] - # gen_ai.system should be set from the model (TestModel -> 'test') - assert "gen_ai.system" in chat_span["attributes"] - assert chat_span["attributes"]["gen_ai.system"] == "test" - else: - events = capture_events() - - await test_agent.run("Test input") + await test_agent.run("Test input") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find chat span - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + # Find chat span + 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] - # gen_ai.system should be set from the model (TestModel -> 'test') - assert "gen_ai.system" in chat_span["data"] - assert chat_span["data"]["gen_ai.system"] == "test" + chat_span = chat_spans[0] + # gen_ai.system should be set from the model (TestModel -> 'test') + assert "gen_ai.system" in chat_span["attributes"] + assert chat_span["attributes"]["gen_ai.system"] == "test" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_include_prompts_false( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that prompts are not captured when include_prompts=False. @@ -1507,54 +984,33 @@ async def test_include_prompts_false( integrations=[PydanticAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=True, # Even with PII enabled, prompts should not be captured - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await test_agent.run("Sensitive prompt") + await test_agent.run("Sensitive prompt") - sentry_sdk.flush() - spans = [item.payload for item in items] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - - # Verify that messages and response text are not captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["attributes"] - assert "gen_ai.response.text" not in span["attributes"] - else: - events = capture_events() - - await test_agent.run("Sensitive prompt") - - (transaction,) = events - spans = transaction["spans"] - - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + # Find child spans (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # Verify that messages and response text are not captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["data"] - assert "gen_ai.response.text" not in span["data"] + # Verify that messages and response text are not captured + for span in chat_spans: + assert "gen_ai.request.messages" not in span["attributes"] + assert "gen_ai.response.text" not in span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_include_prompts_true( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that prompts are captured when include_prompts=True (default). @@ -1563,52 +1019,32 @@ async def test_include_prompts_true( integrations=[PydanticAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await test_agent.run("Test prompt") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - - # Verify that messages are captured in chat spans - assert len(chat_spans) == 1 - assert "gen_ai.request.messages" in chat_spans[0]["attributes"] - else: - events = capture_events() - - await test_agent.run("Test prompt") + await test_agent.run("Test prompt") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + # Find child spans (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # Verify that messages are captured in chat spans - assert len(chat_spans) == 1 - assert "gen_ai.request.messages" in chat_spans[0]["data"] + # Verify that messages are captured in chat spans + assert len(chat_spans) == 1 + assert "gen_ai.request.messages" in chat_spans[0]["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_include_prompts_false_with_tools( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that tool input/output are not captured when include_prompts=False. @@ -1617,8 +1053,7 @@ async def test_include_prompts_false_with_tools( integrations=[PydanticAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -1628,50 +1063,31 @@ def test_tool(value: int) -> int: """A test tool.""" return value * 2 - if span_streaming: - items = capture_items("span") + items = capture_items("span") - await test_agent.run("Use the test tool with value 5") + await test_agent.run("Use the test tool with value 5") - sentry_sdk.flush() - spans = [item.payload for item in items] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find tool spans - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] - - # If tool was executed, verify input/output are not captured - for tool_span in tool_spans: - assert "gen_ai.tool.input" not in tool_span["attributes"] - assert "gen_ai.tool.output" not in tool_span["attributes"] - else: - events = capture_events() - - await test_agent.run("Use the test tool with value 5") - - (transaction,) = events - spans = transaction["spans"] - - # Find tool spans - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + # Find tool spans + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - # If tool was executed, verify input/output are not captured - for tool_span in tool_spans: - assert "gen_ai.tool.input" not in tool_span["data"] - assert "gen_ai.tool.output" not in tool_span["data"] + # If tool was executed, verify input/output are not captured + for tool_span in tool_spans: + assert "gen_ai.tool.input" not in tool_span["attributes"] + assert "gen_ai.tool.output" not in tool_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_include_prompts_requires_pii( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that include_prompts requires send_default_pii=True. @@ -1680,44 +1096,26 @@ async def test_include_prompts_requires_pii( integrations=[PydanticAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=False, # PII disabled - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await test_agent.run("Test prompt") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - - # Even with include_prompts=True, if PII is disabled, messages should not be captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["attributes"] - assert "gen_ai.response.text" not in span["attributes"] - else: - events = capture_events() - - await test_agent.run("Test prompt") + await test_agent.run("Test prompt") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - # Find child spans (invoke_agent is the transaction, not a child span) - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + # Find child spans (invoke_agent is the transaction, not a child span) + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # Even with include_prompts=True, if PII is disabled, messages should not be captured - for span in chat_spans: - assert "gen_ai.request.messages" not in span["data"] - assert "gen_ai.response.text" not in span["data"] + # Even with include_prompts=True, if PII is disabled, messages should not be captured + for span in chat_spans: + assert "gen_ai.request.messages" not in span["attributes"] + assert "gen_ai.response.text" not in span["attributes"] @pytest.mark.asyncio @@ -1730,7 +1128,6 @@ async def test_context_cleanup_after_run(sentry_init, get_test_agent): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Verify context is not set before run @@ -1754,7 +1151,6 @@ def test_context_cleanup_after_run_sync(sentry_init, get_test_agent, sync_event_ sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Verify context is not set before run @@ -1779,7 +1175,6 @@ async def test_context_cleanup_after_streaming(sentry_init, get_test_agent): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Verify context is not set before run @@ -1806,7 +1201,6 @@ async def test_context_cleanup_on_error(sentry_init, get_test_agent): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) test_agent = get_test_agent() @@ -1841,7 +1235,6 @@ async def test_context_isolation_concurrent_agents(sentry_init, get_test_agent): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Create a second agent @@ -1883,13 +1276,10 @@ async def run_and_check_context(agent, agent_name): # ==================== Additional Coverage Tests ==================== -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_invoke_agent_with_list_user_prompt( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent span handles list user prompts correctly. @@ -1903,42 +1293,22 @@ async def test_invoke_agent_with_list_user_prompt( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # Use a list as user prompt - await agent.run(["First part", "Second part"]) + # Use a list as user prompt + await agent.run(["First part", "Second part"]) - sentry_sdk.flush() - spans = [item.payload for item in items] + sentry_sdk.flush() + spans = [item.payload for item in items] - if "gen_ai.request.messages" in spans[0]["attributes"]: - messages_str = spans[0]["attributes"]["gen_ai.request.messages"] - assert "First part" in messages_str - assert "Second part" in messages_str - else: - events = capture_events() - - # Use a list as user prompt - await agent.run(["First part", "Second part"]) - - (transaction,) = events + if "gen_ai.request.messages" in spans[0]["attributes"]: + messages_str = spans[0]["attributes"]["gen_ai.request.messages"] + assert "First part" in messages_str + assert "Second part" in messages_str - # Check that the invoke_agent transaction has messages data - # The invoke_agent is the transaction itself - if "gen_ai.request.messages" in transaction["contexts"]["trace"]["data"]: - messages_str = transaction["contexts"]["trace"]["data"][ - "gen_ai.request.messages" - ] - assert "First part" in messages_str - assert "Second part" in messages_str - -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -1951,11 +1321,9 @@ async def test_invoke_agent_with_list_user_prompt( ) async def test_invoke_agent_with_instructions( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, - span_streaming, ): """ Test that invoke_agent span handles instructions correctly. @@ -1976,65 +1344,36 @@ async def test_invoke_agent_with_instructions( integrations=[PydanticAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") + await agent.run("Test input") - await agent.run("Test input") + sentry_sdk.flush() + spans = [item.payload for item in items] - sentry_sdk.flush() - spans = [item.payload for item in items] - - # The transaction IS the invoke_agent span, check for messages in chat spans instead - 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] + # The transaction IS the invoke_agent span, check for messages in chat spans instead + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) == 1 - if send_default_pii and include_prompts: - system_instructions = chat_span["attributes"][ - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS - ] - assert json.loads(system_instructions) == [ - {"type": "text", "content": "System prompt"}, - { - "type": "text", - "content": f"Instruction 1{instructions_separator}Instruction 2", - }, - ] - else: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["attributes"] + chat_span = chat_spans[0] + if send_default_pii and include_prompts: + system_instructions = chat_span["attributes"][ + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS + ] + assert json.loads(system_instructions) == [ + {"type": "text", "content": "System prompt"}, + { + "type": "text", + "content": f"Instruction 1{instructions_separator}Instruction 2", + }, + ] else: - events = capture_events() - - await agent.run("Test input") - - (transaction,) = events - spans = transaction["spans"] - - # The transaction IS the invoke_agent span, check for messages in chat spans instead - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 - - chat_span = chat_spans[0] - - if send_default_pii and include_prompts: - system_instructions = chat_span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] - assert json.loads(system_instructions) == [ - {"type": "text", "content": "System prompt"}, - { - "type": "text", - "content": f"Instruction 1{instructions_separator}Instruction 2", - }, - ] - else: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["data"] + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["attributes"] @pytest.mark.asyncio @@ -2051,7 +1390,6 @@ async def test_model_name_extraction_with_callable( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Test the utility function directly @@ -2081,7 +1419,6 @@ async def test_model_name_extraction_fallback_to_str( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Test the utility function directly @@ -2113,36 +1450,28 @@ async def test_model_settings_object_style( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Create mock settings object (not a dict) - mock_settings = MagicMock() - mock_settings.temperature = 0.8 - mock_settings.max_tokens = 200 - mock_settings.top_p = 0.95 - mock_settings.frequency_penalty = 0.5 - mock_settings.presence_penalty = 0.3 + span = sentry_sdk.traces.start_span(name="test") - # Set model data with object-style settings - _set_model_data(span, None, mock_settings) + # Create mock settings object (not a dict) + mock_settings = MagicMock() + mock_settings.temperature = 0.8 + mock_settings.max_tokens = 200 + mock_settings.top_p = 0.95 + mock_settings.frequency_penalty = 0.5 + mock_settings.presence_penalty = 0.3 - span.finish() + # Set model data with object-style settings + _set_model_data(span, None, mock_settings) - # Should not crash and should set the settings - assert transaction is not None + span.end() -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_usage_data_partial( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that usage data is correctly handled when only some fields are present. @@ -2155,30 +1484,18 @@ async def test_usage_data_partial( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run("Test input") - - sentry_sdk.flush() - spans = [item.payload for item in items] + await agent.run("Test input") - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] - else: - events = capture_events() - - await agent.run("Test input") - - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] assert len(chat_spans) == 1 @@ -2188,13 +1505,10 @@ async def test_usage_data_partial( assert chat_span is not None -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_data_from_scope( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that agent data can be retrieved from Sentry scope when not passed directly. @@ -2208,41 +1522,24 @@ async def test_agent_data_from_scope( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # The integration automatically sets agent in scope during execution - await agent.run("Test input") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert spans[1]["name"] == "invoke_agent test_scope_agent" - else: - events = capture_events() - - # The integration automatically sets agent in scope during execution - await agent.run("Test input") + # The integration automatically sets agent in scope during execution + await agent.run("Test input") - # Verify agent name is capture - (transaction,) = events + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify agent name is captured - assert transaction["transaction"] == "invoke_agent test_scope_agent" + assert spans[1]["name"] == "invoke_agent test_scope_agent" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_available_tools_without_description( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that available tools are captured even when description is missing. @@ -2250,8 +1547,7 @@ async def test_available_tools_without_description( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -2261,46 +1557,28 @@ def tool_without_desc(x: int) -> int: # No docstring = no description return x * 2 - if span_streaming: - items = capture_items("span") - - await test_agent.run("Use the tool with 5") - - 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" - ] - if chat_spans: - chat_span = chat_spans[0] - if "gen_ai.request.available_tools" in chat_span["attributes"]: - tools_str = chat_span["attributes"]["gen_ai.request.available_tools"] - assert "tool_without_desc" in tools_str - else: - events = capture_events() + items = capture_items("span") - await test_agent.run("Use the tool with 5") + await test_agent.run("Use the tool with 5") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - if chat_spans: - chat_span = chat_spans[0] - if "gen_ai.request.available_tools" in chat_span["data"]: - tools_str = chat_span["data"]["gen_ai.request.available_tools"] - assert "tool_without_desc" in tools_str + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + if chat_spans: + chat_span = chat_spans[0] + if "gen_ai.request.available_tools" in chat_span["attributes"]: + tools_str = chat_span["attributes"]["gen_ai.request.available_tools"] + assert "tool_without_desc" in tools_str -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_output_with_tool_calls( sentry_init, - capture_events, capture_items, get_test_agent, - span_streaming, ): """ Test that tool calls in model response are captured correctly. @@ -2309,8 +1587,7 @@ async def test_output_with_tool_calls( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_agent = get_test_agent() @@ -2320,51 +1597,30 @@ def calc_tool(value: int) -> int: """Calculate something.""" return value + 10 - if span_streaming: - items = capture_items("span") - - await test_agent.run("Use calc_tool with 5") - - 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" - ] - - # At least one chat span should exist - assert len(chat_spans) >= 1 - - # Check if tool calls are captured in response - # Tool calls may or may not be in response depending on TestModel behavior - # Just verify the span was created and has basic data - assert "gen_ai.operation.name" in chat_spans[0]["attributes"] - else: - events = capture_events() + items = capture_items("span") - await test_agent.run("Use calc_tool with 5") + await test_agent.run("Use calc_tool with 5") - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # At least one chat span should exist - assert len(chat_spans) >= 1 + # At least one chat span should exist + assert len(chat_spans) >= 1 - # Check if tool calls are captured in response - # Tool calls may or may not be in response depending on TestModel behavior - # Just verify the span was created and has basic data - assert "gen_ai.operation.name" in chat_spans[0]["data"] + # Check if tool calls are captured in response + # Tool calls may or may not be in response depending on TestModel behavior + # Just verify the span was created and has basic data + assert "gen_ai.operation.name" in chat_spans[0]["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_message_formatting_with_different_parts( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that different message part types are handled correctly in ai_client span. @@ -2380,8 +1636,7 @@ async def test_message_formatting_with_different_parts( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Create message history with different part types @@ -2394,46 +1649,25 @@ async def test_message_formatting_with_different_parts( model_name="test", ), ] + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run("What did I say?", message_history=history) + await agent.run("What did I say?", message_history=history) - sentry_sdk.flush() - spans = [item.payload for item in items] + 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" - ] - - # Should have chat spans - assert len(chat_spans) == 1 - - # Check that messages are captured - chat_span = chat_spans[0] - if "gen_ai.request.messages" in chat_span["attributes"]: - messages_data = chat_span["attributes"]["gen_ai.request.messages"] - assert messages_data is not None - else: - events = capture_events() - - await agent.run("What did I say?", message_history=history) - - (transaction,) = events - spans = transaction["spans"] - - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] - # Should have chat spans - assert len(chat_spans) == 1 + # Should have chat spans + assert len(chat_spans) == 1 - # Check that messages are captured - chat_span = chat_spans[0] - if "gen_ai.request.messages" in chat_span["data"]: - messages_data = chat_span["data"]["gen_ai.request.messages"] - # Should contain message history - assert messages_data is not None + # Check that messages are captured + chat_span = chat_spans[0] + if "gen_ai.request.messages" in chat_span["attributes"]: + messages_data = chat_span["attributes"]["gen_ai.request.messages"] + assert messages_data is not None @pytest.mark.asyncio @@ -2452,19 +1686,14 @@ async def test_update_invoke_agent_span_with_none_output( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Update with None output - should not raise - update_invoke_agent_span(span, None) + span = sentry_sdk.traces.start_span(name="test_span") - span.finish() + # Update with None output - should not raise + update_invoke_agent_span(span, None) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2482,28 +1711,20 @@ async def test_update_ai_client_span_with_none_response( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Update with None response - should not raise - update_ai_client_span(span, None) + span = sentry_sdk.traces.start_span(name="test_span") - span.finish() + # Update with None response - should not raise + update_ai_client_span(span, None) - # Should not crash - assert transaction is not None + span.end() -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_agent_without_name( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that agent without a name is handled correctly. @@ -2514,31 +1735,16 @@ async def test_agent_without_name( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run("Test input") - - sentry_sdk.flush() - spans = [item.payload for item in items] - - assert "invoke_agent" in spans[1]["name"] - else: - events = capture_events() - - await agent.run("Test input") - - (transaction,) = events + await agent.run("Test input") - # Should still create transaction, just with default name - assert transaction["type"] == "transaction" + sentry_sdk.flush() + spans = [item.payload for item in items] - # Transaction name should be "invoke_agent agent" or similar default - assert "invoke_agent" in transaction["transaction"] + assert "invoke_agent" in spans[1]["name"] @pytest.mark.asyncio @@ -2554,22 +1760,17 @@ async def test_input_messages_error_handling( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Pass invalid messages that would cause an error - invalid_messages = [object()] # Plain object without expected attributes + # Pass invalid messages that would cause an error + invalid_messages = [object()] # Plain object without expected attributes - # Should not raise, error is caught internally - _set_input_messages(span, invalid_messages) + # Should not raise, error is caught internally + _set_input_messages(span, invalid_messages) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2587,23 +1788,18 @@ async def test_available_tools_error_handling( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create mock agent with invalid toolset - mock_agent = MagicMock() - mock_agent._function_toolset.tools.items.side_effect = Exception("Error") + # Create mock agent with invalid toolset + mock_agent = MagicMock() + mock_agent._function_toolset.tools.items.side_effect = Exception("Error") - # Should not raise, error is caught internally - _set_available_tools(span, mock_agent) + # Should not raise, error is caught internally + _set_available_tools(span, mock_agent) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2619,19 +1815,14 @@ async def test_set_usage_data_with_none_usage( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Pass None usage - should not raise - _set_usage_data(span, None) + # Pass None usage - should not raise + _set_usage_data(span, None) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2649,34 +1840,26 @@ async def test_set_usage_data_with_partial_fields( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create usage object with only some fields - mock_usage = MagicMock() - mock_usage.input_tokens = 100 - mock_usage.output_tokens = None # Missing - mock_usage.total_tokens = 100 + # Create usage object with only some fields + mock_usage = MagicMock() + mock_usage.input_tokens = 100 + mock_usage.output_tokens = None # Missing + mock_usage.total_tokens = 100 - # Should only set the non-None fields - _set_usage_data(span, mock_usage) + # Should only set the non-None fields + _set_usage_data(span, mock_usage) - span.finish() - - # Should not crash - assert transaction is not None + span.end() -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_message_parts_with_tool_return( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that ToolReturnPart messages are handled correctly. @@ -2697,32 +1880,19 @@ def test_tool(x: int) -> int: integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # Run with history containing tool return - await agent.run("Use test_tool with 5") - - 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" - ] - else: - events = capture_events() + # Run with history containing tool return + await agent.run("Use test_tool with 5") - # Run with history containing tool return - await agent.run("Use test_tool with 5") + sentry_sdk.flush() + spans = [item.payload for item in items] - (transaction,) = events - spans = transaction["spans"] - - chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] # Should have chat spans assert len(chat_spans) == 2 @@ -2743,37 +1913,29 @@ async def test_message_parts_with_list_content( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create message with list content - mock_msg = MagicMock() - mock_part = MagicMock() - mock_part.content = ["item1", "item2", {"complex": "item"}] - mock_msg.parts = [mock_part] - mock_msg.instructions = None + # Create message with list content + mock_msg = MagicMock() + mock_part = MagicMock() + mock_part.content = ["item1", "item2", {"complex": "item"}] + mock_msg.parts = [mock_part] + mock_msg.instructions = None - messages = [mock_msg] + messages = [mock_msg] - # Should handle list content - _set_input_messages(span, messages) + # Should handle list content + _set_input_messages(span, messages) - span.finish() - - # Should not crash - assert transaction is not None + span.end() -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_output_data_transformations( sentry_init, capture_items, - capture_events, - span_streaming, ): """ Test transformation of the model response from `Hooks.on.after_model_request`. @@ -2782,8 +1944,7 @@ async def test_output_data_transformations( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) def response_model(messages, info): @@ -2806,88 +1967,46 @@ def multiply(a: int, b: int) -> int: """Multiply two numbers.""" return a * b - if span_streaming: - items = capture_items("span") - - await agent.run("What is 5 times 3?") - sentry_sdk.flush() + items = capture_items("span") - spans = [item.payload for item in items] + await agent.run("What is 5 times 3?") + sentry_sdk.flush() - invoke_agent_span = next( - span - for span in spans - if span["attributes"].get("sentry.op") == "gen_ai.invoke_agent" - ) - assert invoke_agent_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] == ( - "The answer is 15." - ) - - chat_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == "gen_ai.chat" - ] - assert json.loads( - chat_spans[0]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] - ) == [ - { - "role": "assistant", - "parts": [ - { - "type": "tool_call", - "name": "multiply", - "arguments": '{"a": 5, "b": 3}', - } - ], - } - ] - assert json.loads( - chat_spans[1]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] - ) == [ - { - "role": "assistant", - "parts": [ - {"type": "reasoning", "content": "5 times 3 is 15."}, - {"type": "text", "content": "The answer is 15."}, - ], - } - ] - else: - events = capture_events() - - await agent.run("What is 5 times 3?") + spans = [item.payload for item in items] - (transaction,) = events - assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent" - assert transaction["contexts"]["trace"]["data"][ - SPANDATA.GEN_AI_RESPONSE_TEXT - ] == ("The answer is 15.") + invoke_agent_span = next( + span + for span in spans + if span["attributes"].get("sentry.op") == "gen_ai.invoke_agent" + ) + assert invoke_agent_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] == ( + "The answer is 15." + ) - chat_spans = [ - span for span in transaction["spans"] if span["op"] == "gen_ai.chat" - ] - assert json.loads(chat_spans[0]["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ - { - "role": "assistant", - "parts": [ - { - "type": "tool_call", - "name": "multiply", - "arguments": '{"a": 5, "b": 3}', - } - ], - } - ] - assert json.loads(chat_spans[1]["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ - { - "role": "assistant", - "parts": [ - {"type": "reasoning", "content": "5 times 3 is 15."}, - {"type": "text", "content": "The answer is 15."}, - ], - } - ] + chat_spans = [ + span for span in spans if span["attributes"].get("sentry.op") == "gen_ai.chat" + ] + assert json.loads(chat_spans[0]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ], + } + ] + assert json.loads(chat_spans[1]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "5 times 3 is 15."}, + {"type": "text", "content": "The answer is 15."}, + ], + } + ] @pytest.mark.asyncio @@ -2906,24 +2025,19 @@ async def test_output_data_error_handling( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create mock response that will cause error - mock_response = MagicMock() - mock_response.model_name = "test-model" - mock_response.parts = [MagicMock(side_effect=Exception("Error"))] + # Create mock response that will cause error + mock_response = MagicMock() + mock_response.model_name = "test-model" + mock_response.parts = [MagicMock(side_effect=Exception("Error"))] - # Should catch error and not crash - _set_output_data(span, mock_response) + # Should catch error and not crash + _set_output_data(span, mock_response) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2943,28 +2057,23 @@ async def test_message_with_system_prompt_part( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create message with SystemPromptPart - system_part = messages.SystemPromptPart(content="You are a helpful assistant") + # Create message with SystemPromptPart + system_part = messages.SystemPromptPart(content="You are a helpful assistant") - mock_msg = MagicMock() - mock_msg.parts = [system_part] - mock_msg.instructions = None + mock_msg = MagicMock() + mock_msg.parts = [system_part] + mock_msg.instructions = None - msgs = [mock_msg] + msgs = [mock_msg] - # Should handle system prompt - _set_input_messages(span, msgs) + # Should handle system prompt + _set_input_messages(span, msgs) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -2982,28 +2091,23 @@ async def test_message_with_instructions( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create message with instructions - mock_msg = MagicMock() - mock_msg.instructions = "System instructions here" - mock_part = MagicMock() - mock_part.content = "User message" - mock_msg.parts = [mock_part] + # Create message with instructions + mock_msg = MagicMock() + mock_msg.instructions = "System instructions here" + mock_part = MagicMock() + mock_part.content = "User message" + mock_msg.parts = [mock_part] - msgs = [mock_msg] + msgs = [mock_msg] - # Should extract system prompt from instructions - _set_input_messages(span, msgs) + # Should extract system prompt from instructions + _set_input_messages(span, msgs) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3019,20 +2123,15 @@ async def test_set_input_messages_without_prompts( integrations=[PydanticAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Even with messages, should not set them - messages = ["test"] - _set_input_messages(span, messages) + # Even with messages, should not set them + messages = ["test"] + _set_input_messages(span, messages) - span.finish() - - # Should not crash and should not set messages - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3049,7 +2148,6 @@ async def test_get_model_name_with_exception_in_callable( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Create model with callable name that raises exception @@ -3075,7 +2173,6 @@ async def test_get_model_name_with_string_model( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Pass a string as model @@ -3097,7 +2194,6 @@ async def test_get_model_name_with_none( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Pass None @@ -3122,24 +2218,19 @@ async def test_set_model_data_with_system( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Create model with system - mock_model = MagicMock() - mock_model.system = "openai" - mock_model.model_name = "gpt-4" + span = sentry_sdk.traces.start_span(name="test") - # Set model data - _set_model_data(span, mock_model, None) + # Create model with system + mock_model = MagicMock() + mock_model.system = "openai" + mock_model.model_name = "gpt-4" - span.finish() + # Set model data + _set_model_data(span, mock_model, None) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3157,27 +2248,22 @@ async def test_set_model_data_from_agent_scope( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Set agent in scope - scope = sentry_sdk.get_current_scope() - mock_agent = MagicMock() - mock_agent.model = MagicMock() - mock_agent.model.model_name = "test-model" - mock_agent.model_settings = {"temperature": 0.5} - scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - - span = sentry_sdk.start_span(op="test_span") + # Set agent in scope + scope = sentry_sdk.get_current_scope() + mock_agent = MagicMock() + mock_agent.model = MagicMock() + mock_agent.model.model_name = "test-model" + mock_agent.model_settings = {"temperature": 0.5} + scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - # Pass None for model, should get from scope - _set_model_data(span, None, None) + span = sentry_sdk.traces.start_span(name="test_span") - span.finish() + # Pass None for model, should get from scope + _set_model_data(span, None, None) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3193,26 +2279,21 @@ async def test_set_model_data_with_none_settings_values( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Create settings with None values - settings = { - "temperature": 0.7, - "max_tokens": None, # Should be skipped - "top_p": None, # Should be skipped - } + span = sentry_sdk.traces.start_span(name="test") - # Set model data - _set_model_data(span, None, settings) + # Create settings with None values + settings = { + "temperature": 0.7, + "max_tokens": None, # Should be skipped + "top_p": None, # Should be skipped + } - span.finish() + # Set model data + _set_model_data(span, None, settings) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3231,7 +2312,6 @@ async def test_should_send_prompts_without_pii( integrations=[PydanticAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=False, # PII disabled, - stream_gen_ai_spans=False, ) # Should return False @@ -3252,19 +2332,14 @@ async def test_set_agent_data_without_agent( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Pass None agent, with no agent in scope - _set_agent_data(span, None) + span = sentry_sdk.traces.start_span(name="test_span") - span.finish() + # Pass None agent, with no agent in scope + _set_agent_data(span, None) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3282,25 +2357,20 @@ async def test_set_agent_data_from_scope( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Set agent in scope - scope = sentry_sdk.get_current_scope() - mock_agent = MagicMock() - mock_agent.name = "test_agent_from_scope" - scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - - span = sentry_sdk.start_span(op="test_span") + # Set agent in scope + scope = sentry_sdk.get_current_scope() + mock_agent = MagicMock() + mock_agent.name = "test_agent_from_scope" + scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - # Pass None for agent, should get from scope - _set_agent_data(span, None) + span = sentry_sdk.traces.start_span(name="test_span") - span.finish() + # Pass None for agent, should get from scope + _set_agent_data(span, None) - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3318,23 +2388,18 @@ async def test_set_agent_data_without_name( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create agent without name - mock_agent = MagicMock() - mock_agent.name = None # No name + # Create agent without name + mock_agent = MagicMock() + mock_agent.name = None # No name - # Should not set agent name - _set_agent_data(span, mock_agent) + # Should not set agent name + _set_agent_data(span, mock_agent) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3352,23 +2417,18 @@ async def test_set_available_tools_without_toolset( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create agent without _function_toolset - mock_agent = MagicMock() - del mock_agent._function_toolset + # Create agent without _function_toolset + mock_agent = MagicMock() + del mock_agent._function_toolset - # Should handle gracefully - _set_available_tools(span, mock_agent) + # Should handle gracefully + _set_available_tools(span, mock_agent) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3386,29 +2446,24 @@ async def test_set_available_tools_with_schema( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + span = sentry_sdk.traces.start_span(name="test_span") - # Create agent with toolset containing schema - mock_agent = MagicMock() - mock_tool = MagicMock() - mock_schema = MagicMock() - mock_schema.description = "Test tool description" - mock_schema.json_schema = {"type": "object", "properties": {}} - mock_tool.function_schema = mock_schema + # Create agent with toolset containing schema + mock_agent = MagicMock() + mock_tool = MagicMock() + mock_schema = MagicMock() + mock_schema.description = "Test tool description" + mock_schema.json_schema = {"type": "object", "properties": {}} + mock_tool.function_schema = mock_schema - mock_agent._function_toolset.tools = {"test_tool": mock_tool} + mock_agent._function_toolset.tools = {"test_tool": mock_tool} - # Should extract schema - _set_available_tools(span, mock_agent) + # Should extract schema + _set_available_tools(span, mock_agent) - span.finish() - - # Should not crash - assert transaction is not None + span.end() @pytest.mark.asyncio @@ -3418,7 +2473,6 @@ async def test_execute_tool_span_creation( """ Test direct creation of execute_tool span. """ - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( execute_tool_span, update_execute_tool_span, @@ -3428,17 +2482,12 @@ async def test_execute_tool_span_creation( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create execute_tool span - with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: - # Update with result - update_execute_tool_span(span, {"result": "success"}) - - # Should not crash - assert transaction is not None + # Create execute_tool span + with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: + # Update with result + update_execute_tool_span(span, {"result": "success"}) @pytest.mark.asyncio @@ -3448,7 +2497,6 @@ async def test_execute_tool_span_with_mcp_type( """ Test execute_tool span with MCP tool type. """ - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( execute_tool_span, ) @@ -3457,17 +2505,12 @@ async def test_execute_tool_span_with_mcp_type( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create execute_tool span with mcp type - with execute_tool_span("mcp_tool", {"arg": "value"}, None, "mcp") as span: - # Verify type is set - assert span is not None - - # Should not crash - assert transaction is not None + # Create execute_tool span with mcp type + with execute_tool_span("mcp_tool", {"arg": "value"}, None, "mcp") as span: + # Verify type is set + assert span is not None @pytest.mark.asyncio @@ -3477,7 +2520,6 @@ async def test_execute_tool_span_without_prompts( """ Test that execute_tool span respects _should_send_prompts(). """ - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( execute_tool_span, update_execute_tool_span, @@ -3487,17 +2529,12 @@ async def test_execute_tool_span_without_prompts( integrations=[PydanticAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create execute_tool span - with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: - # Update with result - should not set input/output - update_execute_tool_span(span, {"result": "success"}) - - # Should not crash - assert transaction is not None + # Create execute_tool span + with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: + # Update with result - should not set input/output + update_execute_tool_span(span, {"result": "success"}) @pytest.mark.asyncio @@ -3507,23 +2544,17 @@ async def test_execute_tool_span_with_none_args( """ Test execute_tool span with None args. """ - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import execute_tool_span sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create execute_tool span with None args - with execute_tool_span("test_tool", None, None, "function") as span: - assert span is not None - - # Should not crash - assert transaction is not None + # Create execute_tool span with None args + with execute_tool_span("test_tool", None, None, "function") as span: + assert span is not None @pytest.mark.asyncio @@ -3540,7 +2571,6 @@ async def test_update_execute_tool_span_with_none_span( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Update with None span - should not raise @@ -3557,7 +2587,6 @@ async def test_update_execute_tool_span_with_none_result( """ Test that update_execute_tool_span handles None result gracefully. """ - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( execute_tool_span, update_execute_tool_span, @@ -3567,17 +2596,12 @@ async def test_update_execute_tool_span_with_none_result( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create execute_tool span - with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: - # Update with None result - update_execute_tool_span(span, None) - - # Should not crash - assert transaction is not None + # Create execute_tool span + with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span: + # Update with None result + update_execute_tool_span(span, None) @pytest.mark.asyncio @@ -3593,7 +2617,6 @@ async def test_tool_execution_without_span_context( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Create a simple agent with no tools (won't have function_toolset) @@ -3621,32 +2644,26 @@ async def test_invoke_agent_span_with_callable_instruction( """ from unittest.mock import MagicMock - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create mock agent with callable instruction - mock_agent = MagicMock() - mock_agent.name = "test_agent" - mock_agent._system_prompts = [] + # Create mock agent with callable instruction + mock_agent = MagicMock() + mock_agent.name = "test_agent" + mock_agent._system_prompts = [] - # Add both string and callable instructions - mock_callable = lambda: "Dynamic instruction" - mock_agent._instructions = ["Static instruction", mock_callable] + # Add both string and callable instructions + mock_callable = lambda: "Dynamic instruction" + mock_agent._instructions = ["Static instruction", mock_callable] - # Create span - span = invoke_agent_span("Test prompt", mock_agent, None, None) - span.finish() - - # Should not crash (callable should be skipped) - assert transaction is not None + # Create span + span = invoke_agent_span("Test prompt", mock_agent, None, None) + span.end() @pytest.mark.asyncio @@ -3658,29 +2675,23 @@ async def test_invoke_agent_span_with_string_instructions( """ from unittest.mock import MagicMock - import sentry_sdk from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Create mock agent with string instruction - mock_agent = MagicMock() - mock_agent.name = "test_agent" - mock_agent._system_prompts = [] - mock_agent._instructions = "Single instruction string" + # Create mock agent with string instruction + mock_agent = MagicMock() + mock_agent.name = "test_agent" + mock_agent._system_prompts = [] + mock_agent._instructions = "Single instruction string" - # Create span - span = invoke_agent_span("Test prompt", mock_agent, None, None) - span.finish() - - # Should not crash - assert transaction is not None + # Create span + span = invoke_agent_span("Test prompt", mock_agent, None, None) + span.end() @pytest.mark.asyncio @@ -3696,20 +2707,15 @@ async def test_ai_client_span_with_streaming_flag( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Set streaming flag in scope - scope = sentry_sdk.get_current_scope() - scope._contexts["pydantic_ai_agent"] = {"_streaming": True} - - # Create ai_client span - span = ai_client_span([], None, None, None) - span.finish() + # Set streaming flag in scope + scope = sentry_sdk.get_current_scope() + scope._contexts["pydantic_ai_agent"] = {"_streaming": True} - # Should not crash - assert transaction is not None + # Create ai_client span + span = ai_client_span([], None, None, None) + span.end() @pytest.mark.asyncio @@ -3727,24 +2733,19 @@ async def test_ai_client_span_gets_agent_from_scope( sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - # Set agent in scope - scope = sentry_sdk.get_current_scope() - mock_agent = MagicMock() - mock_agent.name = "test_agent" - mock_agent._function_toolset = MagicMock() - mock_agent._function_toolset.tools = {} - scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - - # Create ai_client span without passing agent - span = ai_client_span([], None, None, None) - span.finish() + # Set agent in scope + scope = sentry_sdk.get_current_scope() + mock_agent = MagicMock() + mock_agent.name = "test_agent" + mock_agent._function_toolset = MagicMock() + mock_agent._function_toolset.tools = {} + scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent} - # Should not crash - assert transaction is not None + # Create ai_client span without passing agent + span = ai_client_span([], None, None, None) + span.end() def _get_messages_from_span(span_data): @@ -3769,165 +2770,91 @@ def _find_binary_content(messages_data, expected_modality, expected_mime_type): return False -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_binary_content_encoding_image( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test that BinaryContent with image data is properly encoded in messages.""" sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span( - name="test", attributes={"sentry.op": "test"} - ): - span = sentry_sdk.traces.start_span( - name="custom span", attributes={"sentry.op": "test_span"} - ) - binary_content = BinaryContent( - data=b"fake_image_data_12345", media_type="image/png" - ) - user_part = UserPromptPart(content=["Look at this image:", binary_content]) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None - - _set_input_messages(span, [mock_msg]) - span.finish() - - sentry_sdk.flush() - spans = [item.payload for item in items] - - span_data = spans[0]["attributes"] - messages_data = _get_messages_from_span(span_data) - assert _find_binary_content(messages_data, "image", "image/png") - else: - events = capture_events() + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.op": "test_span"} + ) + binary_content = BinaryContent( + data=b"fake_image_data_12345", media_type="image/png" + ) + user_part = UserPromptPart(content=["Look at this image:", binary_content]) + mock_msg = MagicMock() + mock_msg.parts = [user_part] + mock_msg.instructions = None - with sentry_sdk.start_transaction(op="test", name="test"): - span = sentry_sdk.start_span(op="test_span") - binary_content = BinaryContent( - data=b"fake_image_data_12345", media_type="image/png" - ) - user_part = UserPromptPart(content=["Look at this image:", binary_content]) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None + _set_input_messages(span, [mock_msg]) + span.end() - _set_input_messages(span, [mock_msg]) - span.finish() + sentry_sdk.flush() + spans = [item.payload for item in items] - (event,) = events - span_data = event["spans"][0]["data"] - messages_data = _get_messages_from_span(span_data) - assert _find_binary_content(messages_data, "image", "image/png") + span_data = spans[0]["attributes"] + messages_data = _get_messages_from_span(span_data) + assert _find_binary_content(messages_data, "image", "image/png") -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_binary_content_encoding_mixed_content( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test that BinaryContent mixed with text content is properly handled.""" sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span( - name="test", attributes={"sentry.op": "test"} - ): - span = sentry_sdk.traces.start_span( - name="custom span", attributes={"sentry.op": "test_span"} - ) - binary_content = BinaryContent( - data=b"fake_image_bytes", media_type="image/jpeg" - ) - user_part = UserPromptPart( - content=["Here is an image:", binary_content, "What do you see?"] - ) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None - - _set_input_messages(span, [mock_msg]) - span.finish() + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.op": "test_span"} + ) + binary_content = BinaryContent(data=b"fake_image_bytes", media_type="image/jpeg") + user_part = UserPromptPart( + content=["Here is an image:", binary_content, "What do you see?"] + ) + mock_msg = MagicMock() + mock_msg.parts = [user_part] + mock_msg.instructions = None - sentry_sdk.flush() - spans = [item.payload for item in items] + _set_input_messages(span, [mock_msg]) + span.end() - span_data = spans[0]["attributes"] - messages_data = _get_messages_from_span(span_data) + sentry_sdk.flush() + spans = [item.payload for item in items] - # Verify both text and binary content are present - found_text = any( - content_item.get("type") == "text" - for msg in messages_data - if "content" in msg - for content_item in msg["content"] - ) - assert found_text, "Text content should be found" - assert _find_binary_content(messages_data, "image", "image/jpeg") - else: - events = capture_events() + span_data = spans[0]["attributes"] + messages_data = _get_messages_from_span(span_data) - with sentry_sdk.start_transaction(op="test", name="test"): - span = sentry_sdk.start_span(op="test_span") - binary_content = BinaryContent( - data=b"fake_image_bytes", media_type="image/jpeg" - ) - user_part = UserPromptPart( - content=["Here is an image:", binary_content, "What do you see?"] - ) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None - - _set_input_messages(span, [mock_msg]) - span.finish() - - (event,) = events - span_data = event["spans"][0]["data"] - messages_data = _get_messages_from_span(span_data) - - # Verify both text and binary content are present - found_text = any( - content_item.get("type") == "text" - for msg in messages_data - if "content" in msg - for content_item in msg["content"] - ) - assert found_text, "Text content should be found" - assert _find_binary_content(messages_data, "image", "image/jpeg") + # Verify both text and binary content are present + found_text = any( + content_item.get("type") == "text" + for msg in messages_data + if "content" in msg + for content_item in msg["content"] + ) + assert found_text, "Text content should be found" + assert _find_binary_content(messages_data, "image", "image/jpeg") -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_binary_content_in_agent_run( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test that BinaryContent in actual agent run is properly captured in spans.""" agent = Agent("test", name="test_binary_agent") @@ -3936,112 +2863,62 @@ async def test_binary_content_in_agent_run( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) binary_content = BinaryContent( data=b"fake_image_data_for_testing", media_type="image/png" ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run(["Analyze this image:", binary_content]) - - 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] - if "gen_ai.request.messages" in chat_span["attributes"]: - messages_str = str(chat_span["attributes"]["gen_ai.request.messages"]) - - assert any( - keyword in messages_str for keyword in ["blob", "image", "base64"] - ) - else: - events = capture_events() + await agent.run(["Analyze this image:", binary_content]) - await agent.run(["Analyze this image:", binary_content]) + 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 - (transaction,) = events - chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 + chat_span = chat_spans[0] + if "gen_ai.request.messages" in chat_span["attributes"]: + messages_str = str(chat_span["attributes"]["gen_ai.request.messages"]) - chat_span = chat_spans[0] - if "gen_ai.request.messages" in chat_span["data"]: - messages_str = str(chat_span["data"]["gen_ai.request.messages"]) - assert any( - keyword in messages_str for keyword in ["blob", "image", "base64"] - ) + assert any(keyword in messages_str for keyword in ["blob", "image", "base64"]) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_set_usage_data_with_cache_tokens( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test that cache_read_tokens and cache_write_tokens are tracked.""" sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span( - name="test", attributes={"sentry.op": "test"} - ): - span = sentry_sdk.traces.start_span( - name="custom span", attributes={"sentry.op": "test_span"} - ) - usage = RequestUsage( - input_tokens=100, - output_tokens=50, - cache_read_tokens=80, - cache_write_tokens=20, - ) - _set_usage_data(span, usage) - span.finish() - - sentry_sdk.flush() - spans = [item.payload for item in items] + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.op": "test_span"} + ) + usage = RequestUsage( + input_tokens=100, + output_tokens=50, + cache_read_tokens=80, + cache_write_tokens=20, + ) + _set_usage_data(span, usage) + span.end() - assert spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 - assert ( - spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 - ) - else: - events = capture_events() - - with sentry_sdk.start_transaction(op="test", name="test"): - span = sentry_sdk.start_span(op="test_span") - usage = RequestUsage( - input_tokens=100, - output_tokens=50, - cache_read_tokens=80, - cache_write_tokens=20, - ) - _set_usage_data(span, usage) - span.finish() + sentry_sdk.flush() + spans = [item.payload for item in items] - (event,) = events - (span_data,) = event["spans"] - assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 - assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 + assert spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 + assert spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "url,image_url_kwargs,expected_content", [ @@ -4085,12 +2962,10 @@ async def test_set_usage_data_with_cache_tokens( ) def test_image_url_base64_content_in_span( sentry_init, - capture_events, capture_items, url, image_url_kwargs, expected_content, - span_streaming, ): from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span @@ -4098,69 +2973,40 @@ def test_image_url_base64_content_in_span( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) found_image = False - if span_streaming: - items = capture_items("span") - - with sentry_sdk.start_transaction(op="test", name="test"): - image_url = ImageUrl(url=url, **image_url_kwargs) - user_part = UserPromptPart(content=["Look at this image:", image_url]) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None - - span = ai_client_span([mock_msg], None, None, None) - span.finish() - - 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 - messages_data = _get_messages_from_span(chat_spans[0]["attributes"]) - - for msg in messages_data: - if "content" not in msg: - continue - for content_item in msg["content"]: - if content_item.get("type") == "image": - found_image = True - assert content_item["content"] == expected_content - else: - events = capture_events() - - with sentry_sdk.start_transaction(op="test", name="test"): - image_url = ImageUrl(url=url, **image_url_kwargs) - user_part = UserPromptPart(content=["Look at this image:", image_url]) - mock_msg = MagicMock() - mock_msg.parts = [user_part] - mock_msg.instructions = None - - span = ai_client_span([mock_msg], None, None, None) - span.finish() - - (event,) = events - chat_spans = [s for s in event["spans"] if s["op"] == "gen_ai.chat"] - assert len(chat_spans) == 1 - messages_data = _get_messages_from_span(chat_spans[0]["data"]) - - for msg in messages_data: - if "content" not in msg: - continue - for content_item in msg["content"]: - if content_item.get("type") == "image": - found_image = True - assert content_item["content"] == expected_content + items = capture_items("span") + + image_url = ImageUrl(url=url, **image_url_kwargs) + user_part = UserPromptPart(content=["Look at this image:", image_url]) + mock_msg = MagicMock() + mock_msg.parts = [user_part] + mock_msg.instructions = None + + span = ai_client_span([mock_msg], None, None, None) + span.end() + + 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 + messages_data = _get_messages_from_span(chat_spans[0]["attributes"]) + + for msg in messages_data: + if "content" not in msg: + continue + for content_item in msg["content"]: + if content_item.get("type") == "image": + found_image = True + assert content_item["content"] == expected_content assert found_image, "Image content item should be found in messages data" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( "url, image_url_kwargs, expected_content", @@ -4193,74 +3039,48 @@ def test_image_url_base64_content_in_span( ) async def test_invoke_agent_image_url( sentry_init, - capture_events, capture_items, url, image_url_kwargs, expected_content, - span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) agent = Agent("test", name="test_image_url_agent") image_url = ImageUrl(url=url, **image_url_kwargs) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - await agent.run([image_url, "Describe this image"]) + await agent.run([image_url, "Describe this image"]) - found_image = False - - 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" - ] - messages_data = _get_messages_from_span(chat_spans[0]["attributes"]) - for msg in messages_data: - if "content" not in msg: - continue - for content_item in msg["content"]: - if content_item.get("type") == "image": - assert content_item["content"] == expected_content - found_image = True - else: - events = capture_events() - - await agent.run([image_url, "Describe this image"]) - - (transaction,) = events - - found_image = False + found_image = False - chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"] - messages_data = _get_messages_from_span(chat_spans[0]["data"]) - for msg in messages_data: - if "content" not in msg: - continue - for content_item in msg["content"]: - if content_item.get("type") == "image": - assert content_item["content"] == expected_content - found_image = True + 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" + ] + messages_data = _get_messages_from_span(chat_spans[0]["attributes"]) + for msg in messages_data: + if "content" not in msg: + continue + for content_item in msg["content"]: + if content_item.get("type") == "image": + assert content_item["content"] == expected_content + found_image = True assert found_image, "Image content item should be found in messages data" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_tool_description_in_execute_tool_span( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that tool description from the tool's docstring is included in execute_tool spans. @@ -4280,80 +3100,51 @@ def multiply_numbers(a: int, b: int) -> int: integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - result = await agent.run("What is 5 times 3?") - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - - tool_spans = [ - s - for s in spans - if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" - ] - - assert len(tool_spans) >= 1 - - tool_span = tool_spans[0] - - assert tool_span["attributes"]["gen_ai.tool.name"] == "multiply_numbers" - assert SPANDATA.GEN_AI_TOOL_DESCRIPTION in tool_span["attributes"] - assert ( - "Multiply two numbers" - in tool_span["attributes"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] - ) - else: - events = capture_events() - - result = await agent.run("What is 5 times 3?") - assert result is not None + result = await agent.run("What is 5 times 3?") + assert result is not None - (transaction,) = events - spans = transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] - tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + tool_spans = [ + s + for s in spans + if s["attributes"].get("sentry.op", "") == "gen_ai.execute_tool" + ] - assert len(tool_spans) >= 1 + assert len(tool_spans) >= 1 - tool_span = tool_spans[0] + tool_span = tool_spans[0] - assert tool_span["data"]["gen_ai.tool.name"] == "multiply_numbers" - assert SPANDATA.GEN_AI_TOOL_DESCRIPTION in tool_span["data"] - assert ( - "Multiply two numbers" - in tool_span["data"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] - ) + assert tool_span["attributes"]["gen_ai.tool.name"] == "multiply_numbers" + assert SPANDATA.GEN_AI_TOOL_DESCRIPTION in tool_span["attributes"] + assert ( + "Multiply two numbers" + in tool_span["attributes"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] + ) -def _spans_by_op(items, events, streaming): +def _spans_by_op(items): """Normalize captured spans to a list of (op, data) tuples. Works for both the span-streaming/gen-AI-span-streaming payloads and the classic transaction payload so data collection assertions can be shared. """ - if streaming: - sentry_sdk.flush() - return [ - ( - item.payload["attributes"].get("sentry.op", ""), - item.payload["attributes"], - ) - for item in items - if item.type == "span" - ] - - (transaction,) = events - return [(span["op"], span["data"]) for span in transaction["spans"]] + sentry_sdk.flush() + return [ + ( + item.payload["attributes"].get("sentry.op", ""), + item.payload["attributes"], + ) + for item in items + if item.type == "span" + ] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expect_inputs,expect_available_tools", [ @@ -4426,14 +3217,12 @@ async def test_data_collection_gen_ai_inputs_gates_request_messages_tool_inputs_ include_prompts, expect_inputs, expect_available_tools, - span_streaming, ): init_kwargs = { "integrations": [PydanticAIIntegration(include_prompts=include_prompts)], "traces_sample_rate": 1.0, "send_default_pii": send_default_pii, - "trace_lifecycle": "stream" if span_streaming else "static", - "stream_gen_ai_spans": False, + "trace_lifecycle": "stream", } if data_collection is not None: init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -4446,18 +3235,12 @@ async def test_data_collection_gen_ai_inputs_gates_request_messages_tool_inputs_ def add_numbers(a: int, b: int) -> int: return a + b - streaming = span_streaming - if streaming: - items = capture_items("span") - events = None - else: - items = None - events = capture_events() + items = capture_items("span") result = await test_agent.run("What is 5 + 3?") assert result is not None - spans = _spans_by_op(items, events, streaming) + spans = _spans_by_op(items) chat_spans = [data for op, data in spans if op == "gen_ai.chat"] tool_spans = [data for op, data in spans if op == "gen_ai.execute_tool"] @@ -4499,7 +3282,6 @@ def add_numbers(a: int, b: int) -> int: assert tool_span[SPANDATA.GEN_AI_TOOL_NAME] == "add_numbers" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expect_outputs", [ @@ -4564,14 +3346,12 @@ async def test_data_collection_gen_ai_outputs_gates_response_text_and_tool_outpu send_default_pii, include_prompts, expect_outputs, - span_streaming, ): init_kwargs = { "integrations": [PydanticAIIntegration(include_prompts=include_prompts)], "traces_sample_rate": 1.0, "send_default_pii": send_default_pii, - "trace_lifecycle": "stream" if span_streaming else "static", - "stream_gen_ai_spans": False, + "trace_lifecycle": "stream", } if data_collection is not None: init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -4584,32 +3364,18 @@ async def test_data_collection_gen_ai_outputs_gates_response_text_and_tool_outpu def add_numbers(a: int, b: int) -> int: return a + b - streaming = span_streaming - if streaming: - items = capture_items("transaction", "span") - events = None - else: - items = None - events = capture_events() + items = capture_items("span") result = await test_agent.run("What is 5 + 3?") assert result is not None - spans = _spans_by_op(items, events, streaming) + spans = _spans_by_op(items) # The invoke_agent span is either a child span or, when it is the segment # span, the transaction itself. invoke_agent_data = next( (data for op, data in spans if op == "gen_ai.invoke_agent"), None ) - if invoke_agent_data is None: - if streaming: - (transaction,) = ( - item.payload for item in items if item.type == "transaction" - ) - else: - (transaction,) = events - invoke_agent_data = transaction["contexts"]["trace"]["data"] chat_spans = [data for op, data in spans if op == "gen_ai.chat"] tool_spans = [data for op, data in spans if op == "gen_ai.execute_tool"] @@ -4647,7 +3413,6 @@ def add_numbers(a: int, b: int) -> int: assert tool_span[SPANDATA.GEN_AI_TOOL_NAME] == "add_numbers" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "gen_ai,expect_outputs", [ @@ -4681,14 +3446,12 @@ async def test_data_collection_gen_ai_output_message_parts_follow_outputs_gate( get_test_agent, gen_ai, expect_outputs, - span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", _experiments={"data_collection": {"gen_ai": gen_ai}}, - stream_gen_ai_spans=False, ) test_agent = get_test_agent() @@ -4697,18 +3460,12 @@ async def test_data_collection_gen_ai_output_message_parts_follow_outputs_gate( def add_numbers(a: int, b: int) -> int: return a + b - streaming = span_streaming - if streaming: - items = capture_items("transaction", "span") - events = None - else: - items = None - events = capture_events() + items = capture_items("transaction", "span") result = await test_agent.run("What is 5 + 3?") assert result is not None - spans = _spans_by_op(items, events, streaming) + spans = _spans_by_op(items) chat_spans = [data for op, data in spans if op == "gen_ai.chat"] # The test model calls the tool on the first response and answers with text @@ -4740,14 +3497,12 @@ def add_numbers(a: int, b: int) -> int: assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in chat_span -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_data_collection_gen_ai_request_messages_keep_tool_returns_when_outputs_disabled( sentry_init, capture_events, capture_items, get_test_agent, - span_streaming, ): """ A tool return value is an output on the `gen_ai.execute_tool` span, so @@ -4759,11 +3514,10 @@ async def test_data_collection_gen_ai_request_messages_keep_tool_returns_when_ou sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", _experiments={ "data_collection": {"gen_ai": {"inputs": True, "outputs": False}} }, - stream_gen_ai_spans=False, ) test_agent = get_test_agent() @@ -4772,18 +3526,12 @@ async def test_data_collection_gen_ai_request_messages_keep_tool_returns_when_ou def add_numbers(a: int, b: int) -> int: return a + b - streaming = span_streaming - if streaming: - items = capture_items("transaction", "span") - events = None - else: - items = None - events = capture_events() + items = capture_items("transaction", "span") result = await test_agent.run("What is 5 + 3?") assert result is not None - spans = _spans_by_op(items, events, streaming) + spans = _spans_by_op(items) chat_spans = [data for op, data in spans if op == "gen_ai.chat"] tool_spans = [data for op, data in spans if op == "gen_ai.execute_tool"]