pydantic-ai-slim 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

pydantic_ai/__init__.py

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

pydantic_ai/_pydantic.py

@@ -8,7 +8,7 @@ from __future__ import annotations as _annotations
 from inspect import Parameter, signature
 from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast, get_origin
 
-from pydantic import ConfigDict, TypeAdapter
+from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
 from pydantic._internal._config import ConfigWrapper
 from pydantic.fields import FieldInfo
@@ -23,7 +23,7 @@ if TYPE_CHECKING:
     from .tools import ObjectJsonSchema
 
 
-__all__ = 'function_schema', 'LazyTypeAdapter'
+__all__ = ('function_schema',)
 
 
 class FunctionSchema(TypedDict):
@@ -138,14 +138,14 @@ def function_schema(function: Callable[..., Any], takes_ctx: bool) -> FunctionSc
     json_schema = GenerateJsonSchema().generate(schema)
 
     # workaround for https://github.com/pydantic/pydantic/issues/10785
-    # if we build a custom TypeDict schema (matches when `single_arg_name` is None), we manually set
+    # if we build a custom TypeDict schema (matches when `single_arg_name is None`), we manually set
     # `additionalProperties` in the JSON Schema
     if single_arg_name is None:
         json_schema['additionalProperties'] = bool(var_kwargs_schema)
-
-    # instead of passing `description` through in core_schema, we just add it here
-    if description:
-        json_schema = {'description': description} | json_schema
+    elif not description:
+        # if the tool description is not set, and we have a single parameter, take the description from that
+        # and set it on the tool
+        description = json_schema.pop('description', None)
 
     return FunctionSchema(
         description=description,
@@ -214,21 +214,3 @@ def _is_call_ctx(annotation: Any) -> bool:
     return annotation is RunContext or (
         _typing_extra.is_generic_alias(annotation) and get_origin(annotation) is RunContext
     )
-
-
-if TYPE_CHECKING:
-    LazyTypeAdapter = TypeAdapter
-else:
-
-    class LazyTypeAdapter:
-        __slots__ = '_args', '_kwargs', '_type_adapter'
-
-        def __init__(self, *args, **kwargs):
-            self._args = args
-            self._kwargs = kwargs
-            self._type_adapter = None
-
-        def __getattr__(self, item):
-            if self._type_adapter is None:
-                self._type_adapter = TypeAdapter(*self._args, **self._kwargs)
-            return getattr(self._type_adapter, item)

pydantic_ai/_result.py

@@ -3,16 +3,15 @@ from __future__ import annotations as _annotations
 import inspect
 import sys
 import types
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Iterable
 from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin
 
 from pydantic import TypeAdapter, ValidationError
 from typing_extensions import Self, TypeAliasType, TypedDict
 
-from . import _utils, messages
+from . import _utils, messages as _messages
 from .exceptions import ModelRetry
-from .messages import ModelStructuredResponse, ToolCall
 from .result import ResultData
 from .tools import AgentDeps, ResultValidatorFunc, RunContext, ToolDefinition
 
@@ -28,21 +27,24 @@ class ResultValidator(Generic[AgentDeps, ResultData]):
         self._is_async = inspect.iscoroutinefunction(self.function)
 
     async def validate(
-        self, result: ResultData, deps: AgentDeps, retry: int, tool_call: messages.ToolCall | None
+        self,
+        result: ResultData,
+        tool_call: _messages.ToolCallPart | None,
+        run_context: RunContext[AgentDeps],
     ) -> ResultData:
         """Validate a result but calling the function.
 
         Args:
             result: The result data after Pydantic validation the message content.
-            deps: The agent dependencies.
-            retry: The current retry number.
             tool_call: The original tool call message, `None` if there was no tool call.
+            run_context: The current run context.
 
         Returns:
             Result of either the validated result data (ok) or a retry message (Err).
         """
         if self._takes_ctx:
-            args = RunContext(deps, retry, tool_call.tool_name if tool_call else None), result
+            ctx = run_context.replace_with(tool_name=tool_call.tool_name if tool_call else None)
+            args = ctx, result
         else:
             args = (result,)
 
@@ -54,10 +56,10 @@ class ResultValidator(Generic[AgentDeps, ResultData]):
                 function = cast(Callable[[Any], ResultData], self.function)
                 result_data = await _utils.run_in_executor(function, *args)
         except ModelRetry as r:
-            m = messages.RetryPrompt(content=r.message)
+            m = _messages.RetryPromptPart(content=r.message)
             if tool_call is not None:
                 m.tool_name = tool_call.tool_name
-                m.tool_id = tool_call.tool_id
+                m.tool_call_id = tool_call.tool_call_id
             raise ToolRetryError(m) from r
         else:
             return result_data
@@ -66,7 +68,7 @@ class ResultValidator(Generic[AgentDeps, ResultData]):
 class ToolRetryError(Exception):
     """Internal exception used to signal a `ToolRetry` message should be returned to the LLM."""
 
-    def __init__(self, tool_retry: messages.RetryPrompt):
+    def __init__(self, tool_retry: _messages.RetryPromptPart):
         self.tool_retry = tool_retry
         super().__init__()
 
@@ -108,11 +110,24 @@ class ResultSchema(Generic[ResultData]):
 
         return cls(tools=tools, allow_text_result=allow_text_result)
 
-    def find_tool(self, message: ModelStructuredResponse) -> tuple[ToolCall, ResultTool[ResultData]] | None:
+    def find_named_tool(
+        self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
+    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultData]] | None:
+        """Find a tool that matches one of the calls, with a specific name."""
+        for part in parts:
+            if isinstance(part, _messages.ToolCallPart):
+                if part.tool_name == tool_name:
+                    return part, self.tools[tool_name]
+
+    def find_tool(
+        self,
+        parts: Iterable[_messages.ModelResponsePart],
+    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultData]] | None:
         """Find a tool that matches one of the calls."""
-        for call in message.calls:
-            if result := self.tools.get(call.tool_name):
-                return call, result
+        for part in parts:
+            if isinstance(part, _messages.ToolCallPart):
+                if result := self.tools.get(part.tool_name):
+                    return part, result
 
     def tool_names(self) -> list[str]:
         """Return the names of the tools."""
@@ -167,7 +182,7 @@ class ResultTool(Generic[ResultData]):
         )
 
     def validate(
-        self, tool_call: messages.ToolCall, allow_partial: bool = False, wrap_validation_errors: bool = True
+        self, tool_call: _messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
     ) -> ResultData:
         """Validate a result message.
 
@@ -181,7 +196,7 @@ class ResultTool(Generic[ResultData]):
         """
         try:
             pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-            if isinstance(tool_call.args, messages.ArgsJson):
+            if isinstance(tool_call.args, _messages.ArgsJson):
                 result = self.type_adapter.validate_json(
                     tool_call.args.args_json or '', experimental_allow_partial=pyd_allow_partial
                 )
@@ -191,10 +206,10 @@ class ResultTool(Generic[ResultData]):
                 )
         except ValidationError as e:
             if wrap_validation_errors:
-                m = messages.RetryPrompt(
+                m = _messages.RetryPromptPart(
                     tool_name=tool_call.tool_name,
                     content=e.errors(include_url=False),
-                    tool_id=tool_call.tool_id,
+                    tool_call_id=tool_call.tool_call_id,
                 )
                 raise ToolRetryError(m) from e
             else:

pydantic_ai/_system_prompt.py

@@ -19,9 +19,9 @@ class SystemPromptRunner(Generic[AgentDeps]):
         self._takes_ctx = len(inspect.signature(self.function).parameters) > 0
         self._is_async = inspect.iscoroutinefunction(self.function)
 
-    async def run(self, deps: AgentDeps) -> str:
+    async def run(self, run_context: RunContext[AgentDeps]) -> str:
         if self._takes_ctx:
-            args = (RunContext(deps, 0),)
+            args = (run_context,)
         else:
             args = ()
 

pydantic_ai/_utils.py

@@ -15,6 +15,7 @@ from pydantic.json_schema import JsonSchemaValue
 from typing_extensions import ParamSpec, TypeAlias, TypeGuard, is_typeddict
 
 if TYPE_CHECKING:
+    from .messages import RetryPromptPart, ToolCallPart, ToolReturnPart
     from .tools import ObjectJsonSchema
 
 _P = ParamSpec('_P')
@@ -87,7 +88,7 @@ class Either(Generic[Left, Right]):
 
     Usage:
 
-    ```py
+    ```python
     if left_thing := either.left:
         use_left(left_thing.value)
     else:
@@ -146,7 +147,7 @@ async def group_by_temporal(
 
     Usage:
 
-    ```py
+    ```python
     async with group_by_temporal(yield_groups(), 0.1) as groups_iter:
         async for groups in groups_iter:
             print(groups)
@@ -254,3 +255,9 @@ def sync_anext(iterator: Iterator[T]) -> T:
 
 def now_utc() -> datetime:
     return datetime.now(tz=timezone.utc)
+
+
+def guard_tool_call_id(t: ToolCallPart | ToolReturnPart | RetryPromptPart, model_source: str) -> str:
+    """Type guard that checks a `tool_call_id` is not None both for static typing and runtime."""
+    assert t.tool_call_id is not None, f'{model_source} requires `tool_call_id` to be set: {t}'
+    return t.tool_call_id