|
12 | 12 | from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration |
13 | 13 | from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages |
14 | 14 | from sentry_sdk.integrations.pydantic_ai.spans.utils import _set_usage_data |
| 15 | +from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import execute_tool_span |
15 | 16 |
|
16 | 17 | from pydantic_ai import Agent |
17 | 18 | from pydantic_ai.messages import BinaryContent, UserPromptPart |
18 | 19 | from pydantic_ai.usage import RequestUsage |
19 | | -from pydantic_ai.models.test import TestModel |
20 | 20 | from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior |
| 21 | +from pydantic_ai._tool_manager import ToolDefinition |
21 | 22 |
|
22 | 23 |
|
23 | 24 | @pytest.fixture |
@@ -2794,3 +2795,142 @@ async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events): |
2794 | 2795 | (span_data,) = event["spans"] |
2795 | 2796 | assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 |
2796 | 2797 | assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 |
| 2798 | + |
| 2799 | + |
| 2800 | +@pytest.mark.asyncio |
| 2801 | +async def test_tool_description_in_execute_tool_span(sentry_init, capture_events): |
| 2802 | + """ |
| 2803 | + Test that tool description from the tool's docstring is included in execute_tool spans. |
| 2804 | + """ |
| 2805 | + agent = Agent( |
| 2806 | + "test", |
| 2807 | + name="test_agent", |
| 2808 | + system_prompt="You are a helpful test assistant.", |
| 2809 | + ) |
| 2810 | + |
| 2811 | + @agent.tool_plain |
| 2812 | + def multiply_numbers(a: int, b: int) -> int: |
| 2813 | + """Multiply two numbers and return the product.""" |
| 2814 | + return a * b |
| 2815 | + |
| 2816 | + sentry_init( |
| 2817 | + integrations=[PydanticAIIntegration()], |
| 2818 | + traces_sample_rate=1.0, |
| 2819 | + send_default_pii=True, |
| 2820 | + ) |
| 2821 | + |
| 2822 | + events = capture_events() |
| 2823 | + |
| 2824 | + result = await agent.run("What is 5 times 3?") |
| 2825 | + assert result is not None |
| 2826 | + |
| 2827 | + (transaction,) = events |
| 2828 | + spans = transaction["spans"] |
| 2829 | + |
| 2830 | + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] |
| 2831 | + assert len(tool_spans) >= 1 |
| 2832 | + |
| 2833 | + tool_span = tool_spans[0] |
| 2834 | + assert tool_span["data"]["gen_ai.tool.name"] == "multiply_numbers" |
| 2835 | + assert SPANDATA.GEN_AI_TOOL_DESCRIPTION in tool_span["data"] |
| 2836 | + assert "Multiply two numbers" in tool_span["data"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] |
| 2837 | + |
| 2838 | + |
| 2839 | +@pytest.mark.asyncio |
| 2840 | +async def test_tool_without_description_omits_tool_description( |
| 2841 | + sentry_init, capture_events |
| 2842 | +): |
| 2843 | + """ |
| 2844 | + Test that execute_tool spans omit tool description when the tool has no docstring. |
| 2845 | + """ |
| 2846 | + agent = Agent( |
| 2847 | + "test", |
| 2848 | + name="test_agent", |
| 2849 | + system_prompt="You are a helpful test assistant.", |
| 2850 | + ) |
| 2851 | + |
| 2852 | + @agent.tool_plain |
| 2853 | + def no_docs_tool(a: int, b: int) -> int: |
| 2854 | + return a + b |
| 2855 | + |
| 2856 | + sentry_init( |
| 2857 | + integrations=[PydanticAIIntegration()], |
| 2858 | + traces_sample_rate=1.0, |
| 2859 | + send_default_pii=True, |
| 2860 | + ) |
| 2861 | + |
| 2862 | + events = capture_events() |
| 2863 | + |
| 2864 | + result = await agent.run("What is 5 + 3?") |
| 2865 | + assert result is not None |
| 2866 | + |
| 2867 | + (transaction,) = events |
| 2868 | + spans = transaction["spans"] |
| 2869 | + |
| 2870 | + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] |
| 2871 | + assert len(tool_spans) >= 1 |
| 2872 | + |
| 2873 | + tool_span = tool_spans[0] |
| 2874 | + assert tool_span["data"]["gen_ai.tool.name"] == "no_docs_tool" |
| 2875 | + assert SPANDATA.GEN_AI_TOOL_DESCRIPTION not in tool_span["data"] |
| 2876 | + |
| 2877 | + |
| 2878 | +@pytest.mark.asyncio |
| 2879 | +async def test_execute_tool_span_with_tool_definition(sentry_init, capture_events): |
| 2880 | + """ |
| 2881 | + Test execute_tool_span helper correctly sets tool description from a ToolDefinition. |
| 2882 | + """ |
| 2883 | + |
| 2884 | + sentry_init( |
| 2885 | + integrations=[PydanticAIIntegration()], |
| 2886 | + traces_sample_rate=1.0, |
| 2887 | + send_default_pii=True, |
| 2888 | + ) |
| 2889 | + |
| 2890 | + events = capture_events() |
| 2891 | + |
| 2892 | + tool_def = ToolDefinition( |
| 2893 | + name="my_tool", |
| 2894 | + description="A tool that does something useful.", |
| 2895 | + parameters_json_schema={"type": "object", "properties": {}}, |
| 2896 | + ) |
| 2897 | + |
| 2898 | + with sentry_sdk.start_transaction(op="test", name="test"): |
| 2899 | + with execute_tool_span( |
| 2900 | + "my_tool", {"arg": "value"}, None, "function", tool_definition=tool_def |
| 2901 | + ) as span: |
| 2902 | + assert span is not None |
| 2903 | + |
| 2904 | + (event,) = events |
| 2905 | + tool_spans = [s for s in event["spans"] if s["op"] == "gen_ai.execute_tool"] |
| 2906 | + assert len(tool_spans) == 1 |
| 2907 | + assert ( |
| 2908 | + tool_spans[0]["data"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] |
| 2909 | + == "A tool that does something useful." |
| 2910 | + ) |
| 2911 | + |
| 2912 | + |
| 2913 | +@pytest.mark.asyncio |
| 2914 | +async def test_execute_tool_span_without_tool_definition(sentry_init, capture_events): |
| 2915 | + """ |
| 2916 | + Test execute_tool_span helper omits tool description when tool_definition is None. |
| 2917 | + """ |
| 2918 | + |
| 2919 | + sentry_init( |
| 2920 | + integrations=[PydanticAIIntegration()], |
| 2921 | + traces_sample_rate=1.0, |
| 2922 | + send_default_pii=True, |
| 2923 | + ) |
| 2924 | + |
| 2925 | + events = capture_events() |
| 2926 | + |
| 2927 | + with sentry_sdk.start_transaction(op="test", name="test"): |
| 2928 | + with execute_tool_span( |
| 2929 | + "my_tool", {"arg": "value"}, None, "function", tool_definition=None |
| 2930 | + ) as span: |
| 2931 | + assert span is not None |
| 2932 | + |
| 2933 | + (event,) = events |
| 2934 | + tool_spans = [s for s in event["spans"] if s["op"] == "gen_ai.execute_tool"] |
| 2935 | + assert len(tool_spans) == 1 |
| 2936 | + assert SPANDATA.GEN_AI_TOOL_DESCRIPTION not in tool_spans[0]["data"] |
0 commit comments