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

pydantic_ai/__init__.py

@@ -1,11 +1,12 @@
 from importlib.metadata import version
 
-from .agent import Agent
+from .agent import Agent, capture_run_messages
 from .exceptions import AgentRunError, ModelRetry, UnexpectedModelBehavior, UsageLimitExceeded, UserError
 from .tools import RunContext, Tool
 
 __all__ = (
     'Agent',
+    'capture_run_messages',
     'RunContext',
     'Tool',
     'AgentRunError',

pydantic_ai/_griffe.py

@@ -4,8 +4,7 @@ import re
 from inspect import Signature
 from typing import Any, Callable, Literal, cast
 
-from _griffe.enumerations import DocstringSectionKind
-from _griffe.models import Docstring, Object as GriffeObject
+from griffe import Docstring, DocstringSectionKind, Object as GriffeObject
 
 DocstringStyle = Literal['google', 'numpy', 'sphinx']
 

pydantic_ai/_result.py

@@ -12,8 +12,8 @@ from typing_extensions import Self, TypeAliasType, TypedDict
 
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
-from .result import ResultData
-from .tools import AgentDeps, ResultValidatorFunc, RunContext, ToolDefinition
+from .result import ResultData, ResultValidatorFunc
+from .tools import AgentDeps, RunContext, ToolDefinition
 
 
 @dataclass

pydantic_ai/agent.py

@@ -5,12 +5,12 @@ import dataclasses
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
-from dataclasses import dataclass, field
+from contextvars import ContextVar
 from types import FrameType
 from typing import Any, Callable, Generic, Literal, cast, final, overload
 
 import logfire_api
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from . import (
     _result,
@@ -35,10 +35,20 @@ from .tools import (
     ToolPrepareFunc,
 )
 
-__all__ = ('Agent',)
+__all__ = 'Agent', 'capture_run_messages', 'EndStrategy'
 
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
+# while waiting for https://github.com/pydantic/logfire/issues/745
+try:
+    import logfire._internal.stack_info
+except ImportError:
+    pass
+else:
+    from pathlib import Path
+
+    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
+
 NoneType = type(None)
 EndStrategy = Literal['early', 'exhaustive']
 """The strategy for handling multiple tool calls when a final result is found.
@@ -49,7 +59,7 @@ EndStrategy = Literal['early', 'exhaustive']
 
 
 @final
-@dataclass(init=False)
+@dataclasses.dataclass(init=False)
 class Agent(Generic[AgentDeps, ResultData]):
     """Class for defining "agents" - a way to have a specific type of "conversation" with an LLM.
 
@@ -89,23 +99,17 @@ class Agent(Generic[AgentDeps, ResultData]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    last_run_messages: list[_messages.ModelMessage] | None
-    """The messages from the last run, useful when a run raised an exception.
-
-    Note: these are not used by the agent, e.g. in future runs, they are just stored for developers' convenience.
-    """
-
-    _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, 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)
-    _max_result_retries: int = field(repr=False)
-    _override_deps: _utils.Option[AgentDeps] = field(default=None, repr=False)
-    _override_model: _utils.Option[models.Model] = field(default=None, repr=False)
+    _result_schema: _result.ResultSchema[ResultData] | None = dataclasses.field(repr=False)
+    _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = dataclasses.field(repr=False)
+    _allow_text_result: bool = dataclasses.field(repr=False)
+    _system_prompts: tuple[str, ...] = dataclasses.field(repr=False)
+    _function_tools: dict[str, Tool[AgentDeps]] = dataclasses.field(repr=False)
+    _default_retries: int = dataclasses.field(repr=False)
+    _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDeps]] = dataclasses.field(repr=False)
+    _deps_type: type[AgentDeps] = dataclasses.field(repr=False)
+    _max_result_retries: int = dataclasses.field(repr=False)
+    _override_deps: _utils.Option[AgentDeps] = dataclasses.field(default=None, repr=False)
+    _override_model: _utils.Option[models.Model] = dataclasses.field(default=None, repr=False)
 
     def __init__(
         self,
@@ -161,7 +165,6 @@ class Agent(Generic[AgentDeps, ResultData]):
         self.end_strategy = end_strategy
         self.name = name
         self.model_settings = model_settings
-        self.last_run_messages = None
         self._result_schema = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
@@ -190,6 +193,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
+        usage: result.Usage | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt in async mode.
@@ -212,6 +216,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
@@ -219,7 +224,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
-        model_used, mode_selection = await self._get_model(model)
+        model_used = await self._get_model(model)
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
@@ -228,40 +233,36 @@ class Agent(Generic[AgentDeps, ResultData]):
             '{agent_name} run {prompt=}',
             prompt=user_prompt,
             agent=self,
-            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
-            run_context = RunContext(deps, 0, [], None, model_used)
+            run_context = RunContext(deps, model_used, usage or result.Usage(), user_prompt)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            self.last_run_messages = run_context.messages = messages
+            run_context.messages = messages
 
             for tool in self._function_tools.values():
                 tool.current_retry = 0
 
-            usage = result.Usage(requests=0)
             model_settings = merge_model_settings(self.model_settings, model_settings)
             usage_limits = usage_limits or UsageLimits()
 
-            run_step = 0
             while True:
-                usage_limits.check_before_request(usage)
+                usage_limits.check_before_request(run_context.usage)
 
-                run_step += 1
-                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                run_context.run_step += 1
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
                     agent_model = await self._prepare_model(run_context)
 
-                with _logfire.span('model request', run_step=run_step) as model_req_span:
+                with _logfire.span('model request', run_step=run_context.run_step) as model_req_span:
                     model_response, request_usage = await agent_model.request(messages, model_settings)
                     model_req_span.set_attribute('response', model_response)
                     model_req_span.set_attribute('usage', request_usage)
 
                 messages.append(model_response)
-                usage += request_usage
-                usage.requests += 1
-                usage_limits.check_tokens(request_usage)
+                run_context.usage.incr(request_usage, requests=1)
+                usage_limits.check_tokens(run_context.usage)
 
-                with _logfire.span('handle model response', run_step=run_step) as handle_span:
+                with _logfire.span('handle model response', run_step=run_context.run_step) as handle_span:
                     final_result, tool_responses = await self._handle_model_response(model_response, run_context)
 
                     if tool_responses:
@@ -272,10 +273,10 @@ class Agent(Generic[AgentDeps, ResultData]):
                     if final_result is not None:
                         result_data = final_result.data
                         run_span.set_attribute('all_messages', messages)
-                        run_span.set_attribute('usage', usage)
+                        run_span.set_attribute('usage', run_context.usage)
                         handle_span.set_attribute('result', result_data)
                         handle_span.message = 'handle model response -> final result'
-                        return result.RunResult(messages, new_message_index, result_data, usage)
+                        return result.RunResult(messages, new_message_index, result_data, run_context.usage)
                     else:
                         # continue the conversation
                         handle_span.set_attribute('tool_responses', tool_responses)
@@ -291,6 +292,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
+        usage: result.Usage | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt synchronously.
@@ -317,6 +319,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
@@ -332,6 +335,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 deps=deps,
                 model_settings=model_settings,
                 usage_limits=usage_limits,
+                usage=usage,
                 infer_name=False,
             )
         )
@@ -346,6 +350,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
+        usage: result.Usage | None = None,
         infer_name: bool = True,
     ) -> AsyncIterator[result.StreamedRunResult[AgentDeps, ResultData]]:
         """Run the agent with a user prompt in async mode, returning a streamed response.
@@ -369,6 +374,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
 
         Returns:
@@ -378,7 +384,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             # f_back because `asynccontextmanager` adds one frame
             if frame := inspect.currentframe():  # pragma: no branch
                 self._infer_name(frame.f_back)
-        model_used, mode_selection = await self._get_model(model)
+        model_used = await self._get_model(model)
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
@@ -387,32 +393,29 @@ class Agent(Generic[AgentDeps, ResultData]):
             '{agent_name} run stream {prompt=}',
             prompt=user_prompt,
             agent=self,
-            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
-            run_context = RunContext(deps, 0, [], None, model_used)
+            run_context = RunContext(deps, model_used, usage or result.Usage(), user_prompt)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            self.last_run_messages = run_context.messages = messages
+            run_context.messages = messages
 
             for tool in self._function_tools.values():
                 tool.current_retry = 0
 
-            usage = result.Usage()
             model_settings = merge_model_settings(self.model_settings, model_settings)
             usage_limits = usage_limits or UsageLimits()
 
-            run_step = 0
             while True:
-                run_step += 1
-                usage_limits.check_before_request(usage)
+                run_context.run_step += 1
+                usage_limits.check_before_request(run_context.usage)
 
-                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
                     agent_model = await self._prepare_model(run_context)
 
-                with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
+                with _logfire.span('model request {run_step=}', run_step=run_context.run_step) as model_req_span:
                     async with agent_model.request_stream(messages, model_settings) as model_response:
-                        usage.requests += 1
+                        run_context.usage.requests += 1
                         model_req_span.set_attribute('response_type', model_response.__class__.__name__)
                         # We want to end the "model request" span here, but we can't exit the context manager
                         # in the traditional way
@@ -448,7 +451,6 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 yield result.StreamedRunResult(
                                     messages,
                                     new_message_index,
-                                    usage,
                                     usage_limits,
                                     result_stream,
                                     self._result_schema,
@@ -472,8 +474,8 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 handle_span.message = f'handle model response -> {tool_responses_str}'
                                 # the model_response should have been fully streamed by now, we can add its usage
                                 model_response_usage = model_response.usage()
-                                usage += model_response_usage
-                                usage_limits.check_tokens(usage)
+                                run_context.usage.incr(model_response_usage)
+                                usage_limits.check_tokens(run_context.usage)
 
     @contextmanager
     def override(
@@ -614,7 +616,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         #> success (no tool calls)
         ```
         """
-        self._result_validators.append(_result.ResultValidator(func))
+        self._result_validators.append(_result.ResultValidator[AgentDeps, Any](func))
         return func
 
     @overload
@@ -784,14 +786,14 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         self._function_tools[tool.name] = tool
 
-    async def _get_model(self, model: models.Model | models.KnownModelName | None) -> tuple[models.Model, str]:
+    async def _get_model(self, model: models.Model | models.KnownModelName | None) -> models.Model:
         """Create a model configured for this agent.
 
         Args:
             model: model to use for this run, required if `model` was not set when creating the agent.
 
         Returns:
-            a tuple of `(model used, how the model was selected)`
+            The model used
         """
         model_: models.Model
         if some_model := self._override_model:
@@ -802,18 +804,15 @@ class Agent(Generic[AgentDeps, ResultData]):
                     '(Even when `override(model=...)` is customizing the model that will actually be called)'
                 )
             model_ = some_model.value
-            mode_selection = 'override-model'
         elif model is not None:
             model_ = models.infer_model(model)
-            mode_selection = 'custom'
         elif self.model is not None:
             # noinspection PyTypeChecker
             model_ = self.model = models.infer_model(self.model)
-            mode_selection = 'from-agent'
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        return model_, mode_selection
+        return model_
 
     async def _prepare_model(self, run_context: RunContext[AgentDeps]) -> models.AgentModel:
         """Build tools and create an agent model."""
@@ -835,14 +834,25 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _prepare_messages(
         self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[AgentDeps]
     ) -> list[_messages.ModelMessage]:
+        try:
+            ctx_messages = _messages_ctx_var.get()
+        except LookupError:
+            messages: list[_messages.ModelMessage] = []
+        else:
+            if ctx_messages.used:
+                messages = []
+            else:
+                messages = ctx_messages.messages
+                ctx_messages.used = True
+
         if message_history:
             # shallow copy messages
-            messages = message_history.copy()
+            messages.extend(message_history)
             messages.append(_messages.ModelRequest([_messages.UserPromptPart(user_prompt)]))
         else:
             parts = await self._sys_parts(run_context)
             parts.append(_messages.UserPromptPart(user_prompt))
-            messages: list[_messages.ModelMessage] = [_messages.ModelRequest(parts)]
+            messages.append(_messages.ModelRequest(parts))
 
         return messages
 
@@ -864,11 +874,15 @@ class Agent(Generic[AgentDeps, ResultData]):
             else:
                 tool_calls.append(part)
 
-        if texts:
+        # At the moment, we prioritize at least executing tool calls if they are present.
+        # In the future, we'd consider making this configurable at the agent or run level.
+        # This accounts for cases like anthropic returns that might contain a text response
+        # and a tool call response, where the text response just indicates the tool call will happen.
+        if tool_calls:
+            return await self._handle_structured_response(tool_calls, run_context)
+        elif texts:
             text = '\n\n'.join(texts)
             return await self._handle_text_response(text, run_context)
-        elif tool_calls:
-            return await self._handle_structured_response(tool_calls, run_context)
         else:
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
@@ -1115,8 +1129,59 @@ class Agent(Generic[AgentDeps, ResultData]):
                             self.name = name
                             return
 
+    @property
+    @deprecated(
+        'The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.', category=None
+    )
+    def last_run_messages(self) -> list[_messages.ModelMessage]:
+        raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
+
+
+@dataclasses.dataclass
+class _RunMessages:
+    messages: list[_messages.ModelMessage]
+    used: bool = False
+
+
+_messages_ctx_var: ContextVar[_RunMessages] = ContextVar('var')
+
+
+@contextmanager
+def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
+    """Context manager to access the messages used in a [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] call.
+
+    Useful when a run may raise an exception, see [model errors](../agents.md#model-errors) for more information.
+
+    Examples:
+    ```python
+    from pydantic_ai import Agent, capture_run_messages
+
+    agent = Agent('test')
+
+    with capture_run_messages() as messages:
+        try:
+            result = agent.run_sync('foobar')
+        except Exception:
+            print(messages)
+            raise
+    ```
+
+    !!! note
+        If you call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context,
+        `messages` will represent the messages exchanged during the first call only.
+    """
+    try:
+        yield _messages_ctx_var.get().messages
+    except LookupError:
+        messages: list[_messages.ModelMessage] = []
+        token = _messages_ctx_var.set(_RunMessages(messages))
+        try:
+            yield messages
+        finally:
+            _messages_ctx_var.reset(token)
+
 
-@dataclass
+@dataclasses.dataclass
 class _MarkFinalResult(Generic[ResultData]):
     """Marker class to indicate that the result is the final result.
 

pydantic_ai/models/gemini.py

@@ -444,7 +444,8 @@ def _content_model_response(m: ModelResponse) -> _GeminiContent:
         if isinstance(item, ToolCallPart):
             parts.append(_function_call_part_from_call(item))
         elif isinstance(item, TextPart):
-            parts.append(_GeminiTextPart(text=item.content))
+            if item.content:
+                parts.append(_GeminiTextPart(text=item.content))
         else:
             assert_never(item)
     return _GeminiContent(role='model', parts=parts)
@@ -701,7 +702,7 @@ class _GeminiJsonSchema:
 
     def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
         schema.pop('title', None)
-        schema.pop('default', None)
+        default = schema.pop('default', _utils.UNSET)
         if ref := schema.pop('$ref', None):
             # noinspection PyTypeChecker
             key = re.sub(r'^#/\$defs/', '', ref)
@@ -714,8 +715,14 @@ class _GeminiJsonSchema:
             return
 
         if any_of := schema.get('anyOf'):
-            for schema in any_of:
-                self._simplify(schema, refs_stack)
+            for item_schema in any_of:
+                self._simplify(item_schema, refs_stack)
+            if len(any_of) == 2 and {'type': 'null'} in any_of and default is None:
+                for item_schema in any_of:
+                    if item_schema != {'type': 'null'}:
+                        schema.clear()
+                        schema.update(item_schema)
+                        return
 
         type_ = schema.get('type')
 

pydantic_ai/models/mistral.py

@@ -8,6 +8,7 @@ from datetime import datetime, timezone
 from itertools import chain
 from typing import Any, Callable, Literal, Union
 
+import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
 from typing_extensions import assert_never
 
@@ -39,7 +40,6 @@ from . import (
 )
 
 try:
-    from json_repair import repair_json
     from mistralai import (
         UNSET,
         CompletionChunk as MistralCompletionChunk,
@@ -198,11 +198,10 @@ class MistralAgentModel(AgentModel):
         """Create a streaming completion request to the Mistral model."""
         response: MistralEventStreamAsync[MistralCompletionEvent] | None
         mistral_messages = list(chain(*(self._map_message(m) for m in messages)))
-
         model_settings = model_settings or {}
 
         if self.result_tools and self.function_tools or self.function_tools:
-            # Function Calling Mode
+            # Function Calling
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -218,9 +217,9 @@ class MistralAgentModel(AgentModel):
         elif self.result_tools:
             # Json Mode
             parameters_json_schemas = [tool.parameters_json_schema for tool in self.result_tools]
-
             user_output_format_message = self._generate_user_output_format(parameters_json_schemas)
             mistral_messages.append(user_output_format_message)
+
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -270,12 +269,13 @@ class MistralAgentModel(AgentModel):
     @staticmethod
     def _process_response(response: MistralChatCompletionResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
+        assert response.choices, 'Unexpected empty response choice.'
+
         if response.created:
             timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         else:
             timestamp = _now_utc()
 
-        assert response.choices, 'Unexpected empty response choice.'
         choice = response.choices[0]
         content = choice.message.content
         tool_calls = choice.message.tool_calls
@@ -546,20 +546,15 @@ class MistralStreamStructuredResponse(StreamStructuredResponse):
                 calls.append(tool)
 
         elif self._delta_content and self._result_tools:
-            # NOTE: Params set for the most efficient and fastest way.
-            output_json = repair_json(self._delta_content, return_objects=True, skip_json_loads=True)
-            assert isinstance(
-                output_json, dict
-            ), f'Expected repair_json as type dict, invalid type: {type(output_json)}'
+            output_json: dict[str, Any] | None = pydantic_core.from_json(
+                self._delta_content, allow_partial='trailing-strings'
+            )
 
             if output_json:
                 for result_tool in self._result_tools.values():
-                    # NOTE: Additional verification to prevent JSON validation to crash in `result.py`
+                    # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
                     # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
-                    # For example, `return_type=list[str]` expects a 'response' key with value type array of str.
-                    # when `{"response":` then `repair_json` sets `{"response": ""}` (type not found default str)
-                    # when `{"response": {` then `repair_json` sets `{"response": {}}` (type found)
-                    # This ensures it's corrected to `{"response": {}}` and other required parameters and type.
+                    # Example with BaseModel and required fields.
                     if not self._validate_required_json_schema(output_json, result_tool.parameters_json_schema):
                         continue
 

pydantic_ai/models/ollama.py

@@ -71,6 +71,7 @@ class OllamaModel(Model):
         model_name: OllamaModelName,
         *,
         base_url: str | None = 'http://localhost:11434/v1/',
+        api_key: str = 'ollama',
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
     ):
@@ -83,6 +84,8 @@ class OllamaModel(Model):
             model_name: The name of the Ollama model to use. List of models available [here](https://ollama.com/library)
                 You must first download the model (`ollama pull <MODEL-NAME>`) in order to use the model
             base_url: The base url for the ollama requests. The default value is the ollama default
+            api_key: The API key to use for authentication. Defaults to 'ollama' for local instances,
+                but can be customized for proxy setups that require authentication
             openai_client: An existing
                 [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
                 client to use, if provided, `base_url` and `http_client` must be `None`.
@@ -96,7 +99,7 @@ class OllamaModel(Model):
         else:
             # API key is not required for ollama but a value is required to create the client
             http_client_ = http_client or cached_async_http_client()
-            oai_client = AsyncOpenAI(base_url=base_url, api_key='ollama', http_client=http_client_)
+            oai_client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client_)
             self.openai_model = OpenAIModel(model_name=model_name, openai_client=oai_client)
 
     async def agent_model(

pydantic_ai/models/test.py

@@ -16,6 +16,7 @@ from ..messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    ModelResponsePart,
     RetryPromptPart,
     TextPart,
     ToolCallPart,
@@ -177,13 +178,23 @@ class TestAgentModel(AgentModel):
             # check if there are any retry prompts, if so retry them
             new_retry_names = {p.tool_name for p in last_message.parts if isinstance(p, RetryPromptPart)}
             if new_retry_names:
-                return ModelResponse(
-                    parts=[
-                        ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
-                        for name, args in self.tool_calls
-                        if name in new_retry_names
-                    ]
-                )
+                # Handle retries for both function tools and result tools
+                # Check function tools first
+                retry_parts: list[ModelResponsePart] = [
+                    ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
+                    for name, args in self.tool_calls
+                    if name in new_retry_names
+                ]
+                # Check result tools
+                if self.result_tools:
+                    retry_parts.extend(
+                        [
+                            ToolCallPart.from_raw_args(tool.name, self.gen_tool_args(tool))
+                            for tool in self.result_tools
+                            if tool.name in new_retry_names
+                        ]
+                    )
+                return ModelResponse(parts=retry_parts)
 
         if response_text := self.result.left:
             if response_text.value is None:

pydantic_ai/result.py

@@ -2,11 +2,13 @@ from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Awaitable, Callable
+from copy import copy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Generic, TypeVar, cast
+from typing import Generic, Union, cast
 
 import logfire_api
+from typing_extensions import TypeVar
 
 from . import _result, _utils, exceptions, messages as _messages, models
 from .settings import UsageLimits
@@ -14,21 +16,37 @@ from .tools import AgentDeps, RunContext
 
 __all__ = (
     'ResultData',
+    'ResultValidatorFunc',
     'Usage',
     'RunResult',
     'StreamedRunResult',
 )
 
 
-ResultData = TypeVar('ResultData')
+ResultData = TypeVar('ResultData', default=str)
 """Type variable for the result data of a run."""
 
+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` and:
+
+* may or may not take [`RunContext`][pydantic_ai.tools.RunContext] as a first argument
+* may or may not be async
+
+Usage `ResultValidatorFunc[AgentDeps, ResultData]`.
+"""
+
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
 
 @dataclass
 class Usage:
-    """LLM usage associated to a request or run.
+    """LLM usage associated with a request or run.
 
     Responsibility for calculating usage is on the model; PydanticAI simply sums the usage information across requests.
 
@@ -36,7 +54,7 @@ class Usage:
     """
 
     requests: int = 0
-    """Number of requests made."""
+    """Number of requests made to the LLM API."""
     request_tokens: int | None = None
     """Tokens used in processing requests."""
     response_tokens: int | None = None
@@ -46,25 +64,33 @@ class Usage:
     details: dict[str, int] | None = None
     """Any extra details returned by the model."""
 
-    def __add__(self, other: Usage) -> Usage:
-        """Add two Usages together.
+    def incr(self, incr_usage: Usage, *, requests: int = 0) -> None:
+        """Increment the usage in place.
 
-        This is provided so it's trivial to sum usage information from multiple requests and runs.
+        Args:
+            incr_usage: The usage to increment by.
+            requests: The number of requests to increment by in addition to `incr_usage.requests`.
         """
-        counts: dict[str, int] = {}
+        self.requests += requests
         for f in 'requests', 'request_tokens', 'response_tokens', 'total_tokens':
             self_value = getattr(self, f)
-            other_value = getattr(other, f)
+            other_value = getattr(incr_usage, f)
             if self_value is not None or other_value is not None:
-                counts[f] = (self_value or 0) + (other_value or 0)
+                setattr(self, f, (self_value or 0) + (other_value or 0))
+
+        if incr_usage.details:
+            self.details = self.details or {}
+            for key, value in incr_usage.details.items():
+                self.details[key] = self.details.get(key, 0) + value
 
-        details = self.details.copy() if self.details is not None else None
-        if other.details is not None:
-            details = details or {}
-            for key, value in other.details.items():
-                details[key] = details.get(key, 0) + value
+    def __add__(self, other: Usage) -> Usage:
+        """Add two Usages together.
 
-        return Usage(**counts, details=details or None)
+        This is provided so it's trivial to sum usage information from multiple requests and runs.
+        """
+        new_usage = copy(self)
+        new_usage.incr(other)
+        return new_usage
 
 
 @dataclass
@@ -119,8 +145,6 @@ class RunResult(_BaseRunResult[ResultData]):
 class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultData]):
     """Result of a streamed run that returns structured data via a tool call."""
 
-    usage_so_far: Usage
-    """Usage of the run up until the last request."""
     _usage_limits: UsageLimits | None
     _stream_response: models.EitherStreamedResponse
     _result_schema: _result.ResultSchema[ResultData] | None
@@ -289,7 +313,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         !!! note
             This won't return the full usage until the stream is finished.
         """
-        return self.usage_so_far + self._stream_response.usage()
+        return self._run_ctx.usage + self._stream_response.usage()
 
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/settings.py

@@ -22,6 +22,7 @@ class ModelSettings(TypedDict, total=False):
     """The maximum number of tokens to generate before stopping.
 
     Supported by:
+
     * Gemini
     * Anthropic
     * OpenAI
@@ -37,6 +38,7 @@ class ModelSettings(TypedDict, total=False):
     Note that even with `temperature` of `0.0`, the results will not be fully deterministic.
 
     Supported by:
+
     * Gemini
     * Anthropic
     * OpenAI
@@ -51,6 +53,7 @@ class ModelSettings(TypedDict, total=False):
     You should either alter `temperature` or `top_p`, but not both.
 
     Supported by:
+
     * Gemini
     * Anthropic
     * OpenAI
@@ -61,6 +64,7 @@ class ModelSettings(TypedDict, total=False):
     """Override the client-level default timeout for a request, in seconds.
 
     Supported by:
+
     * Gemini
     * Anthropic
     * OpenAI
@@ -132,6 +136,6 @@ class UsageLimits:
                 f'Exceeded the response_tokens_limit of {self.response_tokens_limit} ({response_tokens=})'
             )
 
-        total_tokens = request_tokens + response_tokens
+        total_tokens = usage.total_tokens or 0
         if self.total_tokens_limit is not None and total_tokens > self.total_tokens_limit:
             raise UsageLimitExceeded(f'Exceeded the total_tokens_limit of {self.total_tokens_limit} ({total_tokens=})')

pydantic_ai/tools.py

@@ -4,7 +4,7 @@ import dataclasses
 import inspect
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
@@ -13,10 +13,12 @@ from typing_extensions import Concatenate, ParamSpec, TypeAlias
 from . import _pydantic, _utils, messages as _messages, models
 from .exceptions import ModelRetry, UnexpectedModelBehavior
 
+if TYPE_CHECKING:
+    from .result import Usage
+
 __all__ = (
     'AgentDeps',
     'RunContext',
-    'ResultValidatorFunc',
     'SystemPromptFunc',
     'ToolFuncContext',
     'ToolFuncPlain',
@@ -38,14 +40,20 @@ class RunContext(Generic[AgentDeps]):
 
     deps: AgentDeps
     """Dependencies for the agent."""
-    retry: int
-    """Number of retries so far."""
-    messages: list[_messages.ModelMessage]
-    """Messages exchanged in the conversation so far."""
-    tool_name: str | None
-    """Name of the tool being called."""
     model: models.Model
     """The model used in this run."""
+    usage: Usage
+    """LLM usage associated with the run."""
+    prompt: str
+    """The original user prompt passed to the run."""
+    messages: list[_messages.ModelMessage] = field(default_factory=list)
+    """Messages exchanged in the conversation so far."""
+    tool_name: str | None = None
+    """Name of the tool being called."""
+    retry: int = 0
+    """Number of retries so far."""
+    run_step: int = 0
+    """The current step in the run."""
 
     def replace_with(
         self, retry: int | None = None, tool_name: str | None | _utils.Unset = _utils.UNSET
@@ -73,21 +81,6 @@ SystemPromptFunc = Union[
 Usage `SystemPromptFunc[AgentDeps]`.
 """
 
-ResultData = TypeVar('ResultData')
-
-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.
 

{pydantic_ai_slim-0.0.14.dist-info → pydantic_ai_slim-0.0.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.14
+Version: 0.0.16
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -36,7 +36,6 @@ Requires-Dist: groq>=0.12.0; extra == 'groq'
 Provides-Extra: logfire
 Requires-Dist: logfire>=2.3; extra == 'logfire'
 Provides-Extra: mistral
-Requires-Dist: json-repair>=0.30.3; extra == 'mistral'
 Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
 Provides-Extra: openai
 Requires-Dist: openai>=1.54.3; extra == 'openai'

pydantic_ai_slim-0.0.16.dist-info/RECORD

@@ -0,0 +1,26 @@
+pydantic_ai/__init__.py,sha256=FbYetEgT6OO25u2KF5ZnFxKpz5DtnSpfckRXP4mjl8E,489
+pydantic_ai/_griffe.py,sha256=Wqk3AuyeWuPwE5s1GbMeCsERelx1B4QcU9uYZSoko8s,3409
+pydantic_ai/_pydantic.py,sha256=qXi5IsyiYOHeg_-qozCdxkfeqw2z0gBTjqgywBCiJWo,8125
+pydantic_ai/_result.py,sha256=cUSugZQV0n5Z4fFHiMqua-2xs_0S6m-rr-yd6QS3nFE,10317
+pydantic_ai/_system_prompt.py,sha256=MZJWksIoS5GM3Au5lznlcQnC-h7eqwtE7oI5WFgRcOg,1090
+pydantic_ai/_utils.py,sha256=skWNgm89US_x1EpxdRy5wCkghBrm1XgxFCiEh6wAkAo,8753
+pydantic_ai/agent.py,sha256=NJTcPSlqb4Fd-x9pDPuoXGCwFGF1GHcHevutoB0Busw,52333
+pydantic_ai/exceptions.py,sha256=eGDKX6bGhgVxXBzu81Sk3iiAkXr0GUtgT7bD5Rxlqpg,2028
+pydantic_ai/messages.py,sha256=ImbWY8Ft3mxInUQ08EmIWywf4nJBvTiJhmsECRYDkSQ,8968
+pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pydantic_ai/result.py,sha256=LbZVHZnJnQwgegSz5PtwS9r_ifrJnLRpsa9xjYnHg1g,15549
+pydantic_ai/settings.py,sha256=W8krcFsujjhE03rwckrz39F4Dz_9RwdBSeEF3izK0-Y,4918
+pydantic_ai/tools.py,sha256=mnh3Lvs0Ri0FkqpV1MUooExNN4epTCcBKw6DyZvNSQ8,11745
+pydantic_ai/models/__init__.py,sha256=XHt02IDQAircb-lEkIbIcuabSAIh5_UKnz2V1xN0Glw,10926
+pydantic_ai/models/anthropic.py,sha256=EUZgmvT0jhMDbooBp_jfW0z2cM5jTMuAhVws1XKgaNs,13451
+pydantic_ai/models/function.py,sha256=i7qkS_31aHrTbYVh6OzQ7Cwucz44F5PjT2EJK3GMphw,10573
+pydantic_ai/models/gemini.py,sha256=Sr19D2hN8iEAcoLlzv5883pto90TgEr_xiGlV8hMOwA,28572
+pydantic_ai/models/groq.py,sha256=ZoPkuWJrf78JPnTRfZhi7v0ETgxJKNN5dH8BLWagGGk,15770
+pydantic_ai/models/mistral.py,sha256=xGVI6-b8-9vnFickPPI2cRaHEWLc0jKKUM_vMjipf-U,25894
+pydantic_ai/models/ollama.py,sha256=ELqxhcNcnvQBnadd3gukS01zprUp6v8N_h1P5K-uf6c,4188
+pydantic_ai/models/openai.py,sha256=qFFInL3NbgfGcsAWigxMP5mscp76hC-jJimHc9woU6Y,16518
+pydantic_ai/models/test.py,sha256=u2pdZd9OLXQ_jI6CaVt96udXuIcv0Hfnfqd3pFGmeJM,16514
+pydantic_ai/models/vertexai.py,sha256=DBCBfpvpIhZaMG7cKvRl5rugCZqJqqEFm74uBc45weo,9259
+pydantic_ai_slim-0.0.16.dist-info/METADATA,sha256=4udd7j2erIuMC0ekYgmgQAqsKfhA5sLsKzTcD_QyOeo,2730
+pydantic_ai_slim-0.0.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_ai_slim-0.0.16.dist-info/RECORD,,

pydantic_ai_slim-0.0.14.dist-info/RECORD

@@ -1,26 +0,0 @@
-pydantic_ai/__init__.py,sha256=a3ffrVF4eJyylreRa6LbTMF6RI7VIv96zHv2tCWo1kQ,439
-pydantic_ai/_griffe.py,sha256=pRjCJ6B1hhx6k46XJgl9zF6aRYxRmqEZKFok8unp4Iw,3449
-pydantic_ai/_pydantic.py,sha256=qXi5IsyiYOHeg_-qozCdxkfeqw2z0gBTjqgywBCiJWo,8125
-pydantic_ai/_result.py,sha256=iL0oZXvuCEoa37EKGXkEhn90oB_950emKG-uXdlmssM,10317
-pydantic_ai/_system_prompt.py,sha256=MZJWksIoS5GM3Au5lznlcQnC-h7eqwtE7oI5WFgRcOg,1090
-pydantic_ai/_utils.py,sha256=skWNgm89US_x1EpxdRy5wCkghBrm1XgxFCiEh6wAkAo,8753
-pydantic_ai/agent.py,sha256=6iOVgmoGw5wGd0W4ewFgOGJqTLmtJLyl5RLqt3dSegE,49574
-pydantic_ai/exceptions.py,sha256=eGDKX6bGhgVxXBzu81Sk3iiAkXr0GUtgT7bD5Rxlqpg,2028
-pydantic_ai/messages.py,sha256=ImbWY8Ft3mxInUQ08EmIWywf4nJBvTiJhmsECRYDkSQ,8968
-pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=pE8YTCUDS0hXTplEJKSG3UnhMvwGNqaOU-9eC09K4y4,14735
-pydantic_ai/settings.py,sha256=rSNqWYzBoIwC9wZCS2M_aEdCNAy5sOXbi5bXJAuqK08,4923
-pydantic_ai/tools.py,sha256=9ZEhvgylv3pc0_JaahCJgJlFxubCWeEaCNqERl-I9B0,11982
-pydantic_ai/models/__init__.py,sha256=XHt02IDQAircb-lEkIbIcuabSAIh5_UKnz2V1xN0Glw,10926
-pydantic_ai/models/anthropic.py,sha256=EUZgmvT0jhMDbooBp_jfW0z2cM5jTMuAhVws1XKgaNs,13451
-pydantic_ai/models/function.py,sha256=i7qkS_31aHrTbYVh6OzQ7Cwucz44F5PjT2EJK3GMphw,10573
-pydantic_ai/models/gemini.py,sha256=8vdcW4izL9NUGFj6lcD9yIPaakCtsmHauTvKwlTzD14,28207
-pydantic_ai/models/groq.py,sha256=ZoPkuWJrf78JPnTRfZhi7v0ETgxJKNN5dH8BLWagGGk,15770
-pydantic_ai/models/mistral.py,sha256=xZMK2vNLDR4uw1XoAQs-3obeA6c39Q1Qhei8w1JMIow,26458
-pydantic_ai/models/ollama.py,sha256=i3mMXkXu9xL6f4c52Eyx3j4aHKfYoloFondlGHPtkS4,3971
-pydantic_ai/models/openai.py,sha256=qFFInL3NbgfGcsAWigxMP5mscp76hC-jJimHc9woU6Y,16518
-pydantic_ai/models/test.py,sha256=pty5qaudHsSDvdE89HqMj-kmd4UMV9VJI2YGtdfOX1o,15960
-pydantic_ai/models/vertexai.py,sha256=DBCBfpvpIhZaMG7cKvRl5rugCZqJqqEFm74uBc45weo,9259
-pydantic_ai_slim-0.0.14.dist-info/METADATA,sha256=GgTCeoTiFNW1S0pSniLn7ILqGrLK8eaOkOG84aYE06M,2785
-pydantic_ai_slim-0.0.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_ai_slim-0.0.14.dist-info/RECORD,,