pydantic-ai-slim 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

pydantic_ai/models/test.py

@@ -2,27 +2,29 @@ from __future__ import annotations as _annotations
 
 import re
 import string
-from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable, Iterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import date, datetime, timedelta
 from typing import Any, Literal
 
 import pydantic_core
+from typing_extensions import assert_never
 
 from .. import _utils
 from ..messages import (
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    RetryPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
 )
 from ..result import Cost
+from ..settings import ModelSettings
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -55,25 +57,38 @@ class TestModel(Model):
     """If set, these args will be passed to the result tool."""
     seed: int = 0
     """Seed for generating random data."""
-    # these fields are set when the model is called by the agent
-    agent_model_tools: Mapping[str, AbstractToolDefinition] | None = field(default=None, init=False)
+    agent_model_function_tools: list[ToolDefinition] | None = field(default=None, init=False)
+    """Definition of function tools passed to the model.
+
+    This is set when the model is called, so will reflect the function tools from the last step of the last run.
+    """
     agent_model_allow_text_result: bool | None = field(default=None, init=False)
-    agent_model_result_tools: list[AbstractToolDefinition] | None = field(default=None, init=False)
+    """Whether plain text responses from the model are allowed.
+
+    This is set when the model is called, so will reflect the value from the last step of the last run.
+    """
+    agent_model_result_tools: list[ToolDefinition] | None = field(default=None, init=False)
+    """Definition of result tools passed to the model.
+
+    This is set when the model is called, so will reflect the result tools from the last step of the last run.
+    """
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
-        self.agent_model_tools = function_tools
+        self.agent_model_function_tools = function_tools
         self.agent_model_allow_text_result = allow_text_result
-        self.agent_model_result_tools = list(result_tools) if result_tools is not None else None
+        self.agent_model_result_tools = result_tools
 
         if self.call_tools == 'all':
-            tool_calls = [(r.name, r) for r in function_tools.values()]
+            tool_calls = [(r.name, r) for r in function_tools]
         else:
-            tools_to_call = (function_tools[name] for name in self.call_tools)
+            function_tools_lookup = {t.name: t for t in function_tools}
+            tools_to_call = (function_tools_lookup[name] for name in self.call_tools)
             tool_calls = [(r.name, r) for r in tools_to_call]
 
         if self.custom_result_text is not None:
@@ -90,11 +105,12 @@ class TestModel(Model):
                 result = _utils.Either(right=self.custom_result_args)
         elif allow_text_result:
             result = _utils.Either(left=None)
-        elif result_tools is not None:
+        elif result_tools:
             result = _utils.Either(right=None)
         else:
             result = _utils.Either(left=None)
-        return TestAgentModel(tool_calls, result, self.agent_model_result_tools, self.seed)
+
+        return TestAgentModel(tool_calls, result, result_tools, self.seed)
 
     def name(self) -> str:
         return 'test-model'
@@ -107,73 +123,89 @@ class TestAgentModel(AgentModel):
     # NOTE: Avoid test discovery by pytest.
     __test__ = False
 
-    tool_calls: list[tuple[str, AbstractToolDefinition]]
+    tool_calls: list[tuple[str, ToolDefinition]]
     # left means the text is plain text; right means it's a function call
     result: _utils.Either[str | None, Any | None]
-    result_tools: list[AbstractToolDefinition] | None
+    result_tools: list[ToolDefinition]
     seed: int
-    step: int = 0
-    last_message_count: int = 0
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
-        return self._request(messages), Cost()
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, Cost]:
+        return self._request(messages, model_settings), Cost()
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        msg = self._request(messages)
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        msg = self._request(messages, model_settings)
         cost = Cost()
-        if isinstance(msg, ModelTextResponse):
-            yield TestStreamTextResponse(msg.content, cost)
+
+        # TODO: Rework this once we make StreamTextResponse more general
+        texts: list[str] = []
+        tool_calls: list[ToolCallPart] = []
+        for item in msg.parts:
+            if isinstance(item, TextPart):
+                texts.append(item.content)
+            elif isinstance(item, ToolCallPart):
+                tool_calls.append(item)
+            else:
+                assert_never(item)
+
+        if texts:
+            yield TestStreamTextResponse('\n\n'.join(texts), cost)
         else:
             yield TestStreamStructuredResponse(msg, cost)
 
-    def gen_tool_args(self, tool_def: AbstractToolDefinition) -> Any:
-        return _JsonSchemaTestData(tool_def.json_schema, self.seed).generate()
-
-    def _request(self, messages: list[Message]) -> ModelAnyResponse:
-        if self.step == 0 and self.tool_calls:
-            calls = [ToolCall.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
-            self.step += 1
-            self.last_message_count = len(messages)
-            return ModelStructuredResponse(calls=calls)
-
-        new_messages = messages[self.last_message_count :]
-        self.last_message_count = len(messages)
-        new_retry_names = {m.tool_name for m in new_messages if isinstance(m, RetryPrompt)}
-        if new_retry_names:
-            calls = [
-                ToolCall.from_dict(name, self.gen_tool_args(args))
-                for name, args in self.tool_calls
-                if name in new_retry_names
-            ]
-            self.step += 1
-            return ModelStructuredResponse(calls=calls)
-        else:
-            if response_text := self.result.left:
-                self.step += 1
-                if response_text.value is None:
-                    # build up details of tool responses
-                    output: dict[str, Any] = {}
-                    for message in messages:
-                        if isinstance(message, ToolReturn):
-                            output[message.tool_name] = message.content
-                    if output:
-                        return ModelTextResponse(content=pydantic_core.to_json(output).decode())
-                    else:
-                        return ModelTextResponse(content='success (no tool calls)')
+    def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
+        return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
+
+    def _request(self, messages: list[ModelMessage], model_settings: ModelSettings | None) -> ModelResponse:
+        # if there are tools, the first thing we want to do is call all of them
+        if self.tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
+            return ModelResponse(
+                parts=[ToolCallPart.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+            )
+
+        if messages:
+            last_message = messages[-1]
+            assert isinstance(last_message, ModelRequest), 'Expected last message to be a `ModelRequest`.'
+
+            # 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_dict(name, self.gen_tool_args(args))
+                        for name, args in self.tool_calls
+                        if name in new_retry_names
+                    ]
+                )
+
+        if response_text := self.result.left:
+            if response_text.value is None:
+                # build up details of tool responses
+                output: dict[str, Any] = {}
+                for message in messages:
+                    if isinstance(message, ModelRequest):
+                        for part in message.parts:
+                            if isinstance(part, ToolReturnPart):
+                                output[part.tool_name] = part.content
+                if output:
+                    return ModelResponse.from_text(pydantic_core.to_json(output).decode())
                 else:
-                    return ModelTextResponse(content=response_text.value)
+                    return ModelResponse.from_text('success (no tool calls)')
             else:
-                assert self.result_tools is not None, 'No result tools provided'
-                custom_result_args = self.result.right
-                result_tool = self.result_tools[self.seed % len(self.result_tools)]
-                if custom_result_args is not None:
-                    self.step += 1
-                    return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, custom_result_args)])
-                else:
-                    response_args = self.gen_tool_args(result_tool)
-                    self.step += 1
-                    return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, response_args)])
+                return ModelResponse.from_text(response_text.value)
+        else:
+            assert self.result_tools, 'No result tools provided'
+            custom_result_args = self.result.right
+            result_tool = self.result_tools[self.seed % len(self.result_tools)]
+            if custom_result_args is not None:
+                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, custom_result_args)])
+            else:
+                response_args = self.gen_tool_args(result_tool)
+                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, response_args)])
 
 
 @dataclass
@@ -213,7 +245,7 @@ class TestStreamTextResponse(StreamTextResponse):
 class TestStreamStructuredResponse(StreamStructuredResponse):
     """A structured response that streams test data."""
 
-    _structured_response: ModelStructuredResponse
+    _structured_response: ModelResponse
     _cost: Cost
     _iter: Iterator[None] = field(default_factory=lambda: iter([None]))
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
@@ -221,7 +253,7 @@ class TestStreamStructuredResponse(StreamStructuredResponse):
     async def __anext__(self) -> None:
         return _utils.sync_anext(self._iter)
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
+    def get(self, *, final: bool = False) -> ModelResponse:
         return self._structured_response
 
     def cost(self) -> Cost:

pydantic_ai/models/vertexai.py

@@ -1,6 +1,5 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -10,7 +9,8 @@ from httpx import AsyncClient as AsyncHTTPClient
 
 from .._utils import run_in_executor
 from ..exceptions import UserError
-from . import AbstractToolDefinition, Model, cached_async_http_client
+from ..tools import ToolDefinition
+from . import Model, cached_async_http_client
 from .gemini import GeminiAgentModel, GeminiModelName
 
 try:
@@ -18,11 +18,11 @@ try:
     from google.auth.credentials import Credentials as BaseCredentials
     from google.auth.transport.requests import Request
     from google.oauth2.service_account import Credentials as ServiceAccountCredentials
-except ImportError as e:
+except ImportError as _import_error:
     raise ImportError(
         'Please install `google-auth` to use the VertexAI model, '
-        "you can use the `vertexai` optional group — `pip install 'pydantic-ai[vertexai]'`"
-    ) from e
+        "you can use the `vertexai` optional group — `pip install 'pydantic-ai-slim[vertexai]'`"
+    ) from _import_error
 
 VERTEX_AI_URL_TEMPLATE = (
     'https://{region}-aiplatform.googleapis.com/v1'
@@ -109,11 +109,12 @@ class VertexAIModel(Model):
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
-        url, auth = await self._ainit()
+        url, auth = await self.ainit()
         return GeminiAgentModel(
             http_client=self.http_client,
             model_name=self.model_name,
@@ -124,7 +125,11 @@ class VertexAIModel(Model):
             result_tools=result_tools,
         )
 
-    async def _ainit(self) -> tuple[str, BearerTokenAuth]:
+    async def ainit(self) -> tuple[str, BearerTokenAuth]:
+        """Initialize the model, setting the URL and auth.
+
+        This will raise an error if authentication fails.
+        """
         if self.url is not None and self.auth is not None:
             return self.url, self.auth
 

pydantic_ai/result.py

@@ -1,14 +1,14 @@
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic, TypeVar, cast
 
 import logfire_api
 
-from . import _result, _utils, exceptions, messages, models
+from . import _result, _utils, exceptions, messages as _messages, models
 from .tools import AgentDeps
 
 __all__ = (
@@ -71,19 +71,19 @@ class _BaseRunResult(ABC, Generic[ResultData]):
     You should not import or use this type directly, instead use its subclasses `RunResult` and `StreamedRunResult`.
     """
 
-    _all_messages: list[messages.Message]
+    _all_messages: list[_messages.ModelMessage]
     _new_message_index: int
 
-    def all_messages(self) -> list[messages.Message]:
-        """Return the history of messages."""
+    def all_messages(self) -> list[_messages.ModelMessage]:
+        """Return the history of _messages."""
         # this is a method to be consistent with the other methods
         return self._all_messages
 
     def all_messages_json(self) -> bytes:
-        """Return all messages from [`all_messages`][..all_messages] as JSON bytes."""
-        return messages.MessagesTypeAdapter.dump_json(self.all_messages())
+        """Return all messages from [`all_messages`][pydantic_ai.result._BaseRunResult.all_messages] as JSON bytes."""
+        return _messages.ModelMessagesTypeAdapter.dump_json(self.all_messages())
 
-    def new_messages(self) -> list[messages.Message]:
+    def new_messages(self) -> list[_messages.ModelMessage]:
         """Return new messages associated with this run.
 
         System prompts and any messages from older runs are excluded.
@@ -91,8 +91,8 @@ class _BaseRunResult(ABC, Generic[ResultData]):
         return self.all_messages()[self._new_message_index :]
 
     def new_messages_json(self) -> bytes:
-        """Return new messages from [`new_messages`][..new_messages] as JSON bytes."""
-        return messages.MessagesTypeAdapter.dump_json(self.new_messages())
+        """Return new messages from [`new_messages`][pydantic_ai.result._BaseRunResult.new_messages] as JSON bytes."""
+        return _messages.ModelMessagesTypeAdapter.dump_json(self.new_messages())
 
     @abstractmethod
     def cost(self) -> Cost:
@@ -122,7 +122,8 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
     _result_schema: _result.ResultSchema[ResultData] | None
     _deps: AgentDeps
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]]
-    _on_complete: Callable[[list[messages.Message]], None]
+    _result_tool_name: str | None
+    _on_complete: Callable[[], Awaitable[None]]
     is_complete: bool = field(default=False, init=False)
     """Whether the stream has all been received.
 
@@ -205,11 +206,11 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                     combined = await self._validate_text_result(''.join(chunks))
                     yield combined
                 lf_span.set_attribute('combined_text', combined)
-                self._marked_completed(text=combined)
+                await self._marked_completed(_messages.ModelResponse.from_text(combined))
 
     async def stream_structured(
         self, *, debounce_by: float | None = 0.1
-    ) -> AsyncIterator[tuple[messages.ModelStructuredResponse, bool]]:
+    ) -> AsyncIterator[tuple[_messages.ModelResponse, bool]]:
         """Stream the response as an async iterable of Structured LLM Messages.
 
         !!! note
@@ -230,17 +231,21 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
             else:
                 # we should already have a message at this point, yield that first if it has any content
                 msg = self._stream_response.get()
-                if any(call.has_content() for call in msg.calls):
-                    yield msg, False
+                for item in msg.parts:
+                    if isinstance(item, _messages.ToolCallPart) and item.has_content():
+                        yield msg, False
+                        break
                 async with _utils.group_by_temporal(self._stream_response, debounce_by) as group_iter:
                     async for _ in group_iter:
                         msg = self._stream_response.get()
-                        if any(call.has_content() for call in msg.calls):
-                            yield msg, False
+                        for item in msg.parts:
+                            if isinstance(item, _messages.ToolCallPart) and item.has_content():
+                                yield msg, False
+                                break
                 msg = self._stream_response.get(final=True)
                 yield msg, True
                 lf_span.set_attribute('structured_response', msg)
-                self._marked_completed(structured_message=msg)
+                await self._marked_completed(msg)
 
     async def get_data(self) -> ResultData:
         """Stream the whole response, validate and return it."""
@@ -249,12 +254,12 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         if isinstance(self._stream_response, models.StreamTextResponse):
             text = ''.join(self._stream_response.get(final=True))
             text = await self._validate_text_result(text)
-            self._marked_completed(text=text)
+            await self._marked_completed(_messages.ModelResponse.from_text(text))
             return cast(ResultData, text)
         else:
-            structured_message = self._stream_response.get(final=True)
-            self._marked_completed(structured_message=structured_message)
-            return await self.validate_structured_result(structured_message)
+            message = self._stream_response.get(final=True)
+            await self._marked_completed(message)
+            return await self.validate_structured_result(message)
 
     @property
     def is_structured(self) -> bool:
@@ -274,11 +279,12 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         return self._stream_response.timestamp()
 
     async def validate_structured_result(
-        self, message: messages.ModelStructuredResponse, *, allow_partial: bool = False
+        self, message: _messages.ModelResponse, *, allow_partial: bool = False
     ) -> ResultData:
         """Validate a structured result message."""
         assert self._result_schema is not None, 'Expected _result_schema to not be None'
-        match = self._result_schema.find_tool(message)
+        assert self._result_tool_name is not None, 'Expected _result_tool_name to not be None'
+        match = self._result_schema.find_named_tool(message.parts, self._result_tool_name)
         if match is None:
             raise exceptions.UnexpectedModelBehavior(
                 f'Invalid message, unable to find tool: {self._result_schema.tool_names()}'
@@ -288,7 +294,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         result_data = result_tool.validate(call, allow_partial=allow_partial, wrap_validation_errors=False)
 
         for validator in self._result_validators:
-            result_data = await validator.validate(result_data, self._deps, 0, call)
+            result_data = await validator.validate(result_data, self._deps, 0, call, self._all_messages)
         return result_data
 
     async def _validate_text_result(self, text: str) -> str:
@@ -298,19 +304,11 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                 self._deps,
                 0,
                 None,
+                self._all_messages,
             )
         return text
 
-    def _marked_completed(
-        self, *, text: str | None = None, structured_message: messages.ModelStructuredResponse | None = None
-    ) -> None:
+    async def _marked_completed(self, message: _messages.ModelResponse) -> None:
         self.is_complete = True
-        if text is not None:
-            assert structured_message is None, 'Either text or structured_message should provided, not both'
-            self._all_messages.append(
-                messages.ModelTextResponse(content=text, timestamp=self._stream_response.timestamp())
-            )
-        else:
-            assert structured_message is not None, 'Either text or structured_message should provided, not both'
-            self._all_messages.append(structured_message)
-        self._on_complete(self._all_messages)
+        self._all_messages.append(message)
+        await self._on_complete()

pydantic_ai/settings.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from httpx import Timeout
+from typing_extensions import TypedDict
+
+
+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