pydantic-ai 0.2.19__tar.gz → 0.2.20__tar.gz

{pydantic_ai-0.2.19 → pydantic_ai-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.2.19
+Version: 0.2.20
 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,9 +28,9 @@ 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[a2a,anthropic,bedrock,cli,cohere,evals,google,groq,mcp,mistral,openai,vertexai]==0.2.19
+Requires-Dist: pydantic-ai-slim[a2a,anthropic,bedrock,cli,cohere,evals,google,groq,mcp,mistral,openai,vertexai]==0.2.20
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.2.19; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.2.20; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.2.19 → pydantic_ai-0.2.20}/tests/mcp_server.py

@@ -3,6 +3,8 @@ from pathlib import Path
 from typing import Any
 
 from mcp.server.fastmcp import Context, FastMCP, Image
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT, RequestT
 from mcp.types import BlobResourceContents, EmbeddedResource, TextResourceContents
 from pydantic import AnyUrl
 
@@ -118,6 +120,22 @@ async def get_log_level(ctx: Context) -> str:  # type: ignore
     return log_level
 
 
+@mcp.tool()
+async def echo_deps(ctx: Context[ServerSessionT, LifespanContextT, RequestT]) -> dict[str, Any]:
+    """Echo the run context.
+
+    Args:
+        ctx: Context object containing request and session information.
+
+    Returns:
+        Dictionary with an echo message and the deps.
+    """
+    await ctx.info('This is an info message')
+
+    deps: Any = getattr(ctx.request_context.meta, 'deps')
+    return {'echo': 'This is an echo message', 'deps': deps}
+
+
 @mcp._mcp_server.set_logging_level()  # pyright: ignore[reportPrivateUsage]
 async def set_logging_level(level: str) -> None:
     global log_level

pydantic_ai-0.2.20/tests/models/cassettes/test_google/test_google_timeout.yaml

@@ -0,0 +1,64 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '87'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: Hello!
+        role: user
+      generationConfig: {}
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '687'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=304
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - avgLogprobs: -0.0009780019860376012
+        content:
+          parts:
+          - text: |
+              Hello there! How can I help you today?
+          role: model
+        finishReason: STOP
+      modelVersion: gemini-1.5-flash
+      responseId: KKRRaJyZNJiFm9IPzMfNiAk
+      usageMetadata:
+        candidatesTokenCount: 11
+        candidatesTokensDetails:
+        - modality: TEXT
+          tokenCount: 11
+        promptTokenCount: 2
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 2
+        totalTokenCount: 13
+    status:
+      code: 200
+      message: OK
+version: 1

{pydantic_ai-0.2.19 → pydantic_ai-0.2.20}/tests/models/test_google.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from typing import Any, Union
 
 import pytest
-from httpx import Request
+from httpx import Request, Timeout
 from inline_snapshot import Is, snapshot
 from pytest_mock import MockerFixture
 from typing_extensions import TypedDict
@@ -806,3 +806,14 @@ async def test_google_tool_config_any_with_tool_without_args(
             ),
         ]
     )
+
+
+async def test_google_timeout(allow_model_requests: None, google_provider: GoogleProvider):
+    model = GoogleModel('gemini-1.5-flash', provider=google_provider)
+    agent = Agent(model=model)
+
+    result = await agent.run('Hello!', model_settings={'timeout': 10})
+    assert result.output == snapshot('Hello there! How can I help you today?\n')
+
+    with pytest.raises(UserError, match='Google does not support setting ModelSettings.timeout to a httpx.Timeout'):
+        await agent.run('Hello!', model_settings={'timeout': Timeout(10)})

{pydantic_ai-0.2.19 → pydantic_ai-0.2.20}/tests/test_history_processor.py

@@ -1,4 +1,5 @@
 from collections.abc import AsyncIterator
+from typing import Any, cast
 
 import pytest
 from inline_snapshot import snapshot
@@ -6,6 +7,7 @@ from inline_snapshot import snapshot
 from pydantic_ai import Agent
 from pydantic_ai.messages import ModelMessage, ModelRequest, ModelRequestPart, ModelResponse, TextPart, UserPromptPart
 from pydantic_ai.models.function import AgentInfo, FunctionModel
+from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import Usage
 
 from .conftest import IsDatetime
@@ -201,3 +203,101 @@ async def test_history_processor_on_streamed_run(function_model: FunctionModel,
             ModelRequest(parts=[UserPromptPart(content='Question 2', timestamp=IsDatetime())]),
         ]
     )
+
+
+async def test_history_processor_with_context(function_model: FunctionModel, received_messages: list[ModelMessage]):
+    """Test history processor that takes RunContext."""
+
+    def context_processor(ctx: RunContext[str], messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Access deps from context
+        prefix = ctx.deps
+        processed: list[ModelMessage] = []
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                new_parts: list[ModelRequestPart] = []
+                for part in msg.parts:
+                    if isinstance(part, UserPromptPart):
+                        new_parts.append(UserPromptPart(content=f'{prefix}: {part.content}'))
+                    else:
+                        new_parts.append(part)  # pragma: no cover
+                processed.append(ModelRequest(parts=new_parts))
+            else:
+                processed.append(msg)  # pragma: no cover
+        return processed
+
+    agent = Agent(function_model, history_processors=[context_processor], deps_type=str)
+    await agent.run('test', deps='PREFIX')
+
+    # Verify the prefix was added
+    assert len(received_messages) == 1
+    assert isinstance(received_messages[0], ModelRequest)
+    user_part = received_messages[0].parts[0]
+    assert isinstance(user_part, UserPromptPart)
+    assert user_part.content == 'PREFIX: test'
+
+
+async def test_history_processor_with_context_async(
+    function_model: FunctionModel, received_messages: list[ModelMessage]
+):
+    """Test async history processor that takes RunContext."""
+
+    async def async_context_processor(ctx: RunContext[Any], messages: list[ModelMessage]) -> list[ModelMessage]:
+        return messages[-1:]  # Keep only the last message
+
+    message_history = [
+        ModelRequest(parts=[UserPromptPart(content='Question 1')]),
+        ModelResponse(parts=[TextPart(content='Answer 1')]),
+        ModelRequest(parts=[UserPromptPart(content='Question 2')]),
+        ModelResponse(parts=[TextPart(content='Answer 2')]),
+    ]
+
+    agent = Agent(function_model, history_processors=[async_context_processor])
+    await agent.run('Question 3', message_history=message_history)
+
+    # Should have filtered to only recent messages
+    assert len(received_messages) <= 2  # Last message from history + new message
+
+
+async def test_history_processor_mixed_signatures(function_model: FunctionModel, received_messages: list[ModelMessage]):
+    """Test mixing processors with and without context."""
+
+    def simple_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Filter out responses
+        return [msg for msg in messages if isinstance(msg, ModelRequest)]
+
+    def context_processor(ctx: RunContext[Any], messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Add prefix based on deps
+        prefix = getattr(ctx.deps, 'prefix', 'DEFAULT')
+        processed: list[ModelMessage] = []
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                new_parts: list[ModelRequestPart] = []
+                for part in msg.parts:
+                    if isinstance(part, UserPromptPart):
+                        new_parts.append(UserPromptPart(content=f'{prefix}: {part.content}'))
+                    else:
+                        new_parts.append(part)  # pragma: no cover
+                processed.append(ModelRequest(parts=new_parts))
+            else:
+                processed.append(msg)  # pragma: no cover
+        return processed
+
+    message_history = [
+        ModelRequest(parts=[UserPromptPart(content='Question 1')]),
+        ModelResponse(parts=[TextPart(content='Answer 1')]),
+    ]
+
+    # Create deps with prefix attribute
+    class Deps:
+        prefix = 'TEST'
+
+    agent = Agent(function_model, history_processors=[simple_processor, context_processor], deps_type=Deps)
+    await agent.run('Question 2', message_history=message_history, deps=Deps())
+
+    # Should have filtered responses and added prefix
+    assert len(received_messages) == 2
+    for msg in received_messages:
+        assert isinstance(msg, ModelRequest)
+        user_part = msg.parts[0]
+        assert isinstance(user_part, UserPromptPart)
+        assert cast(str, user_part.content).startswith('TEST: ')

{pydantic_ai-0.2.19 → pydantic_ai-0.2.20}/tests/test_mcp.py

@@ -2,12 +2,14 @@
 
 import re
 from pathlib import Path
+from typing import Any, Final
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from inline_snapshot import snapshot
 
 from pydantic_ai.agent import Agent
-from pydantic_ai.exceptions import UserError
+from pydantic_ai.exceptions import ModelRetry, UserError
 from pydantic_ai.messages import (
     BinaryContent,
     ModelRequest,
@@ -18,17 +20,22 @@ from pydantic_ai.messages import (
     ToolReturnPart,
     UserPromptPart,
 )
+from pydantic_ai.models.test import TestModel
+from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import Usage
 
 from .conftest import IsDatetime, IsStr, try_import
 
 with try_import() as imports_successful:
-    from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio
+    from mcp import ErrorData, McpError
+
+    from pydantic_ai.mcp import CallToolFunc, MCPServerSSE, MCPServerStdio, ToolResult
     from pydantic_ai.models.google import GoogleModel
     from pydantic_ai.models.openai import OpenAIModel
     from pydantic_ai.providers.google import GoogleProvider
     from pydantic_ai.providers.openai import OpenAIProvider
 
+TOOL_COUNT: Final[int] = 12
 
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='mcp and openai not installed'),
@@ -48,7 +55,7 @@ async def test_stdio_server():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == 11
+        assert len(tools) == TOOL_COUNT
         assert tools[0].name == 'celsius_to_fahrenheit'
         assert tools[0].description.startswith('Convert Celsius to Fahrenheit.')
 
@@ -69,7 +76,29 @@ async def test_stdio_server_with_cwd():
     server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir)
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == snapshot(11)
+        assert len(tools) == TOOL_COUNT
+
+
+async def test_process_tool_call() -> None:
+    called: bool = False
+
+    async def process_tool_call(
+        ctx: RunContext[int],
+        call_tool: CallToolFunc,
+        tool_name: str,
+        args: dict[str, Any],
+    ) -> ToolResult:
+        """A process_tool_call that sets a flag and sends deps as metadata."""
+        nonlocal called
+        called = True
+        return await call_tool(tool_name, args, {'deps': ctx.deps})
+
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], process_tool_call=process_tool_call)
+    async with server:
+        agent = Agent(deps_type=int, model=TestModel(call_tools=['echo_deps']), mcp_servers=[server])
+        result = await agent.run('Echo with deps set to 42', deps=42)
+        assert result.output == snapshot('{"echo_deps":{"echo":"This is an echo message","deps":42}}')
+        assert called, 'process_tool_call should have been called'
 
 
 def test_sse_server():
@@ -214,7 +243,7 @@ async def test_log_level_unset():
     assert server._get_log_level() is None  # pyright: ignore[reportPrivateUsage]
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == snapshot(11)
+        assert len(tools) == TOOL_COUNT
         assert tools[10].name == 'get_log_level'
 
         result = await server.call_tool('get_log_level', {})
@@ -932,3 +961,18 @@ async def test_tool_returning_multiple_items(allow_model_requests: None, agent:
                 ),
             ]
         )
+
+
+async def test_mcp_server_raises_mcp_error(allow_model_requests: None, agent: Agent) -> None:
+    server = agent._mcp_servers[0]  # pyright: ignore[reportPrivateUsage]
+
+    mcp_error = McpError(error=ErrorData(code=400, message='Test MCP error conversion'))
+
+    async with agent.run_mcp_servers():
+        with patch.object(
+            server._client,  # pyright: ignore[reportPrivateUsage]
+            'send_request',
+            new=AsyncMock(side_effect=mcp_error),
+        ):
+            with pytest.raises(ModelRetry, match='Test MCP error conversion'):
+                await server.call_tool('test_tool', {})