pydantic-ai 0.4.7__tar.gz → 0.4.9__tar.gz

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.4.7
+Version: 0.4.9
 Summary: Agent Framework / shim to use Pydantic with LLMs
 Project-URL: Homepage, https://ai.pydantic.dev
 Project-URL: Source, https://github.com/pydantic/pydantic-ai
@@ -28,11 +28,11 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,mcp,mistral,openai,vertexai]==0.4.7
+Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,mcp,mistral,openai,retries,vertexai]==0.4.9
 Provides-Extra: a2a
 Requires-Dist: fasta2a>=0.4.1; extra == 'a2a'
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.4.7; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.4.9; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/pyproject.toml

@@ -47,7 +47,7 @@ requires-python = ">=3.9"
 
 [tool.hatch.metadata.hooks.uv-dynamic-versioning]
 dependencies = [
-    "pydantic-ai-slim[openai,vertexai,google,groq,anthropic,mistral,cohere,bedrock,huggingface,cli,mcp,evals,ag-ui]=={{ version }}",
+    "pydantic-ai-slim[openai,vertexai,google,groq,anthropic,mistral,cohere,bedrock,huggingface,cli,mcp,evals,ag-ui,retries]=={{ version }}",
 ]
 
 [tool.hatch.metadata.hooks.uv-dynamic-versioning.optional-dependencies]

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/tests/models/test_gemini.py

@@ -194,7 +194,7 @@ async def test_require_response_tool(allow_model_requests: None):
 
 async def test_json_def_replaced(allow_model_requests: None):
     class Axis(BaseModel):
-        label: str
+        label: str = Field(default='<unlabeled axis>', description='The label of the axis')
 
     class Chart(BaseModel):
         x_axis: Axis
@@ -213,8 +213,14 @@ async def test_json_def_replaced(allow_model_requests: None):
         {
             '$defs': {
                 'Axis': {
-                    'properties': {'label': {'title': 'Label', 'type': 'string'}},
-                    'required': ['label'],
+                    'properties': {
+                        'label': {
+                            'default': '<unlabeled axis>',
+                            'description': 'The label of the axis',
+                            'title': 'Label',
+                            'type': 'string',
+                        }
+                    },
                     'title': 'Axis',
                     'type': 'object',
                 },
@@ -268,17 +274,27 @@ async def test_json_def_replaced(allow_model_requests: None):
                                 'items': {
                                     'properties': {
                                         'lat': {'type': 'number'},
-                                        'lng': {'type': 'number'},
+                                        'lng': {'default': 1.1, 'type': 'number'},
                                         'chart': {
                                             'properties': {
                                                 'x_axis': {
-                                                    'properties': {'label': {'type': 'string'}},
-                                                    'required': ['label'],
+                                                    'properties': {
+                                                        'label': {
+                                                            'default': '<unlabeled axis>',
+                                                            'description': 'The label of the axis',
+                                                            'type': 'string',
+                                                        }
+                                                    },
                                                     'type': 'object',
                                                 },
                                                 'y_axis': {
-                                                    'properties': {'label': {'type': 'string'}},
-                                                    'required': ['label'],
+                                                    'properties': {
+                                                        'label': {
+                                                            'default': '<unlabeled axis>',
+                                                            'description': 'The label of the axis',
+                                                            'type': 'string',
+                                                        }
+                                                    },
                                                     'type': 'object',
                                                 },
                                             },

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/tests/models/test_instrumented.py

@@ -26,6 +26,7 @@ from pydantic_ai.messages import (
     SystemPromptPart,
     TextPart,
     TextPartDelta,
+    ThinkingPart,
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
@@ -150,7 +151,7 @@ async def test_instrumented_model(capfire: CaptureLogfire):
                 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                 'parent': None,
                 'start_time': 1000000000,
-                'end_time': 18000000000,
+                'end_time': 16000000000,
                 'attributes': {
                     'gen_ai.operation.name': 'chat',
                     'gen_ai.system': 'my_system',
@@ -284,7 +285,7 @@ Fix the errors and try again.\
                     'index': 0,
                     'message': {
                         'role': 'assistant',
-                        'content': 'text1',
+                        'content': [{'kind': 'text', 'text': 'text1'}, {'kind': 'text', 'text': 'text2'}],
                         'tool_calls': [
                             {
                                 'id': 'tool_call_1',
@@ -308,17 +309,6 @@ Fix the errors and try again.\
                 'span_id': 1,
                 'trace_flags': 1,
             },
-            {
-                'body': {'index': 0, 'message': {'role': 'assistant', 'content': 'text2'}},
-                'severity_number': 9,
-                'severity_text': None,
-                'attributes': {'gen_ai.system': 'my_system', 'event.name': 'gen_ai.choice'},
-                'timestamp': 16000000000,
-                'observed_timestamp': 17000000000,
-                'trace_id': 1,
-                'span_id': 1,
-                'trace_flags': 1,
-            },
         ]
     )
 
@@ -641,11 +631,13 @@ Fix the errors and try again.\
                                     'gen_ai.system': 'my_system',
                                 },
                                 {
-                                    'event.name': 'gen_ai.choice',
                                     'index': 0,
                                     'message': {
                                         'role': 'assistant',
-                                        'content': 'text1',
+                                        'content': [
+                                            {'kind': 'text', 'text': 'text1'},
+                                            {'kind': 'text', 'text': 'text2'},
+                                        ],
                                         'tool_calls': [
                                             {
                                                 'id': 'tool_call_1',
@@ -660,12 +652,7 @@ Fix the errors and try again.\
                                         ],
                                     },
                                     'gen_ai.system': 'my_system',
-                                },
-                                {
                                     'event.name': 'gen_ai.choice',
-                                    'index': 0,
-                                    'message': {'role': 'assistant', 'content': 'text2'},
-                                    'gen_ai.system': 'my_system',
                                 },
                             ]
                         )
@@ -879,6 +866,7 @@ def test_messages_without_content(document_content: BinaryContent):
             },
             {
                 'role': 'assistant',
+                'content': [{'kind': 'text'}],
                 'gen_ai.message.index': 1,
                 'event.name': 'gen_ai.assistant.message',
             },
@@ -897,6 +885,7 @@ def test_messages_without_content(document_content: BinaryContent):
             },
             {
                 'role': 'assistant',
+                'content': [{'kind': 'text'}],
                 'tool_calls': [
                     {
                         'id': IsStr(),
@@ -935,3 +924,38 @@ def test_messages_without_content(document_content: BinaryContent):
             },
         ]
     )
+
+
+def test_message_with_thinking_parts():
+    messages: list[ModelMessage] = [
+        ModelResponse(parts=[TextPart('text1'), ThinkingPart('thinking1'), TextPart('text2')]),
+        ModelResponse(parts=[ThinkingPart('thinking2')]),
+        ModelResponse(parts=[ThinkingPart('thinking3'), TextPart('text3')]),
+    ]
+    settings = InstrumentationSettings()
+    assert [InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(messages)] == snapshot(
+        [
+            {
+                'role': 'assistant',
+                'content': [
+                    {'kind': 'text', 'text': 'text1'},
+                    {'kind': 'thinking', 'text': 'thinking1'},
+                    {'kind': 'text', 'text': 'text2'},
+                ],
+                'gen_ai.message.index': 0,
+                'event.name': 'gen_ai.assistant.message',
+            },
+            {
+                'role': 'assistant',
+                'content': [{'kind': 'thinking', 'text': 'thinking2'}],
+                'gen_ai.message.index': 1,
+                'event.name': 'gen_ai.assistant.message',
+            },
+            {
+                'role': 'assistant',
+                'content': [{'kind': 'thinking', 'text': 'thinking3'}, {'kind': 'text', 'text': 'text3'}],
+                'gen_ai.message.index': 2,
+                'event.name': 'gen_ai.assistant.message',
+            },
+        ]
+    )

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/tests/test_ag_ui.py

@@ -7,6 +7,7 @@ import contextlib
 import json
 import uuid
 from collections.abc import AsyncIterator
+from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any
 
@@ -17,7 +18,9 @@ from dirty_equals import IsStr
 from inline_snapshot import snapshot
 from pydantic import BaseModel
 
+from pydantic_ai._run_context import RunContext
 from pydantic_ai.agent import Agent
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai.models.function import (
     AgentInfo,
@@ -27,8 +30,9 @@ from pydantic_ai.models.function import (
     DeltaToolCalls,
     FunctionModel,
 )
+from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import OutputDataT
-from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from .conftest import IsSameStr
 
@@ -180,7 +184,7 @@ def create_input(
         thread_id=thread_id,
         run_id=uuid_str(),
         messages=list(messages),
-        state=state,
+        state=dict(state) if state else {},
         context=[],
         tools=tools or [],
         forwarded_props=None,
@@ -1050,9 +1054,19 @@ async def test_tool_local_then_ag_ui() -> None:
 async def test_request_with_state() -> None:
     """Test request with state modification."""
 
+    seen_states: list[int] = []
+
+    async def store_state(
+        ctx: RunContext[StateDeps[StateInt]], tool_defs: list[ToolDefinition]
+    ) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state.value)
+        ctx.deps.state.value += 1
+        return tool_defs
+
     agent: Agent[StateDeps[StateInt], str] = Agent(
         model=FunctionModel(stream_function=simple_stream),
         deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
+        prepare_tools=store_state,
     )
     adapter = _Adapter(agent=agent)
     run_inputs = [
@@ -1074,32 +1088,101 @@ async def test_request_with_state() -> None:
                 id='msg_3',
                 content='Hello, how are you?',
             ),
+        ),
+        create_input(
+            UserMessage(
+                id='msg_4',
+                content='Hello, how are you?',
+            ),
             state=StateInt(value=42),
         ),
     ]
 
-    deps = StateDeps(StateInt())
+    deps = StateDeps(StateInt(value=0))
 
-    last_value = deps.state.value
     for run_input in run_inputs:
         events = list[dict[str, Any]]()
         async for event in adapter.run(run_input, deps=deps):
             events.append(json.loads(event.removeprefix('data: ')))
 
         assert events == simple_result()
-        assert deps.state.value == run_input.state.value if run_input.state is not None else last_value
-        last_value = deps.state.value
+    assert seen_states == snapshot(
+        [
+            41,  # run msg_1, prepare_tools call 1
+            42,  # run msg_1, prepare_tools call 2
+            0,  # run msg_2, prepare_tools call 1
+            1,  # run msg_2, prepare_tools call 2
+            0,  # run msg_3, prepare_tools call 1
+            1,  # run msg_3, prepare_tools call 2
+            42,  # run msg_4, prepare_tools call 1
+            43,  # run msg_4, prepare_tools call 2
+        ]
+    )
+
+
+async def test_request_with_state_without_handler() -> None:
+    agent = Agent(model=FunctionModel(stream_function=simple_stream))
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state=StateInt(value=41),
+    )
+
+    with pytest.raises(
+        UserError,
+        match='AG-UI state is provided but `deps` of type `NoneType` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.',
+    ):
+        async for _ in adapter.run(run_input):
+            pass
+
+
+async def test_request_with_state_with_custom_handler() -> None:
+    @dataclass
+    class CustomStateDeps:
+        state: dict[str, Any]
+
+    seen_states: list[dict[str, Any]] = []
+
+    async def store_state(ctx: RunContext[CustomStateDeps], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state)
+        return tool_defs
+
+    agent: Agent[CustomStateDeps, str] = Agent(
+        model=FunctionModel(stream_function=simple_stream),
+        deps_type=CustomStateDeps,
+        prepare_tools=store_state,
+    )
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state={'value': 42},
+    )
+
+    async for _ in adapter.run(run_input, deps=CustomStateDeps(state={'value': 0})):
+        pass
 
-    assert deps.state.value == 42
+    assert seen_states[-1] == {'value': 42}
 
 
 async def test_concurrent_runs() -> None:
     """Test concurrent execution of multiple runs."""
     import asyncio
 
-    agent = Agent(
-        model=FunctionModel(stream_function=simple_stream),
+    agent: Agent[StateDeps[StateInt], str] = Agent(
+        model=TestModel(),
+        deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
     )
+
+    @agent.tool
+    async def get_state(ctx: RunContext[StateDeps[StateInt]]) -> int:
+        return ctx.deps.state.value
+
     adapter = _Adapter(agent=agent)
     concurrent_tasks: list[asyncio.Task[list[dict[str, Any]]]] = []
 
@@ -1109,10 +1192,11 @@ async def test_concurrent_runs() -> None:
                 id=f'msg_{i}',
                 content=f'Message {i}',
             ),
+            state=StateInt(value=i),
             thread_id=f'test_thread_{i}',
         )
 
-        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input))
+        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input, deps=StateDeps(StateInt())))
         concurrent_tasks.append(task)
 
     results = await asyncio.gather(*concurrent_tasks)
@@ -1121,9 +1205,23 @@ async def test_concurrent_runs() -> None:
     for i, events in enumerate(results):
         assert events == [
             {'type': 'RUN_STARTED', 'threadId': f'test_thread_{i}', 'runId': (run_id := IsSameStr())},
+            {
+                'type': 'TOOL_CALL_START',
+                'toolCallId': (tool_call_id := IsSameStr()),
+                'toolCallName': 'get_state',
+                'parentMessageId': IsStr(),
+            },
+            {'type': 'TOOL_CALL_END', 'toolCallId': tool_call_id},
+            {
+                'type': 'TOOL_CALL_RESULT',
+                'messageId': IsStr(),
+                'toolCallId': tool_call_id,
+                'content': str(i),
+                'role': 'tool',
+            },
             {'type': 'TEXT_MESSAGE_START', 'messageId': (message_id := IsSameStr()), 'role': 'assistant'},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'success '},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '(no tool calls)'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '{"get_s'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'tate":' + str(i) + '}'},
             {'type': 'TEXT_MESSAGE_END', 'messageId': message_id},
             {'type': 'RUN_FINISHED', 'threadId': f'test_thread_{i}', 'runId': run_id},
         ]

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/tests/test_mcp.py

@@ -91,6 +91,20 @@ async def test_reentrant_context_manager():
             pass
 
 
+async def test_context_manager_initialization_error() -> None:
+    """Test if streams are closed if client fails to initialize."""
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    from mcp.client.session import ClientSession
+
+    with patch.object(ClientSession, 'initialize', side_effect=Exception):
+        with pytest.raises(Exception):
+            async with server:
+                pass
+
+    assert server._read_stream._closed  # pyright: ignore[reportPrivateUsage]
+    assert server._write_stream._closed  # pyright: ignore[reportPrivateUsage]
+
+
 async def test_stdio_server_with_tool_prefix(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], tool_prefix='foo')
     async with server:

{pydantic_ai-0.4.7 → pydantic_ai-0.4.9}/tests/test_streaming.py

@@ -1108,6 +1108,42 @@ async def test_iter_stream_structured_output():
                     )
 
 
+async def test_iter_stream_output_tool_dont_hit_retry_limit():
+    class CityLocation(BaseModel):
+        city: str
+        country: str | None = None
+
+    async def text_stream(_messages: list[ModelMessage], agent_info: AgentInfo) -> AsyncIterator[DeltaToolCalls]:
+        """Stream partial JSON data that will initially fail validation."""
+        assert agent_info.output_tools is not None
+        assert len(agent_info.output_tools) == 1
+        name = agent_info.output_tools[0].name
+
+        yield {0: DeltaToolCall(name=name)}
+        yield {0: DeltaToolCall(json_args='{"c')}
+        yield {0: DeltaToolCall(json_args='ity":')}
+        yield {0: DeltaToolCall(json_args=' "Mex')}
+        yield {0: DeltaToolCall(json_args='ico City",')}
+        yield {0: DeltaToolCall(json_args=' "cou')}
+        yield {0: DeltaToolCall(json_args='ntry": "Mexico"}')}
+
+    agent = Agent(FunctionModel(stream_function=text_stream), output_type=CityLocation)
+
+    async with agent.iter('Generate city info') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    assert [c async for c in stream.stream_output(debounce_by=None)] == snapshot(
+                        [
+                            CityLocation(city='Mex'),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                        ]
+                    )
+
+
 def test_function_tool_event_tool_call_id_properties():
     """Ensure that the `tool_call_id` property on function tool events mirrors the underlying part's ID."""
     # Prepare a ToolCallPart with a fixed ID