pydantic-ai-slim 0.2.5__tar.gz → 0.2.6__tar.gz

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.2.5
+Version: 0.2.6
 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.2.5
+Requires-Dist: pydantic-graph==0.2.6
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.2.5; extra == 'a2a'
+Requires-Dist: fasta2a==0.2.6; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.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.2.5; extra == 'evals'
+Requires-Dist: pydantic-evals==0.2.6; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/_agent_graph.py

@@ -26,7 +26,7 @@ from . import (
 )
 from .result import OutputDataT, ToolOutput
 from .settings import ModelSettings, merge_model_settings
-from .tools import RunContext, Tool, ToolDefinition
+from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
 if TYPE_CHECKING:
     from .mcp import MCPServer
@@ -97,6 +97,8 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     tracer: Tracer
 
+    prepare_tools: ToolsPrepareFunc[DepsT] | None = None
+
 
 class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
     """The base class for all agent nodes.
@@ -241,6 +243,11 @@ async def _prepare_request_parameters(
         *map(add_mcp_server_tools, ctx.deps.mcp_servers),
     )
 
+    if ctx.deps.prepare_tools:
+        # Prepare the tools using the provided function
+        # This also acts over tool definitions pulled from MCP servers
+        function_tool_defs = await ctx.deps.prepare_tools(run_context, function_tool_defs) or []
+
     output_schema = ctx.deps.output_schema
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/_cli.py

@@ -53,7 +53,7 @@ PYDANTIC_AI_HOME = Path.home() / '.pydantic-ai'
 This folder is used to store the prompt history and configuration.
 """
 
-PROMPT_HISTORY_PATH = PYDANTIC_AI_HOME / 'prompt-history.txt'
+PROMPT_HISTORY_FILENAME = 'prompt-history.txt'
 
 
 class SimpleCodeBlock(CodeBlock):
@@ -211,27 +211,26 @@ Special prompts:
             pass
         return 0
 
-    # Ensure the history directory and file exist
-    PROMPT_HISTORY_PATH.parent.mkdir(parents=True, exist_ok=True)
-    PROMPT_HISTORY_PATH.touch(exist_ok=True)
-
-    # doing this instead of `PromptSession[Any](history=` allows mocking of PromptSession in tests
-    session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
     try:
-        return asyncio.run(run_chat(session, stream, agent, console, code_theme, prog_name))
+        return asyncio.run(run_chat(stream, agent, console, code_theme, prog_name))
     except KeyboardInterrupt:  # pragma: no cover
         return 0
 
 
 async def run_chat(
-    session: PromptSession[Any],
     stream: bool,
     agent: Agent[AgentDepsT, OutputDataT],
     console: Console,
     code_theme: str,
     prog_name: str,
+    config_dir: Path | None = None,
     deps: AgentDepsT = None,
 ) -> int:
+    prompt_history_path = (config_dir or PYDANTIC_AI_HOME) / PROMPT_HISTORY_FILENAME
+    prompt_history_path.parent.mkdir(parents=True, exist_ok=True)
+    prompt_history_path.touch(exist_ok=True)
+    session: PromptSession[Any] = PromptSession(history=FileHistory(str(prompt_history_path)))
+
     multiline = False
     messages: list[ModelMessage] = []
 

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/agent.py

@@ -42,6 +42,7 @@ from .tools import (
     ToolFuncPlain,
     ToolParams,
     ToolPrepareFunc,
+    ToolsPrepareFunc,
 )
 
 # Re-exporting like this improves auto-import behavior in PyCharm
@@ -148,6 +149,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(
         repr=False
     )
+    _prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
     _mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
@@ -172,6 +174,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -200,6 +203,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -223,6 +227,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -251,6 +256,9 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             output_retries: The maximum number of retries to allow for result validation, defaults to `retries`.
             tools: Tools to register with the agent, you can also register tools via the decorators
                 [`@agent.tool`][pydantic_ai.Agent.tool] and [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain].
+            prepare_tools: custom method to prepare the tool definition of all tools for each step.
+                This is useful if you want to customize the definition of multiple tools or you want to register
+                a subset of tools for a given step. See [`ToolsPrepareFunc`][pydantic_ai.tools.ToolsPrepareFunc]
             mcp_servers: MCP servers to register with the agent. You should register a [`MCPServer`][pydantic_ai.mcp.MCPServer]
                 for each server you want the agent to connect to.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
@@ -334,6 +342,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self._default_retries = retries
         self._max_result_retries = output_retries if output_retries is not None else retries
         self._mcp_servers = mcp_servers
+        self._prepare_tools = prepare_tools
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
@@ -694,6 +703,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             mcp_servers=self._mcp_servers,
             default_retries=self._default_retries,
             tracer=tracer,
+            prepare_tools=self._prepare_tools,
             get_instructions=get_instructions,
         )
         start_node = _agent_graph.UserPromptNode[AgentDepsT](
@@ -1763,18 +1773,14 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             await agent.to_cli()
         ```
         """
-        from prompt_toolkit import PromptSession
-        from prompt_toolkit.history import FileHistory
         from rich.console import Console
 
-        from pydantic_ai._cli import PROMPT_HISTORY_PATH, run_chat
+        from pydantic_ai._cli import run_chat
 
         # TODO(Marcelo): We need to refactor the CLI code to be able to be able to just pass `agent`, `deps` and
         # `prog_name` from here.
 
-        session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
         await run_chat(
-            session=session,
             stream=True,
             agent=self,
             deps=deps,

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/models/__init__.py

@@ -491,7 +491,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)
-    elif provider in ('deepseek', 'openai', 'azure'):
+    elif provider in ('deepseek', 'openai', 'azure', 'openrouter'):
         from .openai import OpenAIModel
 
         return OpenAIModel(model_name, provider=provider)

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/models/google.py

@@ -279,7 +279,6 @@ class GoogleModel(Model):
     def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         if not response.candidates or len(response.candidates) != 1:
             raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')  # pragma: no cover
-        print(response.candidates[0].safety_ratings)
         if response.candidates[0].content is None or response.candidates[0].content.parts is None:
             if response.candidates[0].finish_reason == 'SAFETY':
                 raise UnexpectedModelBehavior('Safety settings triggered', str(response))

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/models/openai.py

@@ -170,7 +170,7 @@ class OpenAIModel(Model):
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure'] | Provider[AsyncOpenAI] = 'openai',
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter'] | Provider[AsyncOpenAI] = 'openai',
         system_prompt_role: OpenAISystemPromptRole | None = None,
     ):
         """Initialize an OpenAI model.

{pydantic_ai_slim-0.2.5 → pydantic_ai_slim-0.2.6}/pydantic_ai/tools.py

@@ -29,6 +29,7 @@ __all__ = (
     'ToolFuncEither',
     'ToolParams',
     'ToolPrepareFunc',
+    'ToolsPrepareFunc',
     'Tool',
     'ObjectJsonSchema',
     'ToolDefinition',
@@ -133,6 +134,37 @@ hitchhiker = Tool(hitchhiker, prepare=only_if_42)
 Usage `ToolPrepareFunc[AgentDepsT]`.
 """
 
+ToolsPrepareFunc: TypeAlias = (
+    'Callable[[RunContext[AgentDepsT], list[ToolDefinition]], Awaitable[list[ToolDefinition] | None]]'
+)
+"""Definition of a function that can prepare the tool definition of all tools for each step.
+This is useful if you want to customize the definition of multiple tools or you want to register
+a subset of tools for a given step.
+
+Example — here `turn_on_strict_if_openai` is valid as a `ToolsPrepareFunc`:
+
+```python {noqa="I001"}
+from dataclasses import replace
+from typing import Union
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.tools import ToolDefinition
+
+
+async def turn_on_strict_if_openai(
+    ctx: RunContext[None], tool_defs: list[ToolDefinition]
+) -> Union[list[ToolDefinition], None]:
+    if ctx.model.system == 'openai':
+        return [replace(tool_def, strict=True) for tool_def in tool_defs]
+    return tool_defs
+
+agent = Agent('openai:gpt-4o', prepare_tools=turn_on_strict_if_openai)
+```
+
+Usage `ToolsPrepareFunc[AgentDepsT]`.
+"""
+
+
 DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
 """Supported docstring formats.