pydantic-ai-slim 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

pydantic_ai/__init__.py

@@ -1,8 +1,8 @@
 from importlib.metadata import version
 
 from .agent import Agent
-from .dependencies import RunContext
 from .exceptions import ModelRetry, UnexpectedModelBehavior, UserError
+from .tools import RunContext, Tool
 
-__all__ = 'Agent', 'RunContext', 'ModelRetry', 'UnexpectedModelBehavior', 'UserError', '__version__'
+__all__ = 'Agent', 'Tool', 'RunContext', 'ModelRetry', 'UnexpectedModelBehavior', 'UserError', '__version__'
 __version__ = version('pydantic_ai_slim')

pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 from __future__ import annotations as _annotations
 
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, TypedDict, cast, get_origin
+from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast, get_origin
 
 from pydantic import ConfigDict, TypeAdapter
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -20,8 +20,7 @@ from ._griffe import doc_descriptions
 from ._utils import ObjectJsonSchema, check_object_json_schema, is_model_like
 
 if TYPE_CHECKING:
-    from . import _tool
-    from .dependencies import AgentDeps, ToolParams
+    pass
 
 
 __all__ = 'function_schema', 'LazyTypeAdapter'
@@ -39,17 +38,16 @@ class FunctionSchema(TypedDict):
     var_positional_field: str | None
 
 
-def function_schema(either_function: _tool.ToolEitherFunc[AgentDeps, ToolParams]) -> FunctionSchema:  # noqa: C901
+def function_schema(function: Callable[..., Any], takes_ctx: bool) -> FunctionSchema:  # noqa: C901
     """Build a Pydantic validator and JSON schema from a tool function.
 
     Args:
-        either_function: The function to build a validator and JSON schema for.
+        function: The function to build a validator and JSON schema for.
+        takes_ctx: Whether the function takes a `RunContext` first argument.
 
     Returns:
         A `FunctionSchema` instance.
     """
-    function = either_function.whichever()
-    takes_ctx = either_function.is_left()
     config = ConfigDict(title=function.__name__)
     config_wrapper = ConfigWrapper(config)
     gen_schema = _generate_schema.GenerateSchema(config_wrapper)
@@ -78,13 +76,13 @@ def function_schema(either_function: _tool.ToolEitherFunc[AgentDeps, ToolParams]
 
             if index == 0 and takes_ctx:
                 if not _is_call_ctx(annotation):
-                    errors.append('First argument must be a RunContext instance when using `.tool`')
+                    errors.append('First parameter of tools that take context must be annotated with RunContext[...]')
                 continue
             elif not takes_ctx and _is_call_ctx(annotation):
-                errors.append('RunContext instance can only be used with `.tool`')
+                errors.append('RunContext annotations can only be used with tools that take context')
                 continue
             elif index != 0 and _is_call_ctx(annotation):
-                errors.append('RunContext instance can only be used as the first argument')
+                errors.append('RunContext annotations can only be used as the first argument')
                 continue
 
         field_name = p.name
@@ -159,6 +157,24 @@ def function_schema(either_function: _tool.ToolEitherFunc[AgentDeps, ToolParams]
     )
 
 
+def takes_ctx(function: Callable[..., Any]) -> bool:
+    """Check if a function takes a `RunContext` first argument.
+
+    Args:
+        function: The function to check.
+
+    Returns:
+        `True` if the function takes a `RunContext` as first argument, `False` otherwise.
+    """
+    sig = signature(function)
+    try:
+        _, first_param = next(iter(sig.parameters.items()))
+    except StopIteration:
+        return False
+    else:
+        return first_param.annotation is not sig.empty and _is_call_ctx(first_param.annotation)
+
+
 def _build_schema(
     fields: dict[str, core_schema.TypedDictField],
     var_kwargs_schema: core_schema.CoreSchema | None,
@@ -191,7 +207,7 @@ def _build_schema(
 
 
 def _is_call_ctx(annotation: Any) -> bool:
-    from .dependencies import RunContext
+    from .tools import RunContext
 
     return annotation is RunContext or (
         _typing_extra.is_generic_alias(annotation) and get_origin(annotation) is RunContext

pydantic_ai/_result.py

@@ -11,10 +11,10 @@ from pydantic import TypeAdapter, ValidationError
 from typing_extensions import Self, TypeAliasType, TypedDict
 
 from . import _utils, messages
-from .dependencies import AgentDeps, ResultValidatorFunc, RunContext
 from .exceptions import ModelRetry
 from .messages import ModelStructuredResponse, ToolCall
 from .result import ResultData
+from .tools import AgentDeps, ResultValidatorFunc, RunContext
 
 
 @dataclass

pydantic_ai/_system_prompt.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, cast
 
 from . import _utils
-from .dependencies import AgentDeps, RunContext, SystemPromptFunc
+from .tools import AgentDeps, RunContext, SystemPromptFunc
 
 
 @dataclass

pydantic_ai/agent.py

@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import asyncio
+import dataclasses
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
@@ -12,15 +13,14 @@ from typing_extensions import assert_never
 from . import (
     _result,
     _system_prompt,
-    _tool as _r,
     _utils,
     exceptions,
     messages as _messages,
     models,
     result,
 )
-from .dependencies import AgentDeps, RunContext, ToolContextFunc, ToolParams, ToolPlainFunc
 from .result import ResultData
+from .tools import AgentDeps, RunContext, Tool, ToolFuncContext, ToolFuncEither, ToolFuncPlain, ToolParams
 
 __all__ = ('Agent',)
 
@@ -34,7 +34,7 @@ NoneType = type(None)
 class Agent(Generic[AgentDeps, ResultData]):
     """Class for defining "agents" - a way to have a specific type of "conversation" with an LLM.
 
-    Agents are generic in the dependency type they take [`AgentDeps`][pydantic_ai.dependencies.AgentDeps]
+    Agents are generic in the dependency type they take [`AgentDeps`][pydantic_ai.tools.AgentDeps]
     and the result data type they return, [`ResultData`][pydantic_ai.result.ResultData].
 
     By default, if neither generic parameter is customised, agents have type `Agent[None, str]`.
@@ -58,7 +58,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = field(repr=False)
     _allow_text_result: bool = field(repr=False)
     _system_prompts: tuple[str, ...] = field(repr=False)
-    _function_tools: dict[str, _r.Tool[AgentDeps, Any]] = field(repr=False)
+    _function_tools: dict[str, Tool[AgentDeps]] = field(repr=False)
     _default_retries: int = field(repr=False)
     _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDeps]] = field(repr=False)
     _deps_type: type[AgentDeps] = field(repr=False)
@@ -83,6 +83,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_tool_name: str = 'final_result',
         result_tool_description: str | None = None,
         result_retries: int | None = None,
+        tools: Sequence[Tool[AgentDeps] | ToolFuncEither[AgentDeps, ...]] = (),
         defer_model_check: bool = False,
     ):
         """Create an agent.
@@ -101,6 +102,8 @@ class Agent(Generic[AgentDeps, ResultData]):
             result_tool_name: The name of the tool to use for the final result.
             result_tool_description: The description of the final result tool.
             result_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].
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
                 it's evaluated to create a [`Model`][pydantic_ai.models.Model] instance immediately,
                 which checks for the necessary environment variables. Set this to `false`
@@ -119,9 +122,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         self._allow_text_result = self._result_schema is None or self._result_schema.allow_text_result
 
         self._system_prompts = (system_prompt,) if isinstance(system_prompt, str) else tuple(system_prompt)
-        self._function_tools: dict[str, _r.Tool[AgentDeps, Any]] = {}
-        self._deps_type = deps_type
+        self._function_tools = {}
         self._default_retries = retries
+        for tool in tools:
+            self._register_tool(Tool.infer(tool))
+        self._deps_type = deps_type
         self._system_prompt_functions = []
         self._max_result_retries = result_retries if result_retries is not None else retries
         self._current_result_retry = 0
@@ -206,7 +211,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt synchronously.
 
-        This is a convenience method that wraps `self.run` with `asyncio.run()`.
+        This is a convenience method that wraps `self.run` with `loop.run_until_complete()`.
 
         Args:
             user_prompt: User input to start/continue the conversation.
@@ -217,7 +222,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         Returns:
             The result of the run.
         """
-        return asyncio.run(self.run(user_prompt, message_history=message_history, model=model, deps=deps))
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self.run(user_prompt, message_history=message_history, model=model, deps=deps))
 
     @asynccontextmanager
     async def run_stream(
@@ -284,6 +290,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                                     self._result_schema,
                                     deps,
                                     self._result_validators,
+                                    lambda m: run_span.set_attribute('all_messages', messages),
                                 )
                                 return
                             else:
@@ -353,7 +360,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
         """Decorator to register a system prompt function.
 
-        Optionally takes [`RunContext`][pydantic_ai.dependencies.RunContext] as it's only argument.
+        Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as it's only argument.
         Can decorate a sync or async functions.
 
         Overloads for every possible signature of `system_prompt` are included so the decorator doesn't obscure
@@ -404,7 +411,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     ) -> _result.ResultValidatorFunc[AgentDeps, ResultData]:
         """Decorator to register a result validator function.
 
-        Optionally takes [`RunContext`][pydantic_ai.dependencies.RunContext] as it's first argument.
+        Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as it's first argument.
         Can decorate a sync or async functions.
 
         Overloads for every possible signature of `result_validator` are included so the decorator doesn't obscure
@@ -437,22 +444,22 @@ class Agent(Generic[AgentDeps, ResultData]):
         return func
 
     @overload
-    def tool(self, func: ToolContextFunc[AgentDeps, ToolParams], /) -> ToolContextFunc[AgentDeps, ToolParams]: ...
+    def tool(self, func: ToolFuncContext[AgentDeps, ToolParams], /) -> ToolFuncContext[AgentDeps, ToolParams]: ...
 
     @overload
     def tool(
         self, /, *, retries: int | None = None
-    ) -> Callable[[ToolContextFunc[AgentDeps, ToolParams]], ToolContextFunc[AgentDeps, ToolParams]]: ...
+    ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
 
     def tool(
         self,
-        func: ToolContextFunc[AgentDeps, ToolParams] | None = None,
+        func: ToolFuncContext[AgentDeps, ToolParams] | None = None,
         /,
         *,
         retries: int | None = None,
     ) -> Any:
         """Decorator to register a tool function which takes
-        [`RunContext`][pydantic_ai.dependencies.RunContext] as its first argument.
+        [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
         Can decorate a sync or async functions.
 
@@ -489,27 +496,27 @@ class Agent(Generic[AgentDeps, ResultData]):
         if func is None:
 
             def tool_decorator(
-                func_: ToolContextFunc[AgentDeps, ToolParams],
-            ) -> ToolContextFunc[AgentDeps, ToolParams]:
+                func_: ToolFuncContext[AgentDeps, ToolParams],
+            ) -> ToolFuncContext[AgentDeps, ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_tool(_utils.Either(left=func_), retries)
+                self._register_function(func_, True, retries)
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_tool(_utils.Either(left=func), retries)
+            self._register_function(func, True, retries)
             return func
 
     @overload
-    def tool_plain(self, func: ToolPlainFunc[ToolParams], /) -> ToolPlainFunc[ToolParams]: ...
+    def tool_plain(self, func: ToolFuncPlain[ToolParams], /) -> ToolFuncPlain[ToolParams]: ...
 
     @overload
     def tool_plain(
         self, /, *, retries: int | None = None
-    ) -> Callable[[ToolPlainFunc[ToolParams]], ToolPlainFunc[ToolParams]]: ...
+    ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
-    def tool_plain(self, func: ToolPlainFunc[ToolParams] | None = None, /, *, retries: int | None = None) -> Any:
+    def tool_plain(self, func: ToolFuncPlain[ToolParams] | None = None, /, *, retries: int | None = None) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
         Can decorate a sync or async functions.
@@ -546,28 +553,34 @@ class Agent(Generic[AgentDeps, ResultData]):
         """
         if func is None:
 
-            def tool_decorator(
-                func_: ToolPlainFunc[ToolParams],
-            ) -> ToolPlainFunc[ToolParams]:
+            def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_tool(_utils.Either(right=func_), retries)
+                self._register_function(func_, False, retries)
                 return func_
 
             return tool_decorator
         else:
-            self._register_tool(_utils.Either(right=func), retries)
+            self._register_function(func, False, retries)
             return func
 
-    def _register_tool(self, func: _r.ToolEitherFunc[AgentDeps, ToolParams], retries: int | None) -> None:
-        """Private utility to register a tool function."""
+    def _register_function(
+        self, func: ToolFuncEither[AgentDeps, ToolParams], takes_ctx: bool, retries: int | None
+    ) -> None:
+        """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
-        tool = _r.Tool[AgentDeps, ToolParams](func, retries_)
+        tool = Tool(func, takes_ctx, max_retries=retries_)
+        self._register_tool(tool)
 
-        if self._result_schema and tool.name in self._result_schema.tools:
-            raise ValueError(f'Tool name conflicts with result schema name: {tool.name!r}')
+    def _register_tool(self, tool: Tool[AgentDeps]) -> None:
+        """Private utility to register a tool instance."""
+        if tool.max_retries is None:
+            tool = dataclasses.replace(tool, max_retries=self._default_retries)
 
         if tool.name in self._function_tools:
-            raise ValueError(f'Tool name conflicts with existing tool: {tool.name!r}')
+            raise exceptions.UserError(f'Tool name conflicts with existing tool: {tool.name!r}')
+
+        if self._result_schema and tool.name in self._result_schema.tools:
+            raise exceptions.UserError(f'Tool name conflicts with result schema name: {tool.name!r}')
 
         self._function_tools[tool.name] = tool
 

pydantic_ai/messages.py

@@ -1,15 +1,13 @@
 from __future__ import annotations as _annotations
 
 import json
-from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Union
 
 import pydantic
 import pydantic_core
 from pydantic import TypeAdapter
-from typing_extensions import TypeAlias, TypeAliasType
 
 from . import _pydantic
 from ._utils import now_utc as _now_utc
@@ -44,13 +42,7 @@ class UserPrompt:
     """Message type identifier, this type is available on all message as a discriminator."""
 
 
-JsonData: TypeAlias = 'Union[str, int, float, None, Sequence[JsonData], Mapping[str, JsonData]]'
-if not TYPE_CHECKING:
-    # work around for https://github.com/pydantic/pydantic/issues/10873
-    # this is need for pydantic to work both `json_ta` and `MessagesTypeAdapter` at the bottom of this file
-    JsonData = TypeAliasType('JsonData', 'Union[str, int, float, None, Sequence[JsonData], Mapping[str, JsonData]]')
-
-json_ta: TypeAdapter[JsonData] = TypeAdapter(JsonData)
+tool_return_ta: TypeAdapter[Any] = TypeAdapter(Any)
 
 
 @dataclass
@@ -59,7 +51,7 @@ class ToolReturn:
 
     tool_name: str
     """The name of the "tool" was called."""
-    content: JsonData
+    content: Any
     """The return value."""
     tool_id: str | None = None
     """Optional tool identifier, this is used by some models including OpenAI."""
@@ -72,15 +64,14 @@ class ToolReturn:
         if isinstance(self.content, str):
             return self.content
         else:
-            content = json_ta.validate_python(self.content)
-            return json_ta.dump_json(content).decode()
+            return tool_return_ta.dump_json(self.content).decode()
 
-    def model_response_object(self) -> dict[str, JsonData]:
+    def model_response_object(self) -> dict[str, Any]:
         # gemini supports JSON dict return values, but no other JSON types, hence we wrap anything else in a dict
         if isinstance(self.content, dict):
-            return json_ta.validate_python(self.content)  # pyright: ignore[reportReturnType]
+            return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]
         else:
-            return {'return_value': json_ta.validate_python(self.content)}
+            return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
 
 @dataclass
@@ -139,12 +130,16 @@ class ModelTextResponse:
 
 @dataclass
 class ArgsJson:
+    """Tool arguments as a JSON string."""
+
     args_json: str
     """A JSON string of arguments."""
 
 
 @dataclass
 class ArgsDict:
+    """Tool arguments as a Python dictionary."""
+
     args_dict: dict[str, Any]
     """A python dictionary of arguments."""
 

pydantic_ai/models/__init__.py

@@ -259,20 +259,30 @@ def infer_model(model: Model | KnownModelName) -> Model:
 class AbstractToolDefinition(Protocol):
     """Abstract definition of a function/tool.
 
-    This is used for both tools and result tools.
+    This is used for both function tools and result tools.
     """
 
-    name: str
-    """The name of the tool."""
-    description: str
-    """The description of the tool."""
-    json_schema: ObjectJsonSchema
-    """The JSON schema for the tool's arguments."""
-    outer_typed_dict_key: str | None
-    """The key in the outer [TypedDict] that wraps a result tool.
+    @property
+    def name(self) -> str:
+        """The name of the tool."""
+        ...
 
-    This will only be set for result tools which don't have an `object` JSON schema.
-    """
+    @property
+    def description(self) -> str:
+        """The description of the tool."""
+        ...
+
+    @property
+    def json_schema(self) -> ObjectJsonSchema:
+        """The JSON schema for the tool's arguments."""
+        ...
+
+    @property
+    def outer_typed_dict_key(self) -> str | None:
+        """The key in the outer [TypedDict] that wraps a result tool.
+
+        This will only be set for result tools which don't have an `object` JSON schema.
+        """
 
 
 @cache

pydantic_ai/models/gemini.py

@@ -109,11 +109,15 @@ class GeminiModel(Model):
 
 
 class AuthProtocol(Protocol):
+    """Abstract definition for Gemini authentication."""
+
     async def headers(self) -> dict[str, str]: ...
 
 
 @dataclass
 class ApiKeyAuth:
+    """Authentication using an API key for the `X-Goog-Api-Key` header."""
+
     api_key: str
 
     async def headers(self) -> dict[str, str]:

pydantic_ai/models/test.py

@@ -31,14 +31,6 @@ from . import (
 )
 
 
-class UnSetType:
-    def __repr__(self):
-        return 'UnSet'
-
-
-UnSet = UnSetType()
-
-
 @dataclass
 class TestModel(Model):
     """A model specifically for testing purposes.
@@ -186,6 +178,8 @@ class TestAgentModel(AgentModel):
 
 @dataclass
 class TestStreamTextResponse(StreamTextResponse):
+    """A text response that streams test data."""
+
     _text: str
     _cost: Cost
     _iter: Iterator[str] = field(init=False)
@@ -217,6 +211,8 @@ class TestStreamTextResponse(StreamTextResponse):
 
 @dataclass
 class TestStreamStructuredResponse(StreamStructuredResponse):
+    """A structured response that streams test data."""
+
     _structured_response: ModelStructuredResponse
     _cost: Cost
     _iter: Iterator[None] = field(default_factory=lambda: iter([None]))

pydantic_ai/models/vertexai.py

@@ -182,6 +182,8 @@ MAX_TOKEN_AGE = timedelta(seconds=3000)
 
 @dataclass
 class BearerTokenAuth:
+    """Authentication using a bearer token generated by google-auth."""
+
     credentials: BaseCredentials | ServiceAccountCredentials
     token_created: datetime | None = field(default=None, init=False)
 

pydantic_ai/result.py

@@ -1,15 +1,15 @@
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator
-from dataclasses import dataclass
+from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic, TypeVar, cast
 
 import logfire_api
 
 from . import _result, _utils, exceptions, messages, models
-from .dependencies import AgentDeps
+from .tools import AgentDeps
 
 __all__ = (
     'ResultData',
@@ -49,11 +49,11 @@ class Cost:
         This is provided so it's trivial to sum costs from multiple requests and runs.
         """
         counts: dict[str, int] = {}
-        for field in 'request_tokens', 'response_tokens', 'total_tokens':
-            self_value = getattr(self, field)
-            other_value = getattr(other, field)
+        for f in 'request_tokens', 'response_tokens', 'total_tokens':
+            self_value = getattr(self, f)
+            other_value = getattr(other, f)
             if self_value is not None or other_value is not None:
-                counts[field] = (self_value or 0) + (other_value or 0)
+                counts[f] = (self_value or 0) + (other_value or 0)
 
         details = self.details.copy() if self.details is not None else None
         if other.details is not None:
@@ -122,7 +122,8 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
     _result_schema: _result.ResultSchema[ResultData] | None
     _deps: AgentDeps
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]]
-    is_complete: bool = False
+    _on_complete: Callable[[list[messages.Message]], None]
+    is_complete: bool = field(default=False, init=False)
     """Whether the stream has all been received.
 
     This is set to `True` when one of
@@ -312,3 +313,4 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         else:
             assert structured_message is not None, 'Either text or structured_message should provided, not both'
             self._all_messages.append(structured_message)
+        self._on_complete(self._all_messages)

pydantic_ai/tools.py

@@ -0,0 +1,240 @@
+from __future__ import annotations as _annotations
+
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast
+
+from pydantic import ValidationError
+from pydantic_core import SchemaValidator
+from typing_extensions import Concatenate, ParamSpec, final
+
+from . import _pydantic, _utils, messages
+from .exceptions import ModelRetry, UnexpectedModelBehavior
+
+if TYPE_CHECKING:
+    from .result import ResultData
+else:
+    ResultData = Any
+
+
+__all__ = (
+    'AgentDeps',
+    'RunContext',
+    'ResultValidatorFunc',
+    'SystemPromptFunc',
+    'ToolFuncContext',
+    'ToolFuncPlain',
+    'ToolFuncEither',
+    'ToolParams',
+    'Tool',
+)
+
+AgentDeps = TypeVar('AgentDeps')
+"""Type variable for agent dependencies."""
+
+
+@dataclass
+class RunContext(Generic[AgentDeps]):
+    """Information about the current call."""
+
+    deps: AgentDeps
+    """Dependencies for the agent."""
+    retry: int
+    """Number of retries so far."""
+    tool_name: str | None
+    """Name of the tool being called."""
+
+
+ToolParams = ParamSpec('ToolParams')
+"""Retrieval function param spec."""
+
+SystemPromptFunc = Union[
+    Callable[[RunContext[AgentDeps]], str],
+    Callable[[RunContext[AgentDeps]], Awaitable[str]],
+    Callable[[], str],
+    Callable[[], Awaitable[str]],
+]
+"""A function that may or maybe not take `RunContext` as an argument, and may or may not be async.
+
+Usage `SystemPromptFunc[AgentDeps]`.
+"""
+
+ResultValidatorFunc = Union[
+    Callable[[RunContext[AgentDeps], ResultData], ResultData],
+    Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]],
+    Callable[[ResultData], ResultData],
+    Callable[[ResultData], Awaitable[ResultData]],
+]
+"""
+A function that always takes `ResultData` and returns `ResultData`,
+but may or maybe not take `CallInfo` as a first argument, and may or may not be async.
+
+Usage `ResultValidator[AgentDeps, ResultData]`.
+"""
+
+ToolFuncContext = Callable[Concatenate[RunContext[AgentDeps], ToolParams], Any]
+"""A tool function that takes `RunContext` as the first argument.
+
+Usage `ToolContextFunc[AgentDeps, ToolParams]`.
+"""
+ToolFuncPlain = Callable[ToolParams, Any]
+"""A tool function that does not take `RunContext` as the first argument.
+
+Usage `ToolPlainFunc[ToolParams]`.
+"""
+ToolFuncEither = Union[ToolFuncContext[AgentDeps, ToolParams], ToolFuncPlain[ToolParams]]
+"""Either kind of tool function.
+
+This is just a union of [`ToolFuncContext`][pydantic_ai.tools.ToolFuncContext] and
+[`ToolFuncPlain`][pydantic_ai.tools.ToolFuncPlain].
+
+Usage `ToolFuncEither[AgentDeps, ToolParams]`.
+"""
+
+A = TypeVar('A')
+
+
+@final
+@dataclass(init=False)
+class Tool(Generic[AgentDeps]):
+    """A tool function for an agent."""
+
+    function: ToolFuncEither[AgentDeps, ...]
+    takes_ctx: bool
+    max_retries: int | None
+    name: str
+    description: str
+    _is_async: bool = field(init=False)
+    _single_arg_name: str | None = field(init=False)
+    _positional_fields: list[str] = field(init=False)
+    _var_positional_field: str | None = field(init=False)
+    _validator: SchemaValidator = field(init=False, repr=False)
+    _json_schema: _utils.ObjectJsonSchema = field(init=False)
+    _current_retry: int = field(default=0, init=False)
+
+    def __init__(
+        self,
+        function: ToolFuncEither[AgentDeps, ...],
+        takes_ctx: bool,
+        *,
+        max_retries: int | None = None,
+        name: str | None = None,
+        description: str | None = None,
+    ):
+        """Create a new tool instance.
+
+        Example usage:
+
+        ```py
+        from pydantic_ai import Agent, RunContext, Tool
+
+        async def my_tool(ctx: RunContext[int], x: int, y: int) -> str:
+            return f'{ctx.deps} {x} {y}'
+
+        agent = Agent('test', tools=[Tool(my_tool, True)])
+        ```
+
+        Args:
+            function: The Python function to call as the tool.
+            takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] first argument.
+            max_retries: Maximum number of retries allowed for this tool, set to the agent default if `None`.
+            name: Name of the tool, inferred from the function if `None`.
+            description: Description of the tool, inferred from the function if `None`.
+        """
+        f = _pydantic.function_schema(function, takes_ctx)
+        self.function = function
+        self.takes_ctx = takes_ctx
+        self.max_retries = max_retries
+        self.name = name or function.__name__
+        self.description = description or f['description']
+        self._is_async = inspect.iscoroutinefunction(self.function)
+        self._single_arg_name = f['single_arg_name']
+        self._positional_fields = f['positional_fields']
+        self._var_positional_field = f['var_positional_field']
+        self._validator = f['validator']
+        self._json_schema = f['json_schema']
+
+    @staticmethod
+    def infer(function: ToolFuncEither[A, ...] | Tool[A]) -> Tool[A]:
+        """Create a tool from a pure function, inferring whether it takes `RunContext` as its first argument.
+
+        Args:
+            function: The tool function to wrap; or for convenience, a `Tool` instance.
+
+        Returns:
+            A new `Tool` instance.
+        """
+        if isinstance(function, Tool):
+            return function
+        else:
+            return Tool(function, takes_ctx=_pydantic.takes_ctx(function))
+
+    def reset(self) -> None:
+        """Reset the current retry count."""
+        self._current_retry = 0
+
+    async def run(self, deps: AgentDeps, message: messages.ToolCall) -> messages.Message:
+        """Run the tool function asynchronously."""
+        try:
+            if isinstance(message.args, messages.ArgsJson):
+                args_dict = self._validator.validate_json(message.args.args_json)
+            else:
+                args_dict = self._validator.validate_python(message.args.args_dict)
+        except ValidationError as e:
+            return self._on_error(e, message)
+
+        args, kwargs = self._call_args(deps, args_dict, message)
+        try:
+            if self._is_async:
+                function = cast(Callable[[Any], Awaitable[str]], self.function)
+                response_content = await function(*args, **kwargs)
+            else:
+                function = cast(Callable[[Any], str], self.function)
+                response_content = await _utils.run_in_executor(function, *args, **kwargs)
+        except ModelRetry as e:
+            return self._on_error(e, message)
+
+        self._current_retry = 0
+        return messages.ToolReturn(
+            tool_name=message.tool_name,
+            content=response_content,
+            tool_id=message.tool_id,
+        )
+
+    @property
+    def json_schema(self) -> _utils.ObjectJsonSchema:
+        return self._json_schema
+
+    @property
+    def outer_typed_dict_key(self) -> str | None:
+        return None
+
+    def _call_args(
+        self, deps: AgentDeps, args_dict: dict[str, Any], message: messages.ToolCall
+    ) -> tuple[list[Any], dict[str, Any]]:
+        if self._single_arg_name:
+            args_dict = {self._single_arg_name: args_dict}
+
+        args = [RunContext(deps, self._current_retry, message.tool_name)] if self.takes_ctx else []
+        for positional_field in self._positional_fields:
+            args.append(args_dict.pop(positional_field))
+        if self._var_positional_field:
+            args.extend(args_dict.pop(self._var_positional_field))
+
+        return args, args_dict
+
+    def _on_error(self, exc: ValidationError | ModelRetry, call_message: messages.ToolCall) -> messages.RetryPrompt:
+        self._current_retry += 1
+        if self.max_retries is None or self._current_retry > self.max_retries:
+            raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {self.max_retries}') from exc
+        else:
+            if isinstance(exc, ValidationError):
+                content = exc.errors(include_url=False)
+            else:
+                content = exc.message
+            return messages.RetryPrompt(
+                tool_name=call_message.tool_name,
+                content=content,
+                tool_id=call_message.tool_id,
+            )

{pydantic_ai_slim-0.0.7.dist-info → pydantic_ai_slim-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pydantic-ai-slim
-Version: 0.0.7
+Version: 0.0.9
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License: MIT

pydantic_ai_slim-0.0.9.dist-info/RECORD

@@ -0,0 +1,22 @@
+pydantic_ai/__init__.py,sha256=a29NqQz0JyW4BoCjcRh23fBBfwY17_n57moE4QrFWM4,324
+pydantic_ai/_griffe.py,sha256=pRjCJ6B1hhx6k46XJgl9zF6aRYxRmqEZKFok8unp4Iw,3449
+pydantic_ai/_pydantic.py,sha256=oFfcHDv_wuL1NQ7mCzVHvP1HBaVzyvb7xS-_Iiri_tA,8491
+pydantic_ai/_result.py,sha256=wzcfwDpr_sro1Vkn3DkyIhCXMHTReDxL_ZYm50JzdRI,9667
+pydantic_ai/_system_prompt.py,sha256=vFT0y9Wykl5veGMgLLkGRYiHQrgdW2BZ1rwMn4izjjo,1085
+pydantic_ai/_utils.py,sha256=eNb7f3-ZQC8WDEa87iUcXGQ-lyuutFQG-5yBCMD4Vvs,8227
+pydantic_ai/agent.py,sha256=bHjw1rskPjkAQc8BdoB6AbWjQVTQK0cDDFu7v3EQXqE,34972
+pydantic_ai/exceptions.py,sha256=ko_47M0k6Rhg9mUC9P1cj7N4LCH6cC0pEsF65A2vL-U,1561
+pydantic_ai/messages.py,sha256=I0_CPXDIGGSy-PXHuKq540oAXYOO9uyylpsfSsE4vLs,7032
+pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pydantic_ai/result.py,sha256=UB3vFOcDAceeNLXh_f_3PKy2J0A6FIHdgDFJxPH6OJk,13651
+pydantic_ai/tools.py,sha256=rzchmsEPYtUzBBFoeeJKC1ct36iqOxFOY6kfKqItCCs,8210
+pydantic_ai/models/__init__.py,sha256=_Mz_32WGlAf4NlxXfdQ-EAaY_bDOk10gIc5HmTAO_ts,10318
+pydantic_ai/models/function.py,sha256=Mzc-zXnb2RayWAA8N9NS7KGF49do1S-VW3U9fkc661o,10045
+pydantic_ai/models/gemini.py,sha256=ruO4tnnpDDuHThg7jUOphs8I_KXBJH7gfDMluliED8E,26606
+pydantic_ai/models/groq.py,sha256=Tx2yU3ysmPLBmWGsjzES-XcumzrsoBtB7spCnJBlLiM,14947
+pydantic_ai/models/openai.py,sha256=5ihH25CrS0tnZNW-BZw4GyPe8V-IxIHWw3B9ulPVjQE,14931
+pydantic_ai/models/test.py,sha256=q1wch_E7TSb4qx9PCcP1YyBGZx567MGlAQhlAlON0S8,14463
+pydantic_ai/models/vertexai.py,sha256=5wI8y2YjeRgSE51uKy5OtevQkks65uEbxIUAs5EGBaI,9161
+pydantic_ai_slim-0.0.9.dist-info/METADATA,sha256=NI561tFX5Xwjdak9lEYgMNnXOuKvGY8VZxmTNrt7k5k,2561
+pydantic_ai_slim-0.0.9.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+pydantic_ai_slim-0.0.9.dist-info/RECORD,,

pydantic_ai/_tool.py

@@ -1,112 +0,0 @@
-from __future__ import annotations as _annotations
-
-import inspect
-from collections.abc import Awaitable
-from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, cast
-
-from pydantic import ValidationError
-from pydantic_core import SchemaValidator
-
-from . import _pydantic, _utils, messages
-from .dependencies import AgentDeps, RunContext, ToolContextFunc, ToolParams, ToolPlainFunc
-from .exceptions import ModelRetry, UnexpectedModelBehavior
-
-# Usage `ToolEitherFunc[AgentDependencies, P]`
-ToolEitherFunc = _utils.Either[ToolContextFunc[AgentDeps, ToolParams], ToolPlainFunc[ToolParams]]
-
-
-@dataclass(init=False)
-class Tool(Generic[AgentDeps, ToolParams]):
-    """A tool function for an agent."""
-
-    name: str
-    description: str
-    function: ToolEitherFunc[AgentDeps, ToolParams] = field(repr=False)
-    is_async: bool
-    single_arg_name: str | None
-    positional_fields: list[str]
-    var_positional_field: str | None
-    validator: SchemaValidator = field(repr=False)
-    json_schema: _utils.ObjectJsonSchema
-    max_retries: int
-    _current_retry: int = 0
-    outer_typed_dict_key: str | None = None
-
-    def __init__(self, function: ToolEitherFunc[AgentDeps, ToolParams], retries: int):
-        """Build a Tool dataclass from a function."""
-        self.function = function
-        # noinspection PyTypeChecker
-        f = _pydantic.function_schema(function)
-        raw_function = function.whichever()
-        self.name = raw_function.__name__
-        self.description = f['description']
-        self.is_async = inspect.iscoroutinefunction(raw_function)
-        self.single_arg_name = f['single_arg_name']
-        self.positional_fields = f['positional_fields']
-        self.var_positional_field = f['var_positional_field']
-        self.validator = f['validator']
-        self.json_schema = f['json_schema']
-        self.max_retries = retries
-
-    def reset(self) -> None:
-        """Reset the current retry count."""
-        self._current_retry = 0
-
-    async def run(self, deps: AgentDeps, message: messages.ToolCall) -> messages.Message:
-        """Run the tool function asynchronously."""
-        try:
-            if isinstance(message.args, messages.ArgsJson):
-                args_dict = self.validator.validate_json(message.args.args_json)
-            else:
-                args_dict = self.validator.validate_python(message.args.args_dict)
-        except ValidationError as e:
-            return self._on_error(e, message)
-
-        args, kwargs = self._call_args(deps, args_dict, message)
-        try:
-            if self.is_async:
-                function = cast(Callable[[Any], Awaitable[str]], self.function.whichever())
-                response_content = await function(*args, **kwargs)
-            else:
-                function = cast(Callable[[Any], str], self.function.whichever())
-                response_content = await _utils.run_in_executor(function, *args, **kwargs)
-        except ModelRetry as e:
-            return self._on_error(e, message)
-
-        self._current_retry = 0
-        return messages.ToolReturn(
-            tool_name=message.tool_name,
-            content=response_content,
-            tool_id=message.tool_id,
-        )
-
-    def _call_args(
-        self, deps: AgentDeps, args_dict: dict[str, Any], message: messages.ToolCall
-    ) -> tuple[list[Any], dict[str, Any]]:
-        if self.single_arg_name:
-            args_dict = {self.single_arg_name: args_dict}
-
-        args = [RunContext(deps, self._current_retry, message.tool_name)] if self.function.is_left() else []
-        for positional_field in self.positional_fields:
-            args.append(args_dict.pop(positional_field))
-        if self.var_positional_field:
-            args.extend(args_dict.pop(self.var_positional_field))
-
-        return args, args_dict
-
-    def _on_error(self, exc: ValidationError | ModelRetry, call_message: messages.ToolCall) -> messages.RetryPrompt:
-        self._current_retry += 1
-        if self._current_retry > self.max_retries:
-            # TODO custom error with details of the tool
-            raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {self.max_retries}') from exc
-        else:
-            if isinstance(exc, ValidationError):
-                content = exc.errors(include_url=False)
-            else:
-                content = exc.message
-            return messages.RetryPrompt(
-                tool_name=call_message.tool_name,
-                content=content,
-                tool_id=call_message.tool_id,
-            )

pydantic_ai/dependencies.py

@@ -1,83 +0,0 @@
-from __future__ import annotations as _annotations
-
-from collections.abc import Awaitable, Mapping, Sequence
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
-
-from typing_extensions import Concatenate, ParamSpec, TypeAlias
-
-if TYPE_CHECKING:
-    from .result import ResultData
-else:
-    ResultData = Any
-
-__all__ = (
-    'AgentDeps',
-    'RunContext',
-    'ResultValidatorFunc',
-    'SystemPromptFunc',
-    'ToolReturnValue',
-    'ToolContextFunc',
-    'ToolPlainFunc',
-    'ToolParams',
-    'JsonData',
-)
-
-AgentDeps = TypeVar('AgentDeps')
-"""Type variable for agent dependencies."""
-
-
-@dataclass
-class RunContext(Generic[AgentDeps]):
-    """Information about the current call."""
-
-    deps: AgentDeps
-    """Dependencies for the agent."""
-    retry: int
-    """Number of retries so far."""
-    tool_name: str | None
-    """Name of the tool being called."""
-
-
-ToolParams = ParamSpec('ToolParams')
-"""Retrieval function param spec."""
-
-SystemPromptFunc = Union[
-    Callable[[RunContext[AgentDeps]], str],
-    Callable[[RunContext[AgentDeps]], Awaitable[str]],
-    Callable[[], str],
-    Callable[[], Awaitable[str]],
-]
-"""A function that may or maybe not take `RunContext` as an argument, and may or may not be async.
-
-Usage `SystemPromptFunc[AgentDeps]`.
-"""
-
-ResultValidatorFunc = Union[
-    Callable[[RunContext[AgentDeps], ResultData], ResultData],
-    Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]],
-    Callable[[ResultData], ResultData],
-    Callable[[ResultData], Awaitable[ResultData]],
-]
-"""
-A function that always takes `ResultData` and returns `ResultData`,
-but may or maybe not take `CallInfo` as a first argument, and may or may not be async.
-
-Usage `ResultValidator[AgentDeps, ResultData]`.
-"""
-
-JsonData: TypeAlias = 'None | str | int | float | Sequence[JsonData] | Mapping[str, JsonData]'
-"""Type representing any JSON data."""
-
-ToolReturnValue = Union[JsonData, Awaitable[JsonData]]
-"""Return value of a tool function."""
-ToolContextFunc = Callable[Concatenate[RunContext[AgentDeps], ToolParams], ToolReturnValue]
-"""A tool function that takes `RunContext` as the first argument.
-
-Usage `ToolContextFunc[AgentDeps, ToolParams]`.
-"""
-ToolPlainFunc = Callable[ToolParams, ToolReturnValue]
-"""A tool function that does not take `RunContext` as the first argument.
-
-Usage `ToolPlainFunc[ToolParams]`.
-"""

pydantic_ai_slim-0.0.7.dist-info/RECORD

@@ -1,23 +0,0 @@
-pydantic_ai/__init__.py,sha256=KaTzG8uBSEpfxk_y2a_O_R2Xa53GXYfiigjWtD4PCeI,317
-pydantic_ai/_griffe.py,sha256=pRjCJ6B1hhx6k46XJgl9zF6aRYxRmqEZKFok8unp4Iw,3449
-pydantic_ai/_pydantic.py,sha256=j1kObPIUDwn1VOHbJBwMFbWLxHM_OYRH7GFAda68ZC0,8010
-pydantic_ai/_result.py,sha256=cAqfPipK39cz-p-ftlJ83Q5_LI1TRb3-HH-iivb5rEM,9674
-pydantic_ai/_system_prompt.py,sha256=63egOej8zHsDVOInPayn0EEEDXKd0HVAbbrqXUTV96s,1092
-pydantic_ai/_tool.py,sha256=5Q9XaGOEXbyOLS644osB1AA5EMoJkr4eYK60MVZo0Z8,4528
-pydantic_ai/_utils.py,sha256=eNb7f3-ZQC8WDEa87iUcXGQ-lyuutFQG-5yBCMD4Vvs,8227
-pydantic_ai/agent.py,sha256=RhLl6tf6kogtauqQr3U1ufg712k-g002YPcio9-BPB4,34242
-pydantic_ai/dependencies.py,sha256=EHvD68AFkItxMnfHzJLG7T_AD1RGI2MZOfzm1v89hGQ,2399
-pydantic_ai/exceptions.py,sha256=ko_47M0k6Rhg9mUC9P1cj7N4LCH6cC0pEsF65A2vL-U,1561
-pydantic_ai/messages.py,sha256=tyBs2ucqIrKjjMIX2bvjMaZ6BMxim4VBkmCzTbeH-AI,7493
-pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=Gs2ZjuFJgONGJm8M5R2mGph5-lUFLBg7FxVrs2CVDPs,13525
-pydantic_ai/models/__init__.py,sha256=Cx8PjEsi5gkNOVQic32sf4CmM-A3pRu1LcjpM6poiBI,10138
-pydantic_ai/models/function.py,sha256=Mzc-zXnb2RayWAA8N9NS7KGF49do1S-VW3U9fkc661o,10045
-pydantic_ai/models/gemini.py,sha256=2mD5MJT7qaiMAWrySLjeUFSVoMYhYOS4V7ueTOxzkdA,26472
-pydantic_ai/models/groq.py,sha256=Tx2yU3ysmPLBmWGsjzES-XcumzrsoBtB7spCnJBlLiM,14947
-pydantic_ai/models/openai.py,sha256=5ihH25CrS0tnZNW-BZw4GyPe8V-IxIHWw3B9ulPVjQE,14931
-pydantic_ai/models/test.py,sha256=HOS3_u0n3nYWJmBFPsAFpZ9gfHiZLauz24o825q3e9M,14443
-pydantic_ai/models/vertexai.py,sha256=xHatvwRn7_vyqmp3aDtHqryYAq8NoeXgXVATkj7yHuw,9088
-pydantic_ai_slim-0.0.7.dist-info/METADATA,sha256=UeP3i7vUZv1MCKG-VGLZUwEO8kbcCzzLpD6RsAN_Vbs,2561
-pydantic_ai_slim-0.0.7.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-pydantic_ai_slim-0.0.7.dist-info/RECORD,,