pydantic-ai-slim 0.3.1__tar.gz → 0.3.3__tar.gz

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/.gitignore

@@ -19,3 +19,4 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 node_modules/
 **.idea/
 .coverage*
+/test_tmp/

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.3.1
+Version: 0.3.3
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
 License-Expression: MIT
@@ -30,11 +30,11 @@ 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.3.1
+Requires-Dist: pydantic-graph==0.3.3
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.3.1; extra == 'a2a'
+Requires-Dist: fasta2a==0.3.3; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.52.0; extra == 'anthropic'
 Provides-Extra: bedrock
@@ -48,7 +48,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.3.1; extra == 'evals'
+Requires-Dist: pydantic-evals==0.3.3; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/pydantic_ai/__init__.py

@@ -12,7 +12,7 @@ from .exceptions import (
 )
 from .format_prompt import format_as_xml
 from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
-from .result import ToolOutput
+from .output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
 from .tools import RunContext, Tool
 
 __all__ = (
@@ -41,8 +41,11 @@ __all__ = (
     # tools
     'Tool',
     'RunContext',
-    # result
+    # output
     'ToolOutput',
+    'NativeOutput',
+    'PromptedOutput',
+    'TextOutput',
     # format_prompt
     'format_as_xml',
 )

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/pydantic_ai/_agent_graph.py

@@ -18,7 +18,7 @@ from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
 
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
-from .result import OutputDataT
+from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
@@ -102,7 +102,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     end_strategy: EndStrategy
     get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]]
 
-    output_schema: _output.OutputSchema[OutputDataT] | None
+    output_schema: _output.OutputSchema[OutputDataT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]
@@ -141,6 +141,8 @@ def is_agent_node(
 
 @dataclasses.dataclass
 class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
+    """The node that handles the user prompt and instructions."""
+
     user_prompt: str | Sequence[_messages.UserContent] | None
 
     instructions: str | None
@@ -284,16 +286,29 @@ async def _prepare_request_parameters(
         function_tool_defs = await ctx.deps.prepare_tools(run_context, function_tool_defs) or []
 
     output_schema = ctx.deps.output_schema
+
+    output_tools = []
+    output_object = None
+    if isinstance(output_schema, _output.ToolOutputSchema):
+        output_tools = output_schema.tool_defs()
+    elif isinstance(output_schema, _output.NativeOutputSchema):
+        output_object = output_schema.object_def
+
+    # ToolOrTextOutputSchema, NativeOutputSchema, and PromptedOutputSchema all inherit from TextOutputSchema
+    allow_text_output = isinstance(output_schema, _output.TextOutputSchema)
+
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
-        allow_text_output=_output.allow_text_output(output_schema),
-        output_tools=output_schema.tool_defs() if output_schema is not None else [],
+        output_mode=output_schema.mode,
+        output_tools=output_tools,
+        output_object=output_object,
+        allow_text_output=allow_text_output,
     )
 
 
 @dataclasses.dataclass
 class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
-    """Make a request to the model using the last message in state.message_history."""
+    """The node that makes a request to the model using the last message in state.message_history."""
 
     request: _messages.ModelRequest
 
@@ -412,7 +427,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
 @dataclasses.dataclass
 class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
-    """Process a model response, and decide whether to end the run or make a new request."""
+    """The node that processes a model response, and decides whether to end the run or make a new request."""
 
     model_response: _messages.ModelResponse
 
@@ -482,7 +497,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     # when the model has already returned text along side tool calls
                     # in this scenario, if text responses are allowed, we return text from the most recent model
                     # response, if any
-                    if _output.allow_text_output(ctx.deps.output_schema):
+                    if isinstance(ctx.deps.output_schema, _output.TextOutputSchema):
                         for message in reversed(ctx.state.message_history):
                             if isinstance(message, _messages.ModelResponse):
                                 last_texts = [p.content for p in message.parts if isinstance(p, _messages.TextPart)]
@@ -505,10 +520,11 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         output_schema = ctx.deps.output_schema
         run_context = build_run_context(ctx)
 
-        # first, look for the output tool call
         final_result: result.FinalResult[NodeRunEndT] | None = None
         parts: list[_messages.ModelRequestPart] = []
-        if output_schema is not None:
+
+        # first, look for the output tool call
+        if isinstance(output_schema, _output.ToolOutputSchema):
             for call, output_tool in output_schema.find_tool(tool_calls):
                 try:
                     result_data = await output_tool.process(call, run_context)
@@ -566,9 +582,9 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
 
         text = '\n\n'.join(texts)
         try:
-            if _output.allow_text_output(output_schema):
-                # The following cast is safe because we know `str` is an allowed result type
-                result_data = cast(NodeRunEndT, text)
+            if isinstance(output_schema, _output.TextOutputSchema):
+                run_context = build_run_context(ctx)
+                result_data = await output_schema.process(text, run_context)
             else:
                 m = _messages.RetryPromptPart(
                     content='Plain text responses are not permitted, please include your response in a tool call',
@@ -667,7 +683,7 @@ async def process_function_tools(  # noqa C901
                 yield event
                 call_index_to_event_id[len(calls_to_run)] = event.call_id
                 calls_to_run.append((mcp_tool, call))
-        elif output_schema is not None and call.tool_name in output_schema.tools:
+        elif call.tool_name in output_schema.tools:
             # if tool_name is in output_schema, it means we found a output tool but an error occurred in
             # validation, we don't add another part here
             if output_tool_name is not None:
@@ -807,7 +823,9 @@ def _unknown_tool(
 ) -> _messages.RetryPromptPart:
     ctx.state.increment_retries(ctx.deps.max_result_retries)
     tool_names = list(ctx.deps.function_tools.keys())
-    if output_schema := ctx.deps.output_schema:
+
+    output_schema = ctx.deps.output_schema
+    if isinstance(output_schema, _output.ToolOutputSchema):
         tool_names.extend(output_schema.tool_names())
 
     if tool_names:
@@ -884,7 +902,7 @@ def get_captured_run_messages() -> _RunMessages:
 def build_agent_graph(
     name: str | None,
     deps_type: type[DepsT],
-    output_type: _output.OutputType[OutputT],
+    output_type: OutputSpec[OutputT],
 ) -> Graph[GraphAgentState, GraphAgentDeps[DepsT, result.FinalResult[OutputT]], result.FinalResult[OutputT]]:
     """Build the execution [Graph][pydantic_graph.Graph] for a given agent."""
     nodes = (

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/pydantic_ai/_cli.py

@@ -14,14 +14,13 @@ from typing import Any, cast
 
 from typing_inspection.introspection import get_literal_values
 
-from pydantic_ai.result import OutputDataT
-from pydantic_ai.tools import AgentDepsT
-
 from . import __version__
+from ._run_context import AgentDepsT
 from .agent import Agent
 from .exceptions import UserError
 from .messages import ModelMessage
 from .models import KnownModelName, infer_model
+from .output import OutputDataT
 
 try:
     import argcomplete
@@ -254,6 +253,11 @@ async def run_chat(
                 messages = await ask_agent(agent, text, stream, console, code_theme, deps, messages)
             except CancelledError:  # pragma: no cover
                 console.print('[dim]Interrupted[/dim]')
+            except Exception as e:  # pragma: no cover
+                cause = getattr(e, '__cause__', None)
+                console.print(f'\n[red]{type(e).__name__}:[/red] {e}')
+                if cause:
+                    console.print(f'[dim]Caused by: {cause}[/dim]')
 
 
 async def ask_agent(

{pydantic_ai_slim-0.3.1 → pydantic_ai_slim-0.3.3}/pydantic_ai/_function_schema.py

@@ -19,9 +19,8 @@ from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import Concatenate, ParamSpec, TypeIs, TypeVar, get_origin
 
-from pydantic_ai.tools import RunContext
-
 from ._griffe import doc_descriptions
+from ._run_context import RunContext
 from ._utils import check_object_json_schema, is_async_callable, is_model_like, run_in_executor
 
 if TYPE_CHECKING:
@@ -281,6 +280,4 @@ def _build_schema(
 
 def _is_call_ctx(annotation: Any) -> bool:
     """Return whether the annotation is the `RunContext` class, parameterized or not."""
-    from .tools import RunContext
-
     return annotation is RunContext or get_origin(annotation) is RunContext

pydantic_ai_slim-0.3.3/pydantic_ai/_mcp.py

@@ -0,0 +1,123 @@
+import base64
+from collections.abc import Sequence
+from typing import Literal
+
+from . import exceptions, messages
+
+try:
+    from mcp import types as mcp_types
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install the `mcp` package to use the MCP server, '
+        'you can use the `mcp` optional group — `pip install "pydantic-ai-slim[mcp]"`'
+    ) from _import_error
+
+
+def map_from_mcp_params(params: mcp_types.CreateMessageRequestParams) -> list[messages.ModelMessage]:
+    """Convert from MCP create message request parameters to pydantic-ai messages."""
+    pai_messages: list[messages.ModelMessage] = []
+    request_parts: list[messages.ModelRequestPart] = []
+    if params.systemPrompt:
+        request_parts.append(messages.SystemPromptPart(content=params.systemPrompt))
+    response_parts: list[messages.ModelResponsePart] = []
+    for msg in params.messages:
+        content = msg.content
+        if msg.role == 'user':
+            # if there are any response parts, add a response message wrapping them
+            if response_parts:
+                pai_messages.append(messages.ModelResponse(parts=response_parts))
+                response_parts = []
+
+            # TODO(Marcelo): We can reuse the `_map_tool_result_part` from the mcp module here.
+            if isinstance(content, mcp_types.TextContent):
+                user_part_content: str | Sequence[messages.UserContent] = content.text
+            else:
+                # image content
+                user_part_content = [
+                    messages.BinaryContent(data=base64.b64decode(content.data), media_type=content.mimeType)
+                ]
+
+            request_parts.append(messages.UserPromptPart(content=user_part_content))
+        else:
+            # role is assistant
+            # if there are any request parts, add a request message wrapping them
+            if request_parts:
+                pai_messages.append(messages.ModelRequest(parts=request_parts))
+                request_parts = []
+
+            response_parts.append(map_from_sampling_content(content))
+
+    if response_parts:
+        pai_messages.append(messages.ModelResponse(parts=response_parts))
+    if request_parts:
+        pai_messages.append(messages.ModelRequest(parts=request_parts))
+    return pai_messages
+
+
+def map_from_pai_messages(pai_messages: list[messages.ModelMessage]) -> tuple[str, list[mcp_types.SamplingMessage]]:
+    """Convert from pydantic-ai messages to MCP sampling messages.
+
+    Returns:
+        A tuple containing the system prompt and a list of sampling messages.
+    """
+    sampling_msgs: list[mcp_types.SamplingMessage] = []
+
+    def add_msg(
+        role: Literal['user', 'assistant'],
+        content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent,
+    ):
+        sampling_msgs.append(mcp_types.SamplingMessage(role=role, content=content))
+
+    system_prompt: list[str] = []
+    for pai_message in pai_messages:
+        if isinstance(pai_message, messages.ModelRequest):
+            if pai_message.instructions is not None:
+                system_prompt.append(pai_message.instructions)
+
+            for part in pai_message.parts:
+                if isinstance(part, messages.SystemPromptPart):
+                    system_prompt.append(part.content)
+                if isinstance(part, messages.UserPromptPart):
+                    if isinstance(part.content, str):
+                        add_msg('user', mcp_types.TextContent(type='text', text=part.content))
+                    else:
+                        for chunk in part.content:
+                            if isinstance(chunk, str):
+                                add_msg('user', mcp_types.TextContent(type='text', text=chunk))
+                            elif isinstance(chunk, messages.BinaryContent) and chunk.is_image:
+                                add_msg(
+                                    'user',
+                                    mcp_types.ImageContent(
+                                        type='image',
+                                        data=base64.b64decode(chunk.data).decode(),
+                                        mimeType=chunk.media_type,
+                                    ),
+                                )
+                            # TODO(Marcelo): Add support for audio content.
+                            else:
+                                raise NotImplementedError(f'Unsupported content type: {type(chunk)}')
+        else:
+            add_msg('assistant', map_from_model_response(pai_message))
+    return ''.join(system_prompt), sampling_msgs
+
+
+def map_from_model_response(model_response: messages.ModelResponse) -> mcp_types.TextContent:
+    """Convert from a model response to MCP text content."""
+    text_parts: list[str] = []
+    for part in model_response.parts:
+        if isinstance(part, messages.TextPart):
+            text_parts.append(part.content)
+        # TODO(Marcelo): We should ignore ThinkingPart here.
+        else:
+            raise exceptions.UnexpectedModelBehavior(f'Unexpected part type: {type(part).__name__}, expected TextPart')
+    return mcp_types.TextContent(type='text', text=''.join(text_parts))
+
+
+def map_from_sampling_content(
+    content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent,
+) -> messages.TextPart:
+    """Convert from sampling content to a pydantic-ai text part."""
+    if isinstance(content, mcp_types.TextContent):  # pragma: no branch
+        return messages.TextPart(content=content.text)
+    else:
+        raise NotImplementedError('Image and Audio responses in sampling are not yet supported')