pydantic-ai-slim 0.0.8__py3-none-any.whl → 0.0.10__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,9 +1,12 @@
 from __future__ import annotations as _annotations
 
 import asyncio
+import dataclasses
+import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
+from types import FrameType
 from typing import Any, Callable, Generic, cast, final, overload
 
 import logfire_api
@@ -12,15 +15,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 +36,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]`.
@@ -54,11 +56,16 @@ class Agent(Generic[AgentDeps, ResultData]):
     # dataclass fields mostly for my sanity — knowing what attributes are available
     model: models.Model | models.KnownModelName | None
     """The default model configured for this agent."""
+    name: str | None
+    """The name of the agent, used for logging.
+
+    If `None`, we try to infer the agent name from the call frame when the agent is first run.
+    """
     _result_schema: _result.ResultSchema[ResultData] | None = field(repr=False)
     _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)
@@ -79,10 +86,12 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[ResultData] = str,
         system_prompt: str | Sequence[str] = (),
         deps_type: type[AgentDeps] = NoneType,
+        name: str | None = None,
         retries: int = 1,
         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.
@@ -97,10 +106,14 @@ class Agent(Generic[AgentDeps, ResultData]):
                 parameterize the agent, and therefore get the best out of static type checking.
                 If you're not using deps, but want type checking to pass, you can set `deps=None` to satisfy Pyright
                 or add a type hint `: Agent[None, <return type>]`.
+            name: The name of the agent, used for logging. If `None`, we try to infer the agent name from the call frame
+                when the agent is first run.
             retries: The default number of retries to allow before raising an error.
             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`
@@ -112,6 +125,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         else:
             self.model = models.infer_model(model)
 
+        self.name = name
         self._result_schema = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
@@ -119,9 +133,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
@@ -134,6 +150,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         message_history: list[_messages.Message] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
+        infer_name: bool = True,
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt in async mode.
 
@@ -142,16 +159,19 @@ class Agent(Generic[AgentDeps, ResultData]):
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
             deps: Optional dependencies to use for this run.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
             The result of the run.
         """
+        if infer_name and self.name is None:
+            self._infer_name(inspect.currentframe())
         model_used, custom_model, agent_model = await self._get_agent_model(model)
 
         deps = self._get_deps(deps)
 
         with _logfire.span(
-            'agent run {prompt=}',
+            '{agent.name} run {prompt=}',
             prompt=user_prompt,
             agent=self,
             custom_model=custom_model,
@@ -203,21 +223,28 @@ class Agent(Generic[AgentDeps, ResultData]):
         message_history: list[_messages.Message] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
+        infer_name: bool = True,
     ) -> 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.
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
             deps: Optional dependencies to use for this run.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
             The result of the run.
         """
-        return asyncio.run(self.run(user_prompt, message_history=message_history, model=model, deps=deps))
+        if infer_name and self.name is None:
+            self._infer_name(inspect.currentframe())
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(
+            self.run(user_prompt, message_history=message_history, model=model, deps=deps, infer_name=False)
+        )
 
     @asynccontextmanager
     async def run_stream(
@@ -227,6 +254,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         message_history: list[_messages.Message] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
+        infer_name: bool = True,
     ) -> AsyncIterator[result.StreamedRunResult[AgentDeps, ResultData]]:
         """Run the agent with a user prompt in async mode, returning a streamed response.
 
@@ -235,16 +263,21 @@ class Agent(Generic[AgentDeps, ResultData]):
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
             deps: Optional dependencies to use for this run.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
             The result of the run.
         """
+        if infer_name and self.name is None:
+            # f_back because `asynccontextmanager` adds one frame
+            if frame := inspect.currentframe():  # pragma: no branch
+                self._infer_name(frame.f_back)
         model_used, custom_model, agent_model = await self._get_agent_model(model)
 
         deps = self._get_deps(deps)
 
         with _logfire.span(
-            'agent run stream {prompt=}',
+            '{agent.name} run stream {prompt=}',
             prompt=user_prompt,
             agent=self,
             custom_model=custom_model,
@@ -354,7 +387,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
@@ -405,7 +438,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
@@ -438,22 +471,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.
 
@@ -490,27 +523,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.
@@ -547,28 +580,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
 
@@ -786,6 +825,19 @@ class Agent(Generic[AgentDeps, ResultData]):
         else:
             return deps
 
+    def _infer_name(self, function_frame: FrameType | None) -> None:
+        """Infer the agent name from the call frame.
+
+        Usage should be `self._infer_name(inspect.currentframe())`.
+        """
+        assert self.name is None, 'Name already set'
+        if function_frame is not None:  # pragma: no branch
+            if parent_frame := function_frame.f_back:  # pragma: no branch
+                for name, item in parent_frame.f_locals.items():
+                    if item is self:
+                        self.name = name
+                        return
+
 
 @dataclass
 class _MarkFinalResult(Generic[ResultData]):

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

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/result.py

@@ -9,7 +9,7 @@ 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',

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.8.dist-info → pydantic_ai_slim-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pydantic-ai-slim
-Version: 0.0.8
+Version: 0.0.10
 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.10.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=dB2_JshYBjK04fmzJP79wmoKJwMuEBaLmjIRVdwrISM,36854
+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.10.dist-info/METADATA,sha256=2-sEVelPFDeLZwYC4o8ZgGvyXvJsVXZIkj3Mx3wXs6g,2562
+pydantic_ai_slim-0.0.10.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+pydantic_ai_slim-0.0.10.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.8.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=r5DI4ZBqYE67GOMEEu-LXrTa5ty2AchW4szotwm5Qis,34338
-pydantic_ai/dependencies.py,sha256=EHvD68AFkItxMnfHzJLG7T_AD1RGI2MZOfzm1v89hGQ,2399
-pydantic_ai/exceptions.py,sha256=ko_47M0k6Rhg9mUC9P1cj7N4LCH6cC0pEsF65A2vL-U,1561
-pydantic_ai/messages.py,sha256=FFTQ9Bo2Ct4bLuyJF-M9xkeraw05I--NC_ieR6oGtTM,7587
-pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=qsanb7v4qJ4pJdkdsqpy68kuZ1WNCpQB5jcXeTwGpe0,13658
-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=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.8.dist-info/METADATA,sha256=CmpvlEAUyaWaPbdRCbFuypVLJ8yNC3TwZ0jgvlR9yps,2561
-pydantic_ai_slim-0.0.8.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-pydantic_ai_slim-0.0.8.dist-info/RECORD,,