pydantic-ai-slim 0.0.15__tar.gz → 0.0.17__tar.gz

{pydantic_ai_slim-0.0.15 → pydantic_ai_slim-0.0.17}/PKG-INFO

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

{pydantic_ai_slim-0.0.15 → pydantic_ai_slim-0.0.17}/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_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pydantic_ai/agent.py

@@ -6,7 +6,6 @@ import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
-from dataclasses import dataclass, field
 from types import FrameType
 from typing import Any, Callable, Generic, Literal, cast, final, overload
 
@@ -21,9 +20,10 @@ from . import (
     messages as _messages,
     models,
     result,
+    usage as _usage,
 )
 from .result import ResultData
-from .settings import ModelSettings, UsageLimits, merge_model_settings
+from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDeps,
     RunContext,
@@ -40,6 +40,16 @@ __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.
@@ -50,7 +60,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.
 
@@ -90,17 +100,17 @@ class Agent(Generic[AgentDeps, ResultData]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    _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,
@@ -183,7 +193,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
-        usage_limits: UsageLimits | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt in async mode.
@@ -206,6 +217,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:
@@ -213,7 +225,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
@@ -222,40 +234,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 _usage.Usage(), user_prompt)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
             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()
+            usage_limits = usage_limits or _usage.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:
@@ -265,11 +273,14 @@ class Agent(Generic[AgentDeps, ResultData]):
                     # Check if we got a final result
                     if final_result is not None:
                         result_data = final_result.data
+                        result_tool_name = final_result.tool_name
                         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, result_tool_name, run_context.usage
+                        )
                     else:
                         # continue the conversation
                         handle_span.set_attribute('tool_responses', tool_responses)
@@ -284,7 +295,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
-        usage_limits: UsageLimits | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[ResultData]:
         """Run the agent with a user prompt synchronously.
@@ -311,6 +323,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:
@@ -326,6 +339,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 deps=deps,
                 model_settings=model_settings,
                 usage_limits=usage_limits,
+                usage=usage,
                 infer_name=False,
             )
         )
@@ -339,7 +353,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
         model_settings: ModelSettings | None = None,
-        usage_limits: UsageLimits | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.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.
@@ -363,6 +378,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:
@@ -372,7 +388,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
@@ -381,32 +397,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 _usage.Usage(), user_prompt)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
             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()
+            usage_limits = usage_limits or _usage.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
@@ -442,7 +455,6 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 yield result.StreamedRunResult(
                                     messages,
                                     new_message_index,
-                                    usage,
                                     usage_limits,
                                     result_stream,
                                     self._result_schema,
@@ -466,8 +478,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(
@@ -778,14 +790,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:
@@ -796,18 +808,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."""
@@ -830,15 +839,15 @@ class Agent(Generic[AgentDeps, ResultData]):
         self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[AgentDeps]
     ) -> list[_messages.ModelMessage]:
         try:
-            messages = _messages_ctx_var.get()
+            ctx_messages = _messages_ctx_var.get()
         except LookupError:
-            messages = []
+            messages: list[_messages.ModelMessage] = []
         else:
-            if messages:
-                raise exceptions.UserError(
-                    'The capture_run_messages() context manager may only be used to wrap '
-                    'one call to run(), run_sync(), or run_stream().'
-                )
+            if ctx_messages.used:
+                messages = []
+            else:
+                messages = ctx_messages.messages
+                ctx_messages.used = True
 
         if message_history:
             # shallow copy messages
@@ -1132,7 +1141,13 @@ class Agent(Generic[AgentDeps, ResultData]):
         raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
 
 
-_messages_ctx_var: ContextVar[list[_messages.ModelMessage]] = ContextVar('var')
+@dataclasses.dataclass
+class _RunMessages:
+    messages: list[_messages.ModelMessage]
+    used: bool = False
+
+
+_messages_ctx_var: ContextVar[_RunMessages] = ContextVar('var')
 
 
 @contextmanager
@@ -1156,21 +1171,21 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
     ```
 
     !!! note
-        You may not call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context.
-        If you try to do so, a [`UserError`][pydantic_ai.exceptions.UserError] will be raised.
+        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()
+        yield _messages_ctx_var.get().messages
     except LookupError:
         messages: list[_messages.ModelMessage] = []
-        token = _messages_ctx_var.set(messages)
+        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_slim-0.0.17/pydantic_ai/format_as_xml.py

@@ -0,0 +1,115 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import asdict, dataclass, is_dataclass
+from datetime import date
+from typing import Any
+from xml.etree import ElementTree
+
+from pydantic import BaseModel
+
+__all__ = ('format_as_xml',)
+
+
+def format_as_xml(
+    obj: Any,
+    root_tag: str = 'examples',
+    item_tag: str = 'example',
+    include_root_tag: bool = True,
+    none_str: str = 'null',
+    indent: str | None = '  ',
+) -> str:
+    """Format a Python object as XML.
+
+    This is useful since LLMs often find it easier to read semi-structured data (e.g. examples) as XML,
+    rather than JSON etc.
+
+    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Mapping`,
+    `Iterable`, `dataclass`, and `BaseModel`.
+
+    Args:
+        obj: Python Object to serialize to XML.
+        root_tag: Outer tag to wrap the XML in, use `None` to omit the outer tag.
+        item_tag: Tag to use for each item in an iterable (e.g. list), this is overridden by the class name
+            for dataclasses and Pydantic models.
+        include_root_tag: Whether to include the root tag in the output
+            (The root tag is always included if it includes a body - e.g. when the input is a simple value).
+        none_str: String to use for `None` values.
+        indent: Indentation string to use for pretty printing.
+
+    Returns: XML representation of the object.
+
+    Example:
+    ```python {title="format_as_xml_example.py" lint="skip"}
+    from pydantic_ai.format_as_xml import format_as_xml
+
+    print(format_as_xml({'name': 'John', 'height': 6, 'weight': 200}, root_tag='user'))
+    '''
+    <user>
+      <name>John</name>
+      <height>6</height>
+      <weight>200</weight>
+    </user>
+    '''
+    ```
+    """
+    el = _ToXml(item_tag=item_tag, none_str=none_str).to_xml(obj, root_tag)
+    if not include_root_tag and el.text is None:
+        join = '' if indent is None else '\n'
+        return join.join(_rootless_xml_elements(el, indent))
+    else:
+        if indent is not None:
+            ElementTree.indent(el, space=indent)
+        return ElementTree.tostring(el, encoding='unicode')
+
+
+@dataclass
+class _ToXml:
+    item_tag: str
+    none_str: str
+
+    def to_xml(self, value: Any, tag: str | None) -> ElementTree.Element:
+        element = ElementTree.Element(self.item_tag if tag is None else tag)
+        if value is None:
+            element.text = self.none_str
+        elif isinstance(value, str):
+            element.text = value
+        elif isinstance(value, (bytes, bytearray)):
+            element.text = value.decode(errors='ignore')
+        elif isinstance(value, (bool, int, float)):
+            element.text = str(value)
+        elif isinstance(value, date):
+            element.text = value.isoformat()
+        elif isinstance(value, Mapping):
+            self._mapping_to_xml(element, value)  # pyright: ignore[reportUnknownArgumentType]
+        elif is_dataclass(value) and not isinstance(value, type):
+            if tag is None:
+                element = ElementTree.Element(value.__class__.__name__)
+            dc_dict = asdict(value)
+            self._mapping_to_xml(element, dc_dict)
+        elif isinstance(value, BaseModel):
+            if tag is None:
+                element = ElementTree.Element(value.__class__.__name__)
+            self._mapping_to_xml(element, value.model_dump(mode='python'))
+        elif isinstance(value, Iterable):
+            for item in value:  # pyright: ignore[reportUnknownVariableType]
+                item_el = self.to_xml(item, None)
+                element.append(item_el)
+        else:
+            raise TypeError(f'Unsupported type for XML formatting: {type(value)}')
+        return element
+
+    def _mapping_to_xml(self, element: ElementTree.Element, mapping: Mapping[Any, Any]) -> None:
+        for key, value in mapping.items():
+            if isinstance(key, int):
+                key = str(key)
+            elif not isinstance(key, str):
+                raise TypeError(f'Unsupported key type for XML formatting: {type(key)}, only str and int are allowed')
+            element.append(self.to_xml(value, key))
+
+
+def _rootless_xml_elements(root: ElementTree.Element, indent: str | None) -> Iterator[str]:
+    for sub_element in root:
+        if indent is not None:
+            ElementTree.indent(sub_element, space=indent)
+        yield ElementTree.tostring(sub_element, encoding='unicode')

{pydantic_ai_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pydantic_ai/models/gemini.py

@@ -273,17 +273,26 @@ class GeminiAgentModel(AgentModel):
         contents: list[_GeminiContent] = []
         for m in messages:
             if isinstance(m, ModelRequest):
+                message_parts: list[_GeminiPartUnion] = []
+
                 for part in m.parts:
                     if isinstance(part, SystemPromptPart):
                         sys_prompt_parts.append(_GeminiTextPart(text=part.content))
                     elif isinstance(part, UserPromptPart):
-                        contents.append(_content_user_prompt(part))
+                        message_parts.append(_GeminiTextPart(text=part.content))
                     elif isinstance(part, ToolReturnPart):
-                        contents.append(_content_tool_return(part))
+                        message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object()))
                     elif isinstance(part, RetryPromptPart):
-                        contents.append(_content_retry_prompt(part))
+                        if part.tool_name is None:
+                            message_parts.append(_GeminiTextPart(text=part.model_response()))
+                        else:
+                            response = {'call_error': part.model_response()}
+                            message_parts.append(_response_part_from_response(part.tool_name, response))
                     else:
                         assert_never(part)
+
+                if message_parts:
+                    contents.append(_GeminiContent(role='user', parts=message_parts))
             elif isinstance(m, ModelResponse):
                 contents.append(_content_model_response(m))
             else:
@@ -420,31 +429,14 @@ class _GeminiContent(TypedDict):
     parts: list[_GeminiPartUnion]
 
 
-def _content_user_prompt(m: UserPromptPart) -> _GeminiContent:
-    return _GeminiContent(role='user', parts=[_GeminiTextPart(text=m.content)])
-
-
-def _content_tool_return(m: ToolReturnPart) -> _GeminiContent:
-    f_response = _response_part_from_response(m.tool_name, m.model_response_object())
-    return _GeminiContent(role='user', parts=[f_response])
-
-
-def _content_retry_prompt(m: RetryPromptPart) -> _GeminiContent:
-    if m.tool_name is None:
-        part = _GeminiTextPart(text=m.model_response())
-    else:
-        response = {'call_error': m.model_response()}
-        part = _response_part_from_response(m.tool_name, response)
-    return _GeminiContent(role='user', parts=[part])
-
-
 def _content_model_response(m: ModelResponse) -> _GeminiContent:
     parts: list[_GeminiPartUnion] = []
     for item in m.parts:
         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 +693,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 +706,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_slim-0.0.15 → pydantic_ai_slim-0.0.17}/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_slim-0.0.15 → pydantic_ai_slim-0.0.17}/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_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pydantic_ai/models/vertexai.py

@@ -178,7 +178,7 @@ def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredenti
 # pyright: reportUnknownVariableType=false
 # pyright: reportUnknownArgumentType=false
 async def _async_google_auth() -> tuple[BaseCredentials, str | None]:
-    return await run_in_executor(google.auth.default)
+    return await run_in_executor(google.auth.default, scopes=['https://www.googleapis.com/auth/cloud-platform'])
 
 
 # default expiry is 3600 seconds

{pydantic_ai_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pydantic_ai/result.py

@@ -2,6 +2,7 @@ from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Awaitable, Callable
+from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic, Union, cast
@@ -10,16 +11,10 @@ import logfire_api
 from typing_extensions import TypeVar
 
 from . import _result, _utils, exceptions, messages as _messages, models
-from .settings import UsageLimits
 from .tools import AgentDeps, RunContext
+from .usage import Usage, UsageLimits
 
-__all__ = (
-    'ResultData',
-    'ResultValidatorFunc',
-    'Usage',
-    'RunResult',
-    'StreamedRunResult',
-)
+__all__ = 'ResultData', 'ResultValidatorFunc', 'RunResult', 'StreamedRunResult'
 
 
 ResultData = TypeVar('ResultData', default=str)
@@ -43,47 +38,6 @@ Usage `ResultValidatorFunc[AgentDeps, ResultData]`.
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
 
-@dataclass
-class Usage:
-    """LLM usage associated with a request or run.
-
-    Responsibility for calculating usage is on the model; PydanticAI simply sums the usage information across requests.
-
-    You'll need to look up the documentation of the model you're using to convert usage to monetary costs.
-    """
-
-    requests: int = 0
-    """Number of requests made to the LLM API."""
-    request_tokens: int | None = None
-    """Tokens used in processing requests."""
-    response_tokens: int | None = None
-    """Tokens used in generating responses."""
-    total_tokens: int | None = None
-    """Total tokens used in the whole run, should generally be equal to `request_tokens + response_tokens`."""
-    details: dict[str, int] | None = None
-    """Any extra details returned by the model."""
-
-    def __add__(self, other: Usage) -> Usage:
-        """Add two Usages together.
-
-        This is provided so it's trivial to sum usage information from multiple requests and runs.
-        """
-        counts: dict[str, int] = {}
-        for f in 'requests', 'request_tokens', 'response_tokens', 'total_tokens':
-            self_value = getattr(self, f)
-            other_value = getattr(other, f)
-            if self_value is not None or other_value is not None:
-                counts[f] = (self_value or 0) + (other_value or 0)
-
-        details = self.details.copy() if self.details is not None else None
-        if other.details is not None:
-            details = details or {}
-            for key, value in other.details.items():
-                details[key] = details.get(key, 0) + value
-
-        return Usage(**counts, details=details or None)
-
-
 @dataclass
 class _BaseRunResult(ABC, Generic[ResultData]):
     """Base type for results.
@@ -94,25 +48,70 @@ class _BaseRunResult(ABC, Generic[ResultData]):
     _all_messages: list[_messages.ModelMessage]
     _new_message_index: int
 
-    def all_messages(self) -> list[_messages.ModelMessage]:
-        """Return the history of _messages."""
+    def all_messages(self, *, result_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
+        """Return the history of _messages.
+
+        Args:
+            result_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the result tool call if you want to continue
+                the conversation and want to set the response to the result tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            List of messages.
+        """
         # this is a method to be consistent with the other methods
+        if result_tool_return_content is not None:
+            raise NotImplementedError('Setting result tool return content is not supported for this result type.')
         return self._all_messages
 
-    def all_messages_json(self) -> bytes:
-        """Return all messages from [`all_messages`][pydantic_ai.result._BaseRunResult.all_messages] as JSON bytes."""
-        return _messages.ModelMessagesTypeAdapter.dump_json(self.all_messages())
+    def all_messages_json(self, *, result_tool_return_content: str | None = None) -> bytes:
+        """Return all messages from [`all_messages`][pydantic_ai.result._BaseRunResult.all_messages] as JSON bytes.
+
+        Args:
+            result_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the result tool call if you want to continue
+                the conversation and want to set the response to the result tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            JSON bytes representing the messages.
+        """
+        return _messages.ModelMessagesTypeAdapter.dump_json(
+            self.all_messages(result_tool_return_content=result_tool_return_content)
+        )
 
-    def new_messages(self) -> list[_messages.ModelMessage]:
+    def new_messages(self, *, result_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
         """Return new messages associated with this run.
 
-        System prompts and any messages from older runs are excluded.
+        Messages from older runs are excluded.
+
+        Args:
+            result_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the result tool call if you want to continue
+                the conversation and want to set the response to the result tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            List of new messages.
         """
-        return self.all_messages()[self._new_message_index :]
+        return self.all_messages(result_tool_return_content=result_tool_return_content)[self._new_message_index :]
+
+    def new_messages_json(self, *, result_tool_return_content: str | None = None) -> bytes:
+        """Return new messages from [`new_messages`][pydantic_ai.result._BaseRunResult.new_messages] as JSON bytes.
 
-    def new_messages_json(self) -> bytes:
-        """Return new messages from [`new_messages`][pydantic_ai.result._BaseRunResult.new_messages] as JSON bytes."""
-        return _messages.ModelMessagesTypeAdapter.dump_json(self.new_messages())
+        Args:
+            result_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the result tool call if you want to continue
+                the conversation and want to set the response to the result tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            JSON bytes representing the new messages.
+        """
+        return _messages.ModelMessagesTypeAdapter.dump_json(
+            self.new_messages(result_tool_return_content=result_tool_return_content)
+        )
 
     @abstractmethod
     def usage(self) -> Usage:
@@ -125,19 +124,50 @@ class RunResult(_BaseRunResult[ResultData]):
 
     data: ResultData
     """Data from the final response in the run."""
+    _result_tool_name: str | None
     _usage: Usage
 
     def usage(self) -> Usage:
         """Return the usage of the whole run."""
         return self._usage
 
+    def all_messages(self, *, result_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
+        """Return the history of _messages.
+
+        Args:
+            result_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the result tool call if you want to continue
+                the conversation and want to set the response to the result tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            List of messages.
+        """
+        if result_tool_return_content is not None:
+            return self._set_result_tool_return(result_tool_return_content)
+        else:
+            return self._all_messages
+
+    def _set_result_tool_return(self, return_content: str) -> list[_messages.ModelMessage]:
+        """Set return content for the result tool.
+
+        Useful if you want to continue the conversation and want to set the response to the result tool call.
+        """
+        if not self._result_tool_name:
+            raise ValueError('Cannot set result tool return content when the return type is `str`.')
+        messages = deepcopy(self._all_messages)
+        last_message = messages[-1]
+        for part in last_message.parts:
+            if isinstance(part, _messages.ToolReturnPart) and part.tool_name == self._result_tool_name:
+                part.content = return_content
+                return messages
+        raise LookupError(f'No tool call found with tool name {self._result_tool_name!r}.')
+
 
 @dataclass
 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
@@ -306,7 +336,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_slim-0.0.17/pydantic_ai/settings.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from httpx import Timeout
+from typing_extensions import TypedDict
+
+if TYPE_CHECKING:
+    pass
+
+
+class ModelSettings(TypedDict, total=False):
+    """Settings to configure an LLM.
+
+    Here we include only settings which apply to multiple models / model providers.
+    """
+
+    max_tokens: int
+    """The maximum number of tokens to generate before stopping.
+
+    Supported by:
+
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    temperature: float
+    """Amount of randomness injected into the response.
+
+    Use `temperature` closer to `0.0` for analytical / multiple choice, and closer to a model's
+    maximum `temperature` for creative and generative tasks.
+
+    Note that even with `temperature` of `0.0`, the results will not be fully deterministic.
+
+    Supported by:
+
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    top_p: float
+    """An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass.
+
+    So 0.1 means only the tokens comprising the top 10% probability mass are considered.
+
+    You should either alter `temperature` or `top_p`, but not both.
+
+    Supported by:
+
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    timeout: float | Timeout
+    """Override the client-level default timeout for a request, in seconds.
+
+    Supported by:
+
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+
+def merge_model_settings(base: ModelSettings | None, overrides: ModelSettings | None) -> ModelSettings | None:
+    """Merge two sets of model settings, preferring the overrides.
+
+    A common use case is: merge_model_settings(<agent settings>, <run settings>)
+    """
+    # Note: we may want merge recursively if/when we add non-primitive values
+    if base and overrides:
+        return base | overrides
+    else:
+        return base or overrides

{pydantic_ai_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pydantic_ai/tools.py

@@ -4,15 +4,18 @@ 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, Union, cast
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
-from typing_extensions import Concatenate, ParamSpec, TypeAlias
+from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypeVar
 
 from . import _pydantic, _utils, messages as _messages, models
 from .exceptions import ModelRetry, UnexpectedModelBehavior
 
+if TYPE_CHECKING:
+    from .result import Usage
+
 __all__ = (
     'AgentDeps',
     'RunContext',
@@ -27,7 +30,7 @@ __all__ = (
     'ToolDefinition',
 )
 
-AgentDeps = TypeVar('AgentDeps')
+AgentDeps = TypeVar('AgentDeps', default=None)
 """Type variable for agent dependencies."""
 
 
@@ -37,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
@@ -58,7 +67,7 @@ class RunContext(Generic[AgentDeps]):
         return dataclasses.replace(self, **kwargs)
 
 
-ToolParams = ParamSpec('ToolParams')
+ToolParams = ParamSpec('ToolParams', default=...)
 """Retrieval function param spec."""
 
 SystemPromptFunc = Union[
@@ -83,7 +92,7 @@ ToolFuncPlain = Callable[ToolParams, Any]
 Usage `ToolPlainFunc[ToolParams]`.
 """
 ToolFuncEither = Union[ToolFuncContext[AgentDeps, ToolParams], ToolFuncPlain[ToolParams]]
-"""Either part_kind of tool function.
+"""Either kind of tool function.
 
 This is just a union of [`ToolFuncContext`][pydantic_ai.tools.ToolFuncContext] and
 [`ToolFuncPlain`][pydantic_ai.tools.ToolFuncPlain].
@@ -125,7 +134,7 @@ A = TypeVar('A')
 class Tool(Generic[AgentDeps]):
     """A tool function for an agent."""
 
-    function: ToolFuncEither[AgentDeps, ...]
+    function: ToolFuncEither[AgentDeps]
     takes_ctx: bool
     max_retries: int | None
     name: str
@@ -141,7 +150,7 @@ class Tool(Generic[AgentDeps]):
 
     def __init__(
         self,
-        function: ToolFuncEither[AgentDeps, ...],
+        function: ToolFuncEither[AgentDeps],
         *,
         takes_ctx: bool | None = None,
         max_retries: int | None = None,

pydantic_ai_slim-0.0.15/pydantic_ai/settings.py → pydantic_ai_slim-0.0.17/pydantic_ai/usage.py

@@ -1,87 +1,60 @@
-from __future__ import annotations
+from __future__ import annotations as _annotations
 
+from copy import copy
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-from httpx import Timeout
-from typing_extensions import TypedDict
 
 from .exceptions import UsageLimitExceeded
 
-if TYPE_CHECKING:
-    from .result import Usage
-
-
-class ModelSettings(TypedDict, total=False):
-    """Settings to configure an LLM.
-
-    Here we include only settings which apply to multiple models / model providers.
-    """
-
-    max_tokens: int
-    """The maximum number of tokens to generate before stopping.
-
-    Supported by:
-
-    * Gemini
-    * Anthropic
-    * OpenAI
-    * Groq
-    """
-
-    temperature: float
-    """Amount of randomness injected into the response.
-
-    Use `temperature` closer to `0.0` for analytical / multiple choice, and closer to a model's
-    maximum `temperature` for creative and generative tasks.
-
-    Note that even with `temperature` of `0.0`, the results will not be fully deterministic.
-
-    Supported by:
+__all__ = 'Usage', 'UsageLimits'
 
-    * Gemini
-    * Anthropic
-    * OpenAI
-    * Groq
-    """
-
-    top_p: float
-    """An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass.
-
-    So 0.1 means only the tokens comprising the top 10% probability mass are considered.
-
-    You should either alter `temperature` or `top_p`, but not both.
-
-    Supported by:
-
-    * Gemini
-    * Anthropic
-    * OpenAI
-    * Groq
-    """
 
-    timeout: float | Timeout
-    """Override the client-level default timeout for a request, in seconds.
+@dataclass
+class Usage:
+    """LLM usage associated with a request or run.
 
-    Supported by:
+    Responsibility for calculating usage is on the model; PydanticAI simply sums the usage information across requests.
 
-    * Gemini
-    * Anthropic
-    * OpenAI
-    * Groq
+    You'll need to look up the documentation of the model you're using to convert usage to monetary costs.
     """
 
-
-def merge_model_settings(base: ModelSettings | None, overrides: ModelSettings | None) -> ModelSettings | None:
-    """Merge two sets of model settings, preferring the overrides.
-
-    A common use case is: merge_model_settings(<agent settings>, <run settings>)
-    """
-    # Note: we may want merge recursively if/when we add non-primitive values
-    if base and overrides:
-        return base | overrides
-    else:
-        return base or overrides
+    requests: int = 0
+    """Number of requests made to the LLM API."""
+    request_tokens: int | None = None
+    """Tokens used in processing requests."""
+    response_tokens: int | None = None
+    """Tokens used in generating responses."""
+    total_tokens: int | None = None
+    """Total tokens used in the whole run, should generally be equal to `request_tokens + response_tokens`."""
+    details: dict[str, int] | None = None
+    """Any extra details returned by the model."""
+
+    def incr(self, incr_usage: Usage, *, requests: int = 0) -> None:
+        """Increment the usage in place.
+
+        Args:
+            incr_usage: The usage to increment by.
+            requests: The number of requests to increment by in addition to `incr_usage.requests`.
+        """
+        self.requests += requests
+        for f in 'requests', 'request_tokens', 'response_tokens', 'total_tokens':
+            self_value = getattr(self, f)
+            other_value = getattr(incr_usage, f)
+            if self_value is not None or other_value is not None:
+                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
+
+    def __add__(self, other: Usage) -> Usage:
+        """Add two Usages together.
+
+        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
@@ -136,6 +109,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_slim-0.0.15 → pydantic_ai_slim-0.0.17}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.15"
+version = "0.0.17"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [
     { name = "Samuel Colvin", email = "samuel@pydantic.dev" },