pydantic-ai-slim 0.1.7__tar.gz → 0.1.9__tar.gz

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.1.7
+Version: 0.1.9
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,7 +29,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.1.7
+Requires-Dist: pydantic-graph==0.1.9
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
@@ -45,7 +45,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.1.7; extra == 'evals'
+Requires-Dist: pydantic-evals==0.1.9; extra == 'evals'
 Provides-Extra: groq
 Requires-Dist: groq>=0.15.0; extra == 'groq'
 Provides-Extra: logfire

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/_agent_graph.py

@@ -2,6 +2,7 @@ from __future__ import annotations as _annotations
 
 import asyncio
 import dataclasses
+import hashlib
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
@@ -92,6 +93,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
     mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
+    default_retries: int
 
     tracer: Tracer
 
@@ -546,7 +548,14 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
     )
 
 
-async def process_function_tools(
+def multi_modal_content_identifier(identifier: str | bytes) -> str:
+    """Generate stable identifier for multi-modal content to help LLM in finding a specific file in tool call responses."""
+    if isinstance(identifier, str):
+        identifier = identifier.encode('utf-8')
+    return hashlib.sha1(identifier).hexdigest()[:6]
+
+
+async def process_function_tools(  # noqa C901
     tool_calls: list[_messages.ToolCallPart],
     output_tool_name: str | None,
     output_tool_call_id: str | None,
@@ -632,6 +641,8 @@ async def process_function_tools(
     if not calls_to_run:
         return
 
+    user_parts: list[_messages.UserPromptPart] = []
+
     # Run all tool tasks in parallel
     results_by_index: dict[int, _messages.ModelRequestPart] = {}
     with ctx.deps.tracer.start_as_current_span(
@@ -645,6 +656,7 @@ async def process_function_tools(
             asyncio.create_task(tool.run(call, run_context, ctx.deps.tracer), name=call.tool_name)
             for tool, call in calls_to_run
         ]
+
         pending = tasks
         while pending:
             done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
@@ -652,7 +664,43 @@ async def process_function_tools(
                 index = tasks.index(task)
                 result = task.result()
                 yield _messages.FunctionToolResultEvent(result, tool_call_id=call_index_to_event_id[index])
-                if isinstance(result, (_messages.ToolReturnPart, _messages.RetryPromptPart)):
+
+                if isinstance(result, _messages.RetryPromptPart):
+                    results_by_index[index] = result
+                elif isinstance(result, _messages.ToolReturnPart):
+                    contents: list[Any]
+                    single_content: bool
+                    if isinstance(result.content, list):
+                        contents = result.content  # type: ignore
+                        single_content = False
+                    else:
+                        contents = [result.content]
+                        single_content = True
+
+                    processed_contents: list[Any] = []
+                    for content in contents:
+                        if isinstance(content, _messages.MultiModalContentTypes):
+                            if isinstance(content, _messages.BinaryContent):
+                                identifier = multi_modal_content_identifier(content.data)
+                            else:
+                                identifier = multi_modal_content_identifier(content.url)
+
+                            user_parts.append(
+                                _messages.UserPromptPart(
+                                    content=[f'This is file {identifier}:', content],
+                                    timestamp=result.timestamp,
+                                    part_kind='user-prompt',
+                                )
+                            )
+                            processed_contents.append(f'See file {identifier}')
+                        else:
+                            processed_contents.append(content)
+
+                    if single_content:
+                        result.content = processed_contents[0]
+                    else:
+                        result.content = processed_contents
+
                     results_by_index[index] = result
                 else:
                     assert_never(result)
@@ -662,6 +710,8 @@ async def process_function_tools(
     for k in sorted(results_by_index):
         output_parts.append(results_by_index[k])
 
+    output_parts.extend(user_parts)
+
 
 async def _tool_from_mcp_server(
     tool_name: str,
@@ -688,7 +738,7 @@ async def _tool_from_mcp_server(
     for server in ctx.deps.mcp_servers:
         tools = await server.list_tools()
         if tool_name in {tool.name for tool in tools}:
-            return Tool(name=tool_name, function=run_tool, takes_ctx=True)
+            return Tool(name=tool_name, function=run_tool, takes_ctx=True, max_retries=ctx.deps.default_retries)
     return None
 
 

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/agent.py

@@ -658,6 +658,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             output_validators=output_validators,
             function_tools=self._function_tools,
             mcp_servers=self._mcp_servers,
+            default_retries=self._default_retries,
             tracer=tracer,
             get_instructions=get_instructions,
         )

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/mcp.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import base64
+import json
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Sequence
 from contextlib import AsyncExitStack, asynccontextmanager
@@ -9,16 +11,25 @@ from types import TracebackType
 from typing import Any
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from mcp.types import JSONRPCMessage, LoggingLevel
-from typing_extensions import Self
-
+from mcp.types import (
+    BlobResourceContents,
+    EmbeddedResource,
+    ImageContent,
+    JSONRPCMessage,
+    LoggingLevel,
+    TextContent,
+    TextResourceContents,
+)
+from typing_extensions import Self, assert_never
+
+from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.messages import BinaryContent
 from pydantic_ai.tools import ToolDefinition
 
 try:
     from mcp.client.session import ClientSession
     from mcp.client.sse import sse_client
     from mcp.client.stdio import StdioServerParameters, stdio_client
-    from mcp.types import CallToolResult
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `mcp` package to use the MCP server, '
@@ -74,7 +85,9 @@ class MCPServer(ABC):
             for tool in tools.tools
         ]
 
-    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> CallToolResult:
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any]
+    ) -> str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]:
         """Call a tool on the server.
 
         Args:
@@ -83,8 +96,21 @@ class MCPServer(ABC):
 
         Returns:
             The result of the tool call.
+
+        Raises:
+            ModelRetry: If the tool call fails.
         """
-        return await self._client.call_tool(tool_name, arguments)
+        result = await self._client.call_tool(tool_name, arguments)
+
+        content = [self._map_tool_result_part(part) for part in result.content]
+
+        if result.isError:
+            text = '\n'.join(str(part) for part in content)
+            raise ModelRetry(text)
+
+        if len(content) == 1:
+            return content[0]
+        return content
 
     async def __aenter__(self) -> Self:
         self._exit_stack = AsyncExitStack()
@@ -105,6 +131,35 @@ class MCPServer(ABC):
         await self._exit_stack.aclose()
         self.is_running = False
 
+    def _map_tool_result_part(
+        self, part: TextContent | ImageContent | EmbeddedResource
+    ) -> str | BinaryContent | dict[str, Any] | list[Any]:
+        # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values
+
+        if isinstance(part, TextContent):
+            text = part.text
+            if text.startswith(('[', '{')):
+                try:
+                    return json.loads(text)
+                except ValueError:
+                    pass
+            return text
+        elif isinstance(part, ImageContent):
+            return BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
+        elif isinstance(part, EmbeddedResource):
+            resource = part.resource
+            if isinstance(resource, TextResourceContents):
+                return resource.text
+            elif isinstance(resource, BlobResourceContents):
+                return BinaryContent(
+                    data=base64.b64decode(resource.blob),
+                    media_type=resource.mimeType or 'application/octet-stream',
+                )
+            else:
+                assert_never(resource)
+        else:
+            assert_never(part)
+
 
 @dataclass
 class MCPServerStdio(MCPServer):

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/messages.py

@@ -253,6 +253,9 @@ class BinaryContent:
 
 UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | DocumentUrl | VideoUrl | BinaryContent'
 
+# Ideally this would be a Union of types, but Python 3.9 requires it to be a string, and strings don't work with `isinstance``.
+MultiModalContentTypes = (ImageUrl, AudioUrl, DocumentUrl, VideoUrl, BinaryContent)
+
 
 def _document_format(media_type: str) -> DocumentFormat:
     if media_type == 'application/pdf':
@@ -829,4 +832,6 @@ class FunctionToolResultEvent:
     """Event type identifier, used as a discriminator."""
 
 
-HandleResponseEvent = Annotated[Union[FunctionToolCallEvent, FunctionToolResultEvent], pydantic.Discriminator('kind')]
+HandleResponseEvent = Annotated[
+    Union[FunctionToolCallEvent, FunctionToolResultEvent], pydantic.Discriminator('event_kind')
+]

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/models/anthropic.py

@@ -109,10 +109,6 @@ class AnthropicModel(Model):
     Internally, this uses the [Anthropic Python client](https://github.com/anthropics/anthropic-sdk-python) to interact with the API.
 
     Apart from `__init__`, all methods are private or match those of the base class.
-
-    !!! note
-        The `AnthropicModel` class does not yet support streaming responses.
-        We anticipate adding support for streaming responses in a near-term future release.
     """
 
     client: AsyncAnthropic = field(repr=False)
@@ -409,13 +405,27 @@ def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage
     if response_usage is None:
         return usage.Usage()
 
-    request_tokens = getattr(response_usage, 'input_tokens', None)
+    # Store all integer-typed usage values in the details dict
+    response_usage_dict = response_usage.model_dump()
+    details: dict[str, int] = {}
+    for key, value in response_usage_dict.items():
+        if isinstance(value, int):
+            details[key] = value
+
+    # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence the getattr call
+    # Tokens are only counted once between input_tokens, cache_creation_input_tokens, and cache_read_input_tokens
+    # This approach maintains request_tokens as the count of all input tokens, with cached counts as details
+    request_tokens = (
+        getattr(response_usage, 'input_tokens', 0)
+        + (getattr(response_usage, 'cache_creation_input_tokens', 0) or 0)  # These can be missing, None, or int
+        + (getattr(response_usage, 'cache_read_input_tokens', 0) or 0)
+    )
 
     return usage.Usage(
-        # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence this getattr
-        request_tokens=request_tokens,
+        request_tokens=request_tokens or None,
         response_tokens=response_usage.output_tokens,
-        total_tokens=(request_tokens or 0) + response_usage.output_tokens,
+        total_tokens=request_tokens + response_usage.output_tokens,
+        details=details or None,
     )
 
 

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/models/mistral.py

@@ -483,7 +483,20 @@ class MistralModel(Model):
                 assert_never(message)
         if instructions := self._get_instructions(messages):
             mistral_messages.insert(0, MistralSystemMessage(content=instructions))
-        return mistral_messages
+
+        # Post-process messages to insert fake assistant message after tool message if followed by user message
+        # to work around `Unexpected role 'user' after role 'tool'` error.
+        processed_messages: list[MistralMessages] = []
+        for i, current_message in enumerate(mistral_messages):
+            processed_messages.append(current_message)
+
+            if isinstance(current_message, MistralToolMessage) and i + 1 < len(mistral_messages):
+                next_message = mistral_messages[i + 1]
+                if isinstance(next_message, MistralUserMessage):
+                    # Insert a dummy assistant message
+                    processed_messages.append(MistralAssistantMessage(content=[MistralTextChunk(text='OK')]))
+
+        return processed_messages
 
     def _map_user_prompt(self, part: UserPromptPart) -> MistralUserMessage:
         content: str | list[MistralContentChunk]

{pydantic_ai_slim-0.1.7 → pydantic_ai_slim-0.1.9}/pydantic_ai/providers/mistral.py

@@ -44,6 +44,7 @@ class MistralProvider(Provider[Mistral]):
         *,
         api_key: str | None = None,
         mistral_client: Mistral | None = None,
+        base_url: str | None = None,
         http_client: AsyncHTTPClient | None = None,
     ) -> None:
         """Create a new Mistral provider.
@@ -52,11 +53,13 @@ class MistralProvider(Provider[Mistral]):
             api_key: The API key to use for authentication, if not provided, the `MISTRAL_API_KEY` environment variable
                 will be used if available.
             mistral_client: An existing `Mistral` client to use, if provided, `api_key` and `http_client` must be `None`.
+            base_url: The base url for the Mistral requests.
             http_client: An existing async client to use for making HTTP requests.
         """
         if mistral_client is not None:
             assert http_client is None, 'Cannot provide both `mistral_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `mistral_client` and `api_key`'
+            assert base_url is None, 'Cannot provide both `mistral_client` and `base_url`'
             self._client = mistral_client
         else:
             api_key = api_key or os.environ.get('MISTRAL_API_KEY')
@@ -67,7 +70,7 @@ class MistralProvider(Provider[Mistral]):
                     'to use the Mistral provider.'
                 )
             elif http_client is not None:
-                self._client = Mistral(api_key=api_key, async_client=http_client)
+                self._client = Mistral(api_key=api_key, async_client=http_client, server_url=base_url)
             else:
                 http_client = cached_async_http_client(provider='mistral')
-                self._client = Mistral(api_key=api_key, async_client=http_client)
+                self._client = Mistral(api_key=api_key, async_client=http_client, server_url=base_url)