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

pydantic_ai/models/test.py

@@ -0,0 +1,389 @@
+"""Utilities for testing apps built with PydanticAI."""
+
+from __future__ import annotations as _annotations
+
+import re
+import string
+from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Literal
+
+import pydantic_core
+
+from .. import _utils
+from ..messages import (
+    Message,
+    ModelAnyResponse,
+    ModelStructuredResponse,
+    ModelTextResponse,
+    RetryPrompt,
+    ToolCall,
+    ToolReturn,
+)
+from ..result import Cost
+from . import (
+    AbstractToolDefinition,
+    AgentModel,
+    EitherStreamedResponse,
+    Model,
+    StreamStructuredResponse,
+    StreamTextResponse,
+)
+
+
+class UnSetType:
+    def __repr__(self):
+        return 'UnSet'
+
+
+UnSet = UnSetType()
+
+
+@dataclass
+class TestModel(Model):
+    """A model specifically for testing purposes.
+
+    This will (by default) call all retrievers in the agent model, then return a tool response if possible,
+    otherwise a plain response.
+
+    How useful this function will be is unknown, it may be useless, it may require significant changes to be useful.
+
+    Apart from `__init__` derived by the `dataclass` decorator, all methods are private or match those
+    of the base class.
+    """
+
+    # NOTE: Avoid test discovery by pytest.
+    __test__ = False
+
+    call_retrievers: list[str] | Literal['all'] = 'all'
+    """List of retrievers to call. If `'all'`, all retrievers will be called."""
+    custom_result_text: str | None = None
+    """If set, this text is return as the final result."""
+    custom_result_args: Any | None = None
+    """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_retrievers: Mapping[str, AbstractToolDefinition] | None = field(default=None, init=False)
+    agent_model_allow_text_result: bool | None = field(default=None, init=False)
+    agent_model_result_tools: list[AbstractToolDefinition] | None = field(default=None, init=False)
+
+    async def agent_model(
+        self,
+        retrievers: Mapping[str, AbstractToolDefinition],
+        allow_text_result: bool,
+        result_tools: Sequence[AbstractToolDefinition] | None,
+    ) -> AgentModel:
+        self.agent_model_retrievers = retrievers
+        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
+
+        if self.call_retrievers == 'all':
+            retriever_calls = [(r.name, r) for r in retrievers.values()]
+        else:
+            retrievers_to_call = (retrievers[name] for name in self.call_retrievers)
+            retriever_calls = [(r.name, r) for r in retrievers_to_call]
+
+        if self.custom_result_text is not None:
+            assert allow_text_result, 'Plain response not allowed, but `custom_result_text` is set.'
+            assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
+            result: _utils.Either[str | None, Any | None] = _utils.Either(left=self.custom_result_text)
+        elif self.custom_result_args is not None:
+            assert result_tools is not None, 'No result tools provided, but `custom_result_args` is set.'
+            result_tool = result_tools[0]
+
+            if k := result_tool.outer_typed_dict_key:
+                result = _utils.Either(right={k: self.custom_result_args})
+            else:
+                result = _utils.Either(right=self.custom_result_args)
+        elif allow_text_result:
+            result = _utils.Either(left=None)
+        elif result_tools is not None:
+            result = _utils.Either(right=None)
+        else:
+            result = _utils.Either(left=None)
+        return TestAgentModel(retriever_calls, result, self.agent_model_result_tools, self.seed)
+
+    def name(self) -> str:
+        return 'test-model'
+
+
+@dataclass
+class TestAgentModel(AgentModel):
+    """Implementation of `AgentModel` for testing purposes."""
+
+    # NOTE: Avoid test discovery by pytest.
+    __test__ = False
+
+    retriever_calls: list[tuple[str, AbstractToolDefinition]]
+    # 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
+    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()
+
+    @asynccontextmanager
+    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
+        msg = self._request(messages)
+        cost = Cost()
+        if isinstance(msg, ModelTextResponse):
+            yield TestStreamTextResponse(msg.content, cost)
+        else:
+            yield TestStreamStructuredResponse(msg, cost)
+
+    def gen_retriever_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.retriever_calls:
+            calls = [ToolCall.from_object(name, self.gen_retriever_args(args)) for name, args in self.retriever_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_object(name, self.gen_retriever_args(args))
+                for name, args in self.retriever_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 retriever 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 retriever calls)')
+                else:
+                    return ModelTextResponse(content=response_text.value)
+            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_object(result_tool.name, custom_result_args)])
+                else:
+                    response_args = self.gen_retriever_args(result_tool)
+                    self.step += 1
+                    return ModelStructuredResponse(calls=[ToolCall.from_object(result_tool.name, response_args)])
+
+
+@dataclass
+class TestStreamTextResponse(StreamTextResponse):
+    _text: str
+    _cost: Cost
+    _iter: Iterator[str] = field(init=False)
+    _timestamp: datetime = field(default_factory=_utils.now_utc)
+    _buffer: list[str] = field(default_factory=list, init=False)
+
+    def __post_init__(self):
+        *words, last_word = self._text.split(' ')
+        words = [f'{word} ' for word in words]
+        words.append(last_word)
+        if len(words) == 1 and len(self._text) > 2:
+            mid = len(self._text) // 2
+            words = [self._text[:mid], self._text[mid:]]
+        self._iter = iter(words)
+
+    async def __anext__(self) -> None:
+        self._buffer.append(_utils.sync_anext(self._iter))
+
+    def get(self, *, final: bool = False) -> Iterable[str]:
+        yield from self._buffer
+        self._buffer.clear()
+
+    def cost(self) -> Cost:
+        return self._cost
+
+    def timestamp(self) -> datetime:
+        return self._timestamp
+
+
+@dataclass
+class TestStreamStructuredResponse(StreamStructuredResponse):
+    _structured_response: ModelStructuredResponse
+    _cost: Cost
+    _iter: Iterator[None] = field(default_factory=lambda: iter([None]))
+    _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
+
+    async def __anext__(self) -> None:
+        return _utils.sync_anext(self._iter)
+
+    def get(self, *, final: bool = False) -> ModelStructuredResponse:
+        return self._structured_response
+
+    def cost(self) -> Cost:
+        return self._cost
+
+    def timestamp(self) -> datetime:
+        return self._timestamp
+
+
+_chars = string.ascii_letters + string.digits + string.punctuation
+
+
+class _JsonSchemaTestData:
+    """Generate data that matches a JSON schema.
+
+    This tries to generate the minimal viable data for the schema.
+    """
+
+    def __init__(self, schema: _utils.ObjectJsonSchema, seed: int = 0):
+        self.schema = schema
+        self.defs = schema.get('$defs', {})
+        self.seed = seed
+
+    def generate(self) -> Any:
+        """Generate data for the JSON schema."""
+        return self._gen_any(self.schema)
+
+    def _gen_any(self, schema: dict[str, Any]) -> Any:
+        """Generate data for any JSON Schema."""
+        if const := schema.get('const'):
+            return const
+        elif enum := schema.get('enum'):
+            return enum[self.seed % len(enum)]
+        elif examples := schema.get('examples'):
+            return examples[self.seed % len(examples)]
+        elif ref := schema.get('$ref'):
+            key = re.sub(r'^#/\$defs/', '', ref)
+            js_def = self.defs[key]
+            return self._gen_any(js_def)
+        elif any_of := schema.get('anyOf'):
+            return self._gen_any(any_of[self.seed % len(any_of)])
+
+        type_ = schema.get('type')
+        if type_ is None:
+            # if there's no type or ref, we can't generate anything
+            return self._char()
+        elif type_ == 'object':
+            return self._object_gen(schema)
+        elif type_ == 'string':
+            return self._str_gen(schema)
+        elif type_ == 'integer':
+            return self._int_gen(schema)
+        elif type_ == 'number':
+            return float(self._int_gen(schema))
+        elif type_ == 'boolean':
+            return self._bool_gen()
+        elif type_ == 'array':
+            return self._array_gen(schema)
+        elif type_ == 'null':
+            return None
+        else:
+            raise NotImplementedError(f'Unknown type: {type_}, please submit a PR to extend JsonSchemaTestData!')
+
+    def _object_gen(self, schema: dict[str, Any]) -> dict[str, Any]:
+        """Generate data for a JSON Schema object."""
+        required = set(schema.get('required', []))
+
+        data: dict[str, Any] = {}
+        if properties := schema.get('properties'):
+            for key, value in properties.items():
+                if key in required:
+                    data[key] = self._gen_any(value)
+
+        if addition_props := schema.get('additionalProperties'):
+            add_prop_key = 'additionalProperty'
+            while add_prop_key in data:
+                add_prop_key += '_'
+            if addition_props is True:
+                data[add_prop_key] = self._char()
+            else:
+                data[add_prop_key] = self._gen_any(addition_props)
+
+        return data
+
+    def _str_gen(self, schema: dict[str, Any]) -> str:
+        """Generate a string from a JSON Schema string."""
+        min_len = schema.get('minLength')
+        if min_len is not None:
+            return self._char() * min_len
+
+        if schema.get('maxLength') == 0:
+            return ''
+        else:
+            return self._char()
+
+    def _int_gen(self, schema: dict[str, Any]) -> int:
+        """Generate an integer from a JSON Schema integer."""
+        maximum = schema.get('maximum')
+        if maximum is None:
+            exc_max = schema.get('exclusiveMaximum')
+            if exc_max is not None:
+                maximum = exc_max - 1
+
+        minimum = schema.get('minimum')
+        if minimum is None:
+            exc_min = schema.get('exclusiveMinimum')
+            if exc_min is not None:
+                minimum = exc_min + 1
+
+        if minimum is not None and maximum is not None:
+            return minimum + self.seed % (maximum - minimum)
+        elif minimum is not None:
+            return minimum + self.seed
+        elif maximum is not None:
+            return maximum - self.seed
+        else:
+            return self.seed
+
+    def _bool_gen(self) -> bool:
+        """Generate a boolean from a JSON Schema boolean."""
+        return bool(self.seed % 2)
+
+    def _array_gen(self, schema: dict[str, Any]) -> list[Any]:
+        """Generate an array from a JSON Schema array."""
+        data: list[Any] = []
+        unique_items = schema.get('uniqueItems')
+        if prefix_items := schema.get('prefixItems'):
+            for item in prefix_items:
+                data.append(self._gen_any(item))
+                if unique_items:
+                    self.seed += 1
+
+        items_schema = schema.get('items', {})
+        min_items = schema.get('minItems', 0)
+        if min_items > len(data):
+            for _ in range(min_items - len(data)):
+                data.append(self._gen_any(items_schema))
+                if unique_items:
+                    self.seed += 1
+        elif items_schema:
+            # if there is an `items` schema, add an item unless it would break `maxItems` rule
+            max_items = schema.get('maxItems')
+            if max_items is None or max_items > len(data):
+                data.append(self._gen_any(items_schema))
+                if unique_items:
+                    self.seed += 1
+
+        return data
+
+    def _char(self) -> str:
+        """Generate a character on the same principle as Excel columns, e.g. a-z, aa-az..."""
+        chars = len(_chars)
+        s = ''
+        rem = self.seed // chars
+        while rem > 0:
+            s += _chars[(rem - 1) % chars]
+            rem //= chars
+        s += _chars[self.seed % chars]
+        return s

pydantic_ai/models/vertexai.py

@@ -0,0 +1,306 @@
+"""Custom interface to the `*-aiplatform.googleapis.com` API for Gemini models.
+
+This model uses [`GeminiAgentModel`][pydantic_ai.models.gemini.GeminiAgentModel] with just the URL and auth method
+changed from the default `GeminiModel`, it relies on the VertexAI
+[`generateContent` function endpoint](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent)
+and `streamGenerateContent` function endpoints
+having the same schemas as the equivalent [Gemini endpoints][pydantic_ai.models.gemini.GeminiModel].
+
+There are four advantages of using this API over the `generativelanguage.googleapis.com` API which
+[`GeminiModel`][pydantic_ai.models.gemini.GeminiModel] uses, and one big disadvantage.
+
+Advantages:
+
+1. The VertexAI API seems to be less flakey, less likely to occasionally return a 503 response.
+2. You can
+   [purchase provisioned throughput](https://cloud.google.com/vertex-ai/generative-ai/docs/provisioned-throughput#purchase-provisioned-throughput)
+   with VertexAI.
+3. If you're running PydanticAI inside GCP, you don't need to set up authentication, it should "just work".
+4. You can decide which region to use, which might be important from a regulatory perspective,
+   and might improve latency.
+
+Disadvantage:
+
+1. When authorization doesn't just work, it's much more painful to set up than an API key.
+
+## Example Usage
+
+With the default google project already configured in your environment:
+
+```py title="vertex_example_env.py"
+from pydantic_ai import Agent
+from pydantic_ai.models.vertexai import VertexAIModel
+
+model = VertexAIModel('gemini-1.5-flash')
+agent = Agent(model)
+result = agent.run_sync('Tell me a joke.')
+print(result.data)
+#> Did you hear about the toothpaste scandal? They called it Colgate.
+```
+
+Or using a service account JSON file:
+
+```py title="vertex_example_service_account.py"
+from pydantic_ai import Agent
+from pydantic_ai.models.vertexai import VertexAIModel
+
+model = VertexAIModel(
+    'gemini-1.5-flash',
+    service_account_file='path/to/service-account.json',
+)
+agent = Agent(model)
+result = agent.run_sync('Tell me a joke.')
+print(result.data)
+#> Did you hear about the toothpaste scandal? They called it Colgate.
+```
+"""
+
+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
+from typing import Literal
+
+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 .gemini import GeminiAgentModel, GeminiModelName
+
+try:
+    import google.auth
+    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:
+    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
+
+VERTEX_AI_URL_TEMPLATE = (
+    'https://{region}-aiplatform.googleapis.com/v1'
+    '/projects/{project_id}'
+    '/locations/{region}'
+    '/publishers/{model_publisher}'
+    '/models/{model}'
+    ':'
+)
+"""URL template for Vertex AI.
+
+See
+[`generateContent` docs](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent)
+and
+[`streamGenerateContent` docs](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/streamGenerateContent)
+for more information.
+
+The template is used thus:
+
+* `region` is substituted with the `region` argument,
+  see [available regions][pydantic_ai.models.vertexai.VertexAiRegion]
+* `model_publisher` is substituted with the `model_publisher` argument
+* `model` is substituted with the `model_name` argument
+* `project_id` is substituted with the `project_id` from auth/credentials
+* `function` (`generateContent` or `streamGenerateContent`) is added to the end of the URL
+"""
+
+
+@dataclass(init=False)
+class VertexAIModel(Model):
+    """A model that uses Gemini via the `*-aiplatform.googleapis.com` VertexAI API."""
+
+    model_name: GeminiModelName
+    service_account_file: Path | str | None
+    project_id: str | None
+    region: VertexAiRegion
+    model_publisher: Literal['google']
+    http_client: AsyncHTTPClient
+    url_template: str
+
+    auth: BearerTokenAuth | None
+    url: str | None
+
+    # TODO __init__ can be removed once we drop 3.9 and we can set kw_only correctly on the dataclass
+    def __init__(
+        self,
+        model_name: GeminiModelName,
+        *,
+        service_account_file: Path | str | None = None,
+        project_id: str | None = None,
+        region: VertexAiRegion = 'us-central1',
+        model_publisher: Literal['google'] = 'google',
+        http_client: AsyncHTTPClient | None = None,
+        url_template: str = VERTEX_AI_URL_TEMPLATE,
+    ):
+        """Initialize a Vertex AI Gemini model.
+
+        Args:
+            model_name: The name of the model to use. I couldn't find a list of supported Google models, in VertexAI
+                so for now this uses the same models as the [Gemini model][pydantic_ai.models.gemini.GeminiModel].
+            service_account_file: Path to a service account file.
+                If not provided, the default environment credentials will be used.
+            project_id: The project ID to use, if not provided it will be taken from the credentials.
+            region: The region to make requests to.
+            model_publisher: The model publisher to use, I couldn't find a good list of available publishers,
+                and from trial and error it seems non-google models don't work with the `generateContent` and
+                `streamGenerateContent` functions, hence only `google` is currently supported.
+                Please create an issue or PR if you know how to use other publishers.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+            url_template: URL template for Vertex AI, see
+                [`VERTEX_AI_URL_TEMPLATE` docs][pydantic_ai.models.vertexai.VERTEX_AI_URL_TEMPLATE]
+                for more information.
+        """
+        self.model_name = model_name
+        self.service_account_file = service_account_file
+        self.project_id = project_id
+        self.region = region
+        self.model_publisher = model_publisher
+        self.http_client = http_client or cached_async_http_client()
+        self.url_template = url_template
+
+        self.auth = None
+        self.url = None
+
+    async def agent_model(
+        self,
+        retrievers: Mapping[str, AbstractToolDefinition],
+        allow_text_result: bool,
+        result_tools: Sequence[AbstractToolDefinition] | None,
+    ) -> GeminiAgentModel:
+        url, auth = await self._ainit()
+        return GeminiAgentModel(
+            http_client=self.http_client,
+            model_name=self.model_name,
+            auth=auth,
+            url=url,
+            retrievers=retrievers,
+            allow_text_result=allow_text_result,
+            result_tools=result_tools,
+        )
+
+    async def _ainit(self) -> tuple[str, BearerTokenAuth]:
+        if self.url is not None and self.auth is not None:
+            return self.url, self.auth
+
+        if self.service_account_file is not None:
+            creds: BaseCredentials | ServiceAccountCredentials = _creds_from_file(self.service_account_file)
+            assert creds.project_id is None or isinstance(creds.project_id, str)
+            creds_project_id: str | None = creds.project_id
+            creds_source = 'service account file'
+        else:
+            creds, creds_project_id = await _async_google_auth()
+            creds_source = '`google.auth.default()`'
+
+        if self.project_id is None:
+            if creds_project_id is None:
+                raise UserError(f'No project_id provided and none found in {creds_source}')
+            project_id = creds_project_id
+        else:
+            if creds_project_id is not None and self.project_id != creds_project_id:
+                raise UserError(
+                    f'The project_id you provided does not match the one from {creds_source}: '
+                    f'{self.project_id!r} != {creds_project_id!r}'
+                )
+            project_id = self.project_id
+
+        self.url = url = self.url_template.format(
+            region=self.region,
+            project_id=project_id,
+            model_publisher=self.model_publisher,
+            model=self.model_name,
+        )
+        self.auth = auth = BearerTokenAuth(creds)
+        return url, auth
+
+    def name(self) -> str:
+        return f'vertexai:{self.model_name}'
+
+
+# pyright: reportUnknownMemberType=false
+def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredentials:
+    return ServiceAccountCredentials.from_service_account_file(
+        str(service_account_file), scopes=['https://www.googleapis.com/auth/cloud-platform']
+    )
+
+
+# pyright: reportReturnType=false
+# pyright: reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false
+async def _async_google_auth() -> tuple[BaseCredentials, str | None]:
+    return await run_in_executor(google.auth.default)
+
+
+# default expiry is 3600 seconds
+MAX_TOKEN_AGE = timedelta(seconds=3000)
+
+
+@dataclass
+class BearerTokenAuth:
+    credentials: BaseCredentials | ServiceAccountCredentials
+    token_created: datetime | None = field(default=None, init=False)
+
+    async def headers(self) -> dict[str, str]:
+        if self.credentials.token is None or self._token_expired():
+            await run_in_executor(self._refresh_token)
+            self.token_created = datetime.now()
+        return {'Authorization': f'Bearer {self.credentials.token}'}
+
+    def _token_expired(self) -> bool:
+        if self.token_created is None:
+            return True
+        else:
+            return (datetime.now() - self.token_created) > MAX_TOKEN_AGE
+
+    def _refresh_token(self) -> str:
+        self.credentials.refresh(Request())
+        assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'
+        return self.credentials.token
+
+
+VertexAiRegion = Literal[
+    'us-central1',
+    'us-east1',
+    'us-east4',
+    'us-south1',
+    'us-west1',
+    'us-west2',
+    'us-west3',
+    'us-west4',
+    'us-east5',
+    'europe-central2',
+    'europe-north1',
+    'europe-southwest1',
+    'europe-west1',
+    'europe-west2',
+    'europe-west3',
+    'europe-west4',
+    'europe-west6',
+    'europe-west8',
+    'europe-west9',
+    'europe-west12',
+    'africa-south1',
+    'asia-east1',
+    'asia-east2',
+    'asia-northeast1',
+    'asia-northeast2',
+    'asia-northeast3',
+    'asia-south1',
+    'asia-southeast1',
+    'asia-southeast2',
+    'australia-southeast1',
+    'australia-southeast2',
+    'me-central1',
+    'me-central2',
+    'me-west1',
+    'northamerica-northeast1',
+    'northamerica-northeast2',
+    'southamerica-east1',
+    'southamerica-west1',
+]
+"""Regions available for Vertex AI.
+
+More details [here](https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints).
+"""