pydantic-ai-slim 0.0.6__py3-none-any.whl

pydantic_ai/dependencies.py

@@ -0,0 +1,83 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Awaitable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
+
+from typing_extensions import Concatenate, ParamSpec, TypeAlias
+
+if TYPE_CHECKING:
+    from .result import ResultData
+else:
+    ResultData = Any
+
+__all__ = (
+    'AgentDeps',
+    'CallContext',
+    'ResultValidatorFunc',
+    'SystemPromptFunc',
+    'RetrieverReturnValue',
+    'RetrieverContextFunc',
+    'RetrieverPlainFunc',
+    'RetrieverParams',
+    'JsonData',
+)
+
+AgentDeps = TypeVar('AgentDeps')
+"""Type variable for agent dependencies."""
+
+
+@dataclass
+class CallContext(Generic[AgentDeps]):
+    """Information about the current call."""
+
+    deps: AgentDeps
+    """Dependencies for the agent."""
+    retry: int
+    """Number of retries so far."""
+    tool_name: str | None
+    """Name of the tool being called."""
+
+
+RetrieverParams = ParamSpec('RetrieverParams')
+"""Retrieval function param spec."""
+
+SystemPromptFunc = Union[
+    Callable[[CallContext[AgentDeps]], str],
+    Callable[[CallContext[AgentDeps]], Awaitable[str]],
+    Callable[[], str],
+    Callable[[], Awaitable[str]],
+]
+"""A function that may or maybe not take `CallContext` as an argument, and may or may not be async.
+
+Usage `SystemPromptFunc[AgentDeps]`.
+"""
+
+ResultValidatorFunc = Union[
+    Callable[[CallContext[AgentDeps], ResultData], ResultData],
+    Callable[[CallContext[AgentDeps], ResultData], Awaitable[ResultData]],
+    Callable[[ResultData], ResultData],
+    Callable[[ResultData], Awaitable[ResultData]],
+]
+"""
+A function that always takes `ResultData` and returns `ResultData`,
+but may or maybe not take `CallInfo` as a first argument, and may or may not be async.
+
+Usage `ResultValidator[AgentDeps, ResultData]`.
+"""
+
+JsonData: TypeAlias = 'None | str | int | float | Sequence[JsonData] | Mapping[str, JsonData]'
+"""Type representing any JSON data."""
+
+RetrieverReturnValue = Union[JsonData, Awaitable[JsonData]]
+"""Return value of a retriever function."""
+RetrieverContextFunc = Callable[Concatenate[CallContext[AgentDeps], RetrieverParams], RetrieverReturnValue]
+"""A retriever function that takes `CallContext` as the first argument.
+
+Usage `RetrieverContextFunc[AgentDeps, RetrieverParams]`.
+"""
+RetrieverPlainFunc = Callable[RetrieverParams, RetrieverReturnValue]
+"""A retriever function that does not take `CallContext` as the first argument.
+
+Usage `RetrieverPlainFunc[RetrieverParams]`.
+"""

pydantic_ai/exceptions.py

@@ -0,0 +1,56 @@
+from __future__ import annotations as _annotations
+
+import json
+
+__all__ = 'ModelRetry', 'UserError', 'UnexpectedModelBehavior'
+
+
+class ModelRetry(Exception):
+    """Exception raised when a retriever function should be retried.
+
+    The agent will return the message to the model and ask it to try calling the function/tool again.
+    """
+
+    message: str
+    """The message to return to the model."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class UserError(RuntimeError):
+    """Error caused by a usage mistake by the application developer — You!"""
+
+    message: str
+    """Description of the mistake."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class UnexpectedModelBehavior(RuntimeError):
+    """Error caused by unexpected Model behavior, e.g. an unexpected response code."""
+
+    message: str
+    """Description of the unexpected behavior."""
+    body: str | None
+    """The body of the response, if available."""
+
+    def __init__(self, message: str, body: str | None = None):
+        self.message = message
+        if body is None:
+            self.body: str | None = None
+        else:
+            try:
+                self.body = json.dumps(json.loads(body), indent=2)
+            except ValueError:
+                self.body = body
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        if self.body:
+            return f'{self.message}, body:\n{self.body}'
+        else:
+            return self.message

pydantic_ai/messages.py

@@ -0,0 +1,205 @@
+from __future__ import annotations as _annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
+
+import pydantic
+import pydantic_core
+from pydantic import TypeAdapter
+from typing_extensions import TypeAlias, TypeAliasType
+
+from . import _pydantic
+from ._utils import now_utc as _now_utc
+
+
+@dataclass
+class SystemPrompt:
+    """A system prompt, generally written by the application developer.
+
+    This gives the model context and guidance on how to respond.
+    """
+
+    content: str
+    """The content of the prompt."""
+    role: Literal['system'] = 'system'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+
+@dataclass
+class UserPrompt:
+    """A user prompt, generally written by the end user.
+
+    Content comes from the `user_prompt` parameter of [`Agent.run`][pydantic_ai.Agent.run],
+    [`Agent.run_sync`][pydantic_ai.Agent.run_sync], and [`Agent.run_stream`][pydantic_ai.Agent.run_stream].
+    """
+
+    content: str
+    """The content of the prompt."""
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp of the prompt."""
+    role: Literal['user'] = 'user'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+
+JsonData: TypeAlias = 'Union[str, int, float, None, Sequence[JsonData], Mapping[str, JsonData]]'
+if not TYPE_CHECKING:
+    # work around for https://github.com/pydantic/pydantic/issues/10873
+    # this is need for pydantic to work both `json_ta` and `MessagesTypeAdapter` at the bottom of this file
+    JsonData = TypeAliasType('JsonData', 'Union[str, int, float, None, Sequence[JsonData], Mapping[str, JsonData]]')
+
+json_ta: TypeAdapter[JsonData] = TypeAdapter(JsonData)
+
+
+@dataclass
+class ToolReturn:
+    """A tool return message, this encodes the result of running a retriever."""
+
+    tool_name: str
+    """The name of the "tool" was called."""
+    content: JsonData
+    """The return value."""
+    tool_id: str | None = None
+    """Optional tool identifier, this is used by some models including OpenAI."""
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp, when the tool returned."""
+    role: Literal['tool-return'] = 'tool-return'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+    def model_response_str(self) -> str:
+        if isinstance(self.content, str):
+            return self.content
+        else:
+            content = json_ta.validate_python(self.content)
+            return json_ta.dump_json(content).decode()
+
+    def model_response_object(self) -> dict[str, JsonData]:
+        # gemini supports JSON dict return values, but no other JSON types, hence we wrap anything else in a dict
+        if isinstance(self.content, dict):
+            return json_ta.validate_python(self.content)  # pyright: ignore[reportReturnType]
+        else:
+            return {'return_value': json_ta.validate_python(self.content)}
+
+
+@dataclass
+class RetryPrompt:
+    """A message back to a model asking it to try again.
+
+    This can be sent for a number of reasons:
+
+    * Pydantic validation of retriever arguments failed, here content is derived from a Pydantic
+      [`ValidationError`][pydantic_core.ValidationError]
+    * a retriever raised a [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] exception
+    * no retriever was found for the tool name
+    * the model returned plain text when a structured response was expected
+    * Pydantic validation of a structured response failed, here content is derived from a Pydantic
+      [`ValidationError`][pydantic_core.ValidationError]
+    * a result validator raised a [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] exception
+    """
+
+    content: list[pydantic_core.ErrorDetails] | str
+    """Details of why and how the model should retry.
+
+    If the retry was triggered by a [`ValidationError`][pydantic_core.ValidationError], this will be a list of
+    error details.
+    """
+    tool_name: str | None = None
+    """The name of the tool that was called, if any."""
+    tool_id: str | None = None
+    """The tool identifier, if any."""
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp, when the retry was triggered."""
+    role: Literal['retry-prompt'] = 'retry-prompt'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+    def model_response(self) -> str:
+        if isinstance(self.content, str):
+            description = self.content
+        else:
+            description = f'{len(self.content)} validation errors: {json.dumps(self.content, indent=2)}'
+        return f'{description}\n\nFix the errors and try again.'
+
+
+@dataclass
+class ModelTextResponse:
+    """A plain text response from a model."""
+
+    content: str
+    """The text content of the response."""
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp of the response.
+
+    If the model provides a timestamp in the response (as OpenAI does) that will be used.
+    """
+    role: Literal['model-text-response'] = 'model-text-response'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+
+@dataclass
+class ArgsJson:
+    args_json: str
+    """A JSON string of arguments."""
+
+
+@dataclass
+class ArgsObject:
+    args_object: dict[str, Any]
+    """A python dictionary of arguments."""
+
+
+@dataclass
+class ToolCall:
+    """Either a tool call from the agent."""
+
+    tool_name: str
+    """The name of the tool to call."""
+    args: ArgsJson | ArgsObject
+    """The arguments to pass to the tool.
+
+    Either as JSON or a Python dictionary depending on how data was returned.
+    """
+    tool_id: str | None = None
+    """Optional tool identifier, this is used by some models including OpenAI."""
+
+    @classmethod
+    def from_json(cls, tool_name: str, args_json: str, tool_id: str | None = None) -> ToolCall:
+        return cls(tool_name, ArgsJson(args_json), tool_id)
+
+    @classmethod
+    def from_object(cls, tool_name: str, args_object: dict[str, Any]) -> ToolCall:
+        return cls(tool_name, ArgsObject(args_object))
+
+    def has_content(self) -> bool:
+        if isinstance(self.args, ArgsObject):
+            return any(self.args.args_object.values())
+        else:
+            return bool(self.args.args_json)
+
+
+@dataclass
+class ModelStructuredResponse:
+    """A structured response from a model.
+
+    This is used either to call a retriever or to return a structured response from an agent run.
+    """
+
+    calls: list[ToolCall]
+    """The tool calls being made."""
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp of the response.
+
+    If the model provides a timestamp in the response (as OpenAI does) that will be used.
+    """
+    role: Literal['model-structured-response'] = 'model-structured-response'
+    """Message type identifier, this type is available on all message as a discriminator."""
+
+
+ModelAnyResponse = Union[ModelTextResponse, ModelStructuredResponse]
+"""Any response from a model."""
+Message = Union[SystemPrompt, UserPrompt, ToolReturn, RetryPrompt, ModelTextResponse, ModelStructuredResponse]
+"""Any message send to or returned by a model."""
+
+MessagesTypeAdapter = _pydantic.LazyTypeAdapter(list[Annotated[Message, pydantic.Field(discriminator='role')]])
+"""Pydantic [`TypeAdapter`][pydantic.type_adapter.TypeAdapter] for (de)serializing messages."""

pydantic_ai/models/__init__.py

@@ -0,0 +1,300 @@
+"""Logic related to making requests to an LLM.
+
+The aim here is to make a common interface for different LLMs, so that the rest of the code can be agnostic to the
+specific LLM being used.
+"""
+
+from __future__ import annotations as _annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
+from functools import cache
+from typing import TYPE_CHECKING, Literal, Protocol, Union
+
+import httpx
+
+from ..exceptions import UserError
+from ..messages import Message, ModelAnyResponse, ModelStructuredResponse
+
+if TYPE_CHECKING:
+    from .._utils import ObjectJsonSchema
+    from ..result import Cost
+
+
+KnownModelName = Literal[
+    'openai:gpt-4o',
+    'openai:gpt-4o-mini',
+    'openai:gpt-4-turbo',
+    'openai:gpt-4',
+    'openai:o1-preview',
+    'openai:o1-mini',
+    'openai:gpt-3.5-turbo',
+    'groq:llama-3.1-70b-versatile',
+    'groq:llama3-groq-70b-8192-tool-use-preview',
+    'groq:llama3-groq-8b-8192-tool-use-preview',
+    'groq:llama-3.1-70b-specdec',
+    'groq:llama-3.1-8b-instant',
+    'groq:llama-3.2-1b-preview',
+    'groq:llama-3.2-3b-preview',
+    'groq:llama-3.2-11b-vision-preview',
+    'groq:llama-3.2-90b-vision-preview',
+    'groq:llama3-70b-8192',
+    'groq:llama3-8b-8192',
+    'groq:mixtral-8x7b-32768',
+    'groq:gemma2-9b-it',
+    'groq:gemma-7b-it',
+    'gemini-1.5-flash',
+    'gemini-1.5-pro',
+    'vertexai:gemini-1.5-flash',
+    'vertexai:gemini-1.5-pro',
+    'test',
+]
+"""Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
+
+`KnownModelName` is provided as a concise way to specify a model.
+"""
+
+
+class Model(ABC):
+    """Abstract class for a model."""
+
+    @abstractmethod
+    async def agent_model(
+        self,
+        retrievers: Mapping[str, AbstractToolDefinition],
+        allow_text_result: bool,
+        result_tools: Sequence[AbstractToolDefinition] | None,
+    ) -> AgentModel:
+        """Create an agent model.
+
+        This is async in case slow/async config checks need to be performed that can't be done in `__init__`.
+
+        Args:
+            retrievers: The retrievers available to the agent.
+            allow_text_result: Whether a plain text final response/result is permitted.
+            result_tools: Tool definitions for the final result tool(s), if any.
+
+        Returns:
+            An agent model.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def name(self) -> str:
+        raise NotImplementedError()
+
+
+class AgentModel(ABC):
+    """Model configured for a specific agent."""
+
+    @abstractmethod
+    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
+        """Make a request to the model."""
+        raise NotImplementedError()
+
+    @asynccontextmanager
+    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
+        """Make a request to the model and return a streaming response."""
+        raise NotImplementedError(f'Streamed requests not supported by this {self.__class__.__name__}')
+        # yield is required to make this a generator for type checking
+        # noinspection PyUnreachableCode
+        yield  # pragma: no cover
+
+
+class StreamTextResponse(ABC):
+    """Streamed response from an LLM when returning text."""
+
+    def __aiter__(self) -> AsyncIterator[None]:
+        """Stream the response as an async iterable, building up the text as it goes.
+
+        This is an async iterator that yields `None` to avoid doing the work of validating the input and
+        extracting the text field when it will often be thrown away.
+        """
+        return self
+
+    @abstractmethod
+    async def __anext__(self) -> None:
+        """Process the next chunk of the response, see above for why this returns `None`."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get(self, *, final: bool = False) -> Iterable[str]:
+        """Returns an iterable of text since the last call to `get()` — e.g. the text delta.
+
+        Args:
+            final: If True, this is the final call, after iteration is complete, the response should be fully validated
+                and all text extracted.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def cost(self) -> Cost:
+        """Return the cost of the request.
+
+        NOTE: this won't return the ful cost until the stream is finished.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
+        raise NotImplementedError()
+
+
+class StreamStructuredResponse(ABC):
+    """Streamed response from an LLM when calling a tool."""
+
+    def __aiter__(self) -> AsyncIterator[None]:
+        """Stream the response as an async iterable, building up the tool call as it goes.
+
+        This is an async iterator that yields `None` to avoid doing the work of building the final tool call when
+        it will often be thrown away.
+        """
+        return self
+
+    @abstractmethod
+    async def __anext__(self) -> None:
+        """Process the next chunk of the response, see above for why this returns `None`."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get(self, *, final: bool = False) -> ModelStructuredResponse:
+        """Get the `ModelStructuredResponse` at this point.
+
+        The `ModelStructuredResponse` may or may not be complete, depending on whether the stream is finished.
+
+        Args:
+            final: If True, this is the final call, after iteration is complete, the response should be fully validated.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def cost(self) -> Cost:
+        """Get the cost of the request.
+
+        NOTE: this won't return the full cost until the stream is finished.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
+        raise NotImplementedError()
+
+
+EitherStreamedResponse = Union[StreamTextResponse, StreamStructuredResponse]
+
+
+ALLOW_MODEL_REQUESTS = True
+"""Whether to allow requests to models.
+
+This global setting allows you to disable request to most models, e.g. to make sure you don't accidentally
+make costly requests to a model during tests.
+
+The testing models [`TestModel`][pydantic_ai.models.test.TestModel] and
+[`FunctionModel`][pydantic_ai.models.function.FunctionModel] are no affected by this setting.
+"""
+
+
+def check_allow_model_requests() -> None:
+    """Check if model requests are allowed.
+
+    If you're defining your own models that have cost or latency associated with their use, you should call this in
+    [`Model.agent_model`][pydantic_ai.models.Model.agent_model].
+
+    Raises:
+        RuntimeError: If model requests are not allowed.
+    """
+    if not ALLOW_MODEL_REQUESTS:
+        raise RuntimeError('Model requests are not allowed, since ALLOW_MODEL_REQUESTS is False')
+
+
+@contextmanager
+def override_allow_model_requests(allow_model_requests: bool) -> Iterator[None]:
+    """Context manager to temporarily override [`ALLOW_MODEL_REQUESTS`][pydantic_ai.models.ALLOW_MODEL_REQUESTS].
+
+    Args:
+        allow_model_requests: Whether to allow model requests within the context.
+    """
+    global ALLOW_MODEL_REQUESTS
+    old_value = ALLOW_MODEL_REQUESTS
+    ALLOW_MODEL_REQUESTS = allow_model_requests  # pyright: ignore[reportConstantRedefinition]
+    try:
+        yield
+    finally:
+        ALLOW_MODEL_REQUESTS = old_value  # pyright: ignore[reportConstantRedefinition]
+
+
+def infer_model(model: Model | KnownModelName) -> Model:
+    """Infer the model from the name."""
+    if isinstance(model, Model):
+        return model
+    elif model == 'test':
+        from .test import TestModel
+
+        return TestModel()
+    elif model.startswith('openai:'):
+        from .openai import OpenAIModel
+
+        return OpenAIModel(model[7:])  # pyright: ignore[reportArgumentType]
+    elif model.startswith('gemini'):
+        from .gemini import GeminiModel
+
+        # noinspection PyTypeChecker
+        return GeminiModel(model)  # pyright: ignore[reportArgumentType]
+    elif model.startswith('groq:'):
+        from .groq import GroqModel
+
+        return GroqModel(model[5:])  # pyright: ignore[reportArgumentType]
+    elif model.startswith('vertexai:'):
+        from .vertexai import VertexAIModel
+
+        return VertexAIModel(model[9:])  # pyright: ignore[reportArgumentType]
+    else:
+        raise UserError(f'Unknown model: {model}')
+
+
+class AbstractToolDefinition(Protocol):
+    """Abstract definition of a function/tool.
+
+    This is used for both retrievers and result tools.
+    """
+
+    name: str
+    """The name of the tool."""
+    description: str
+    """The description of the tool."""
+    json_schema: ObjectJsonSchema
+    """The JSON schema for the tool's arguments."""
+    outer_typed_dict_key: str | None
+    """The key in the outer [TypedDict] that wraps a result tool.
+
+    This will only be set for result tools which don't have an `object` JSON schema.
+    """
+
+
+@cache
+def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
+    """Cached HTTPX async client so multiple agents and calls can share the same client.
+
+    There are good reasons why in production you should use a `httpx.AsyncClient` as an async context manager as
+    described in [encode/httpx#2026](https://github.com/encode/httpx/pull/2026), but when experimenting or showing
+    examples, it's very useful not to, this allows multiple Agents to use a single client.
+
+    The default timeouts match those of OpenAI,
+    see <https://github.com/openai/openai-python/blob/v1.54.4/src/openai/_constants.py#L9>.
+    """
+    return httpx.AsyncClient(
+        timeout=httpx.Timeout(timeout=timeout, connect=connect),
+        headers={'User-Agent': get_user_agent()},
+    )
+
+
+@cache
+def get_user_agent() -> str:
+    """Get the user agent string for the HTTP client."""
+    from .. import __version__
+
+    return f'pydantic-ai/{__version__}'