From 5c3f4ebb23926bd4be9350f3ebae38ade12b6c65 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 8 Apr 2026 13:47:46 +0200 Subject: [PATCH 1/5] fix(pydantic-ai): Stop double reporting model requests --- .../integrations/pydantic_ai/__init__.py | 46 +++++++- .../pydantic_ai/patches/__init__.py | 1 - .../pydantic_ai/patches/graph_nodes.py | 106 ------------------ .../pydantic_ai/test_pydantic_ai.py | 45 ++++---- 4 files changed, 65 insertions(+), 133 deletions(-) delete mode 100644 sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index eb40732d7a..6cf333d304 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -1,17 +1,57 @@ +from functools import wraps from sentry_sdk.integrations import DidNotEnable, Integration - try: import pydantic_ai # type: ignore # noqa: F401 + from pydantic_ai.capabilities.combined import CombinedCapability # type: ignore except ImportError: raise DidNotEnable("pydantic-ai not installed") from .patches import ( _patch_agent_run, - _patch_graph_nodes, _patch_tool_execution, ) +from .spans import ( + ai_client_span, + update_ai_client_span, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Awaitable, Callable + + from pydantic_ai._run_context import RunContext + from pydantic_ai.models import ModelRequestContext + from pydantic_ai.messages import ModelResponse + + +def _patch_wrap_model_request(): + original_wrap_model_request = CombinedCapability.wrap_model_request + + @wraps(original_wrap_model_request) + async def wrapped_wrap_model_request( + self, + ctx: "RunContext[Any]", + *, + request_context: "ModelRequestContext", + handler: "Callable[[ModelRequestContext], Awaitable[ModelResponse]]", + ) -> "Any": + with ai_client_span( + request_context.messages, + None, + request_context.model, + request_context.model_settings, + ) as span: + result = await original_wrap_model_request( + self, ctx, request_context=request_context, handler=handler + ) + + update_ai_client_span(span, result) + return result + + CombinedCapability.wrap_model_request = wrapped_wrap_model_request class PydanticAIIntegration(Integration): @@ -44,5 +84,5 @@ def setup_once() -> None: - Tool executions """ _patch_agent_run() - _patch_graph_nodes() + _patch_wrap_model_request() _patch_tool_execution() diff --git a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py index d0ea6242b4..ad6c22216b 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py @@ -1,3 +1,2 @@ from .agent_run import _patch_agent_run # noqa: F401 -from .graph_nodes import _patch_graph_nodes # noqa: F401 from .tools import _patch_tool_execution # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py deleted file mode 100644 index afb10395f4..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ /dev/null @@ -1,106 +0,0 @@ -from contextlib import asynccontextmanager -from functools import wraps - -from sentry_sdk.integrations import DidNotEnable - -from ..spans import ( - ai_client_span, - update_ai_client_span, -) - -try: - from pydantic_ai._agent_graph import ModelRequestNode # type: ignore -except ImportError: - raise DidNotEnable("pydantic-ai not installed") - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Any, Callable - - -def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]": - """Extract common data needed for creating chat spans. - - Returns: - Tuple of (messages, model, model_settings) - """ - # Extract model and settings from context - model = None - model_settings = None - if hasattr(ctx, "deps"): - model = getattr(ctx.deps, "model", None) - model_settings = getattr(ctx.deps, "model_settings", None) - - # Build full message list: history + current request - messages = [] - if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"): - messages.extend(ctx.state.message_history) - - current_request = getattr(node, "request", None) - if current_request: - messages.append(current_request) - - return messages, model, model_settings - - -def _patch_graph_nodes() -> None: - """ - Patches the graph node execution to create appropriate spans. - - ModelRequestNode -> Creates ai_client span for model requests - CallToolsNode -> Handles tool calls (spans created in tool patching) - """ - - # Patch ModelRequestNode to create ai_client spans - original_model_request_run = ModelRequestNode.run - - @wraps(original_model_request_run) - async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": - messages, model, model_settings = _extract_span_data(self, ctx) - - with ai_client_span(messages, None, model, model_settings) as span: - result = await original_model_request_run(self, ctx) - - # Extract response from result if available - model_response = None - if hasattr(result, "model_response"): - model_response = result.model_response - - update_ai_client_span(span, model_response) - return result - - ModelRequestNode.run = wrapped_model_request_run - - # Patch ModelRequestNode.stream for streaming requests - original_model_request_stream = ModelRequestNode.stream - - def create_wrapped_stream( - original_stream_method: "Callable[..., Any]", - ) -> "Callable[..., Any]": - """Create a wrapper for ModelRequestNode.stream that creates chat spans.""" - - @asynccontextmanager - @wraps(original_stream_method) - async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": - messages, model, model_settings = _extract_span_data(self, ctx) - - # Create chat span for streaming request - with ai_client_span(messages, None, model, model_settings) as span: - # Call the original stream method - async with original_stream_method(self, ctx) as stream: - yield stream - - # After streaming completes, update span with response data - # The ModelRequestNode stores the final response in _result - model_response = None - if hasattr(self, "_result") and self._result is not None: - # _result is a NextNode containing the model_response - if hasattr(self._result, "model_response"): - model_response = self._result.model_response - - update_ai_client_span(span, model_response) - - return wrapped_model_request_stream - - ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 15627a705a..25b608761d 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -75,7 +75,7 @@ async def test_agent_run_async(sentry_init, capture_events, test_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 + assert len(chat_spans) == 1 # Check chat span chat_span = chat_spans[0] @@ -158,7 +158,7 @@ def test_agent_run_sync(sentry_init, capture_events, test_agent): # Find span types chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 # Verify streaming flag is False for sync for chat_span in chat_spans: @@ -192,7 +192,7 @@ async def test_agent_run_stream(sentry_init, capture_events, test_agent): # Find chat spans chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 # Verify streaming flag is True for streaming for chat_span in chat_spans: @@ -231,9 +231,8 @@ async def test_agent_run_stream_events(sentry_init, capture_events, 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 + assert len(chat_spans) == 1 - # run_stream_events uses run() internally, so streaming should be False for chat_span in chat_spans: assert chat_span["data"]["gen_ai.response.streaming"] is False @@ -269,7 +268,7 @@ def add_numbers(a: int, b: int) -> int: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) >= 1 + assert len(tool_spans) == 1 # Check tool span tool_span = tool_spans[0] @@ -342,7 +341,7 @@ def add_numbers(a: int, b: int) -> float: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) >= 1 + assert len(tool_spans) == 2 # Check tool spans model_retry_tool_span = tool_spans[0] @@ -421,7 +420,7 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) >= 1 + assert len(tool_spans) == 1 # Check tool spans model_retry_tool_span = tool_spans[0] @@ -470,7 +469,7 @@ def multiply(a: int, b: int) -> int: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) >= 1 + assert len(tool_spans) == 1 # Verify streaming flag is True for chat_span in chat_spans: @@ -502,7 +501,7 @@ async def test_model_settings(sentry_init, capture_events, test_agent_with_setti # Find chat span chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 chat_span = chat_spans[0] # Check that model settings are captured @@ -548,7 +547,7 @@ async def test_system_prompt_attribute( # 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 len(chat_spans) == 1 chat_span = chat_spans[0] @@ -587,7 +586,7 @@ async def test_error_handling(sentry_init, capture_events): await agent.run("Hello") # At minimum, we should have a transaction - assert len(events) >= 1 + assert len(events) == 1 transaction = [e for e in events if e.get("type") == "transaction"][0] assert transaction["transaction"] == "invoke_agent test_error" # Transaction should complete successfully (status key may not exist if no error) @@ -681,7 +680,7 @@ async def run_agent(input_text): assert transaction["type"] == "transaction" assert transaction["transaction"] == "invoke_agent test_agent" # Each should have its own spans - assert len(transaction["spans"]) >= 1 + assert len(transaction["spans"]) == 1 @pytest.mark.asyncio @@ -721,7 +720,7 @@ async def test_message_history(sentry_init, capture_events): await agent.run("What is my name?", message_history=history) # We should have 2 transactions - assert len(events) >= 2 + assert len(events) == 2 # Check the second transaction has the full history second_transaction = events[1] @@ -755,7 +754,7 @@ async def test_gen_ai_system(sentry_init, capture_events, test_agent): # Find chat span chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 chat_span = chat_spans[0] # gen_ai.system should be set from the model (TestModel -> 'test') @@ -812,7 +811,7 @@ async def test_include_prompts_true(sentry_init, capture_events, test_agent): chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] # Verify that messages are captured in chat spans - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 for chat_span in chat_spans: assert "gen_ai.request.messages" in chat_span["data"] @@ -1242,7 +1241,7 @@ async def test_invoke_agent_with_instructions( # 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 len(chat_spans) == 1 chat_span = chat_spans[0] @@ -1366,7 +1365,7 @@ async def test_usage_data_partial(sentry_init, capture_events): spans = transaction["spans"] chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 # Check that usage data fields exist (they may or may not be set depending on TestModel) chat_span = chat_spans[0] @@ -1461,7 +1460,7 @@ def calc_tool(value: int) -> int: chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] # At least one chat span should exist - assert len(chat_spans) >= 1 + assert len(chat_spans) == 2 # Check if tool calls are captured in response for chat_span in chat_spans: @@ -1509,7 +1508,7 @@ async def test_message_formatting_with_different_parts(sentry_init, capture_even chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] # Should have chat spans - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 # Check that messages are captured chat_span = chat_spans[0] @@ -1781,7 +1780,7 @@ def test_tool(x: int) -> int: chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] # Should have chat spans - assert len(chat_spans) >= 1 + assert len(chat_spans) == 2 @pytest.mark.asyncio @@ -2762,7 +2761,7 @@ async def test_binary_content_in_agent_run(sentry_init, capture_events): (transaction,) = events chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"] - assert len(chat_spans) >= 1 + assert len(chat_spans) == 1 chat_span = chat_spans[0] if "gen_ai.request.messages" in chat_span["data"]: @@ -2906,7 +2905,7 @@ def multiply_numbers(a: int, b: int) -> int: spans = transaction["spans"] tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] - assert len(tool_spans) >= 1 + assert len(tool_spans) == 1 tool_span = tool_spans[0] assert tool_span["data"]["gen_ai.tool.name"] == "multiply_numbers" From cf39d825b1361742045950165b321f97d1623f2a Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 8 Apr 2026 14:38:29 +0200 Subject: [PATCH 2/5] typing --- sentry_sdk/integrations/pydantic_ai/__init__.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 6cf333d304..9107d5b7f8 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -22,17 +22,17 @@ if TYPE_CHECKING: from typing import Any, Awaitable, Callable - from pydantic_ai._run_context import RunContext - from pydantic_ai.models import ModelRequestContext - from pydantic_ai.messages import ModelResponse + from pydantic_ai._run_context import RunContext # type: ignore + from pydantic_ai.models import ModelRequestContext # type: ignore + from pydantic_ai.messages import ModelResponse # type: ignore -def _patch_wrap_model_request(): +def _patch_wrap_model_request() -> None: original_wrap_model_request = CombinedCapability.wrap_model_request @wraps(original_wrap_model_request) async def wrapped_wrap_model_request( - self, + self: "CombinedCapability", ctx: "RunContext[Any]", *, request_context: "ModelRequestContext", From 82901868f1319f96b246d9e59e0e7382c7e87485 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 8 Apr 2026 14:52:51 +0200 Subject: [PATCH 3/5] . --- .../integrations/pydantic_ai/__init__.py | 46 +------ .../pydantic_ai/patches/__init__.py | 1 + .../pydantic_ai/patches/graph_nodes.py | 116 ++++++++++++++++++ 3 files changed, 120 insertions(+), 43 deletions(-) create mode 100644 sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 9107d5b7f8..eb40732d7a 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -1,57 +1,17 @@ -from functools import wraps from sentry_sdk.integrations import DidNotEnable, Integration + try: import pydantic_ai # type: ignore # noqa: F401 - from pydantic_ai.capabilities.combined import CombinedCapability # type: ignore except ImportError: raise DidNotEnable("pydantic-ai not installed") from .patches import ( _patch_agent_run, + _patch_graph_nodes, _patch_tool_execution, ) -from .spans import ( - ai_client_span, - update_ai_client_span, -) - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Any, Awaitable, Callable - - from pydantic_ai._run_context import RunContext # type: ignore - from pydantic_ai.models import ModelRequestContext # type: ignore - from pydantic_ai.messages import ModelResponse # type: ignore - - -def _patch_wrap_model_request() -> None: - original_wrap_model_request = CombinedCapability.wrap_model_request - - @wraps(original_wrap_model_request) - async def wrapped_wrap_model_request( - self: "CombinedCapability", - ctx: "RunContext[Any]", - *, - request_context: "ModelRequestContext", - handler: "Callable[[ModelRequestContext], Awaitable[ModelResponse]]", - ) -> "Any": - with ai_client_span( - request_context.messages, - None, - request_context.model, - request_context.model_settings, - ) as span: - result = await original_wrap_model_request( - self, ctx, request_context=request_context, handler=handler - ) - - update_ai_client_span(span, result) - return result - - CombinedCapability.wrap_model_request = wrapped_wrap_model_request class PydanticAIIntegration(Integration): @@ -84,5 +44,5 @@ def setup_once() -> None: - Tool executions """ _patch_agent_run() - _patch_wrap_model_request() + _patch_graph_nodes() _patch_tool_execution() diff --git a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py index ad6c22216b..d0ea6242b4 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py @@ -1,2 +1,3 @@ from .agent_run import _patch_agent_run # noqa: F401 +from .graph_nodes import _patch_graph_nodes # noqa: F401 from .tools import _patch_tool_execution # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py new file mode 100644 index 0000000000..fa4410d662 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -0,0 +1,116 @@ +from contextlib import asynccontextmanager +from functools import wraps + +from sentry_sdk.integrations import DidNotEnable + +from ..spans import ( + ai_client_span, + update_ai_client_span, +) + +try: + from pydantic_ai._agent_graph import ModelRequestNode # type: ignore +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable + + +def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]": + """Extract common data needed for creating chat spans. + + Returns: + Tuple of (messages, model, model_settings) + """ + # Extract model and settings from context + model = None + model_settings = None + if hasattr(ctx, "deps"): + model = getattr(ctx.deps, "model", None) + model_settings = getattr(ctx.deps, "model_settings", None) + + # Build full message list: history + current request + messages = [] + if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"): + messages.extend(ctx.state.message_history) + + current_request = getattr(node, "request", None) + if current_request: + messages.append(current_request) + + return messages, model, model_settings + + +def _patch_graph_nodes() -> None: + """ + Patches the graph node execution to create appropriate spans. + + ModelRequestNode -> Creates ai_client span for model requests + CallToolsNode -> Handles tool calls (spans created in tool patching) + """ + + # Patch ModelRequestNode to create ai_client spans + original_model_request_run = ModelRequestNode._make_request + + @wraps(original_model_request_run) + async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": + did_stream = getattr(self, "_did_stream", None) + cached_result = getattr(self, "_result", None) + if did_stream or cached_result is not None: + return await original_model_request_run(self, ctx) + + messages, model, model_settings = _extract_span_data(self, ctx) + + with ai_client_span(messages, None, model, model_settings) as span: + result = await original_model_request_run(self, ctx) + + # Extract response from result if available + model_response = None + if hasattr(result, "model_response"): + model_response = result.model_response + + update_ai_client_span(span, model_response) + return result + + ModelRequestNode.run = wrapped_model_request_run + + # Patch ModelRequestNode.stream for streaming requests + original_model_request_stream = ModelRequestNode.stream + + def create_wrapped_stream( + original_stream_method: "Callable[..., Any]", + ) -> "Callable[..., Any]": + """Create a wrapper for ModelRequestNode.stream that creates chat spans.""" + + @asynccontextmanager + @wraps(original_stream_method) + async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": + did_stream = getattr(self, "_did_stream", None) + if did_stream: + async with original_stream_method(self, ctx) as stream: + yield stream + + messages, model, model_settings = _extract_span_data(self, ctx) + + # Create chat span for streaming request + with ai_client_span(messages, None, model, model_settings) as span: + # Call the original stream method + async with original_stream_method(self, ctx) as stream: + yield stream + + # After streaming completes, update span with response data + # The ModelRequestNode stores the final response in _result + model_response = None + if hasattr(self, "_result") and self._result is not None: + # _result is a NextNode containing the model_response + if hasattr(self._result, "model_response"): + model_response = self._result.model_response + + update_ai_client_span(span, model_response) + + return wrapped_model_request_stream + + ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) From ba7b2d475bdc3f291a7968b0b572fe1e27323040 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 8 Apr 2026 14:53:23 +0200 Subject: [PATCH 4/5] . --- sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py index fa4410d662..6e638505a6 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -53,7 +53,7 @@ def _patch_graph_nodes() -> None: """ # Patch ModelRequestNode to create ai_client spans - original_model_request_run = ModelRequestNode._make_request + original_model_request_run = ModelRequestNode.run @wraps(original_model_request_run) async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": From 2f6c932aa11d78ffbe619933763a11c221b50f9a Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 8 Apr 2026 15:20:17 +0200 Subject: [PATCH 5/5] restore tool assertions --- tests/integrations/pydantic_ai/test_pydantic_ai.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 25b608761d..6776a45039 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -341,7 +341,7 @@ def add_numbers(a: int, b: int) -> float: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) == 2 + assert len(tool_spans) >= 1 # Check tool spans model_retry_tool_span = tool_spans[0] @@ -420,7 +420,7 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) == 1 + assert len(tool_spans) >= 1 # Check tool spans model_retry_tool_span = tool_spans[0] @@ -469,7 +469,7 @@ def multiply(a: int, b: int) -> int: tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] # Should have tool spans - assert len(tool_spans) == 1 + assert len(tool_spans) >= 1 # Verify streaming flag is True for chat_span in chat_spans: @@ -2905,7 +2905,7 @@ def multiply_numbers(a: int, b: int) -> int: spans = transaction["spans"] tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] - assert len(tool_spans) == 1 + assert len(tool_spans) >= 1 tool_span = tool_spans[0] assert tool_span["data"]["gen_ai.tool.name"] == "multiply_numbers"