pydantic-ai-slim 0.0.32__tar.gz → 0.0.33__tar.gz

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.32
+Version: 0.0.33
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,8 +29,9 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.32
+Requires-Dist: pydantic-graph==0.0.33
 Requires-Dist: pydantic>=2.10
+Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
 Provides-Extra: cohere

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 from __future__ import annotations as _annotations
 
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast, get_origin
+from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -15,6 +15,7 @@ from pydantic.fields import FieldInfo
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
+from typing_extensions import get_origin
 
 from ._griffe import doc_descriptions
 from ._utils import check_object_json_schema, is_model_like
@@ -223,8 +224,7 @@ def _build_schema(
 
 
 def _is_call_ctx(annotation: Any) -> bool:
+    """Return whether the annotation is the `RunContext` class, parameterized or not."""
     from .tools import RunContext
 
-    return annotation is RunContext or (
-        _typing_extra.is_generic_alias(annotation) and get_origin(annotation) is RunContext
-    )
+    return annotation is RunContext or get_origin(annotation) is RunContext

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/_result.py

@@ -1,14 +1,14 @@
 from __future__ import annotations as _annotations
 
 import inspect
-import sys
-import types
 from collections.abc import Awaitable, Iterable, Iterator
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin
+from typing import Any, Callable, Generic, Literal, Union, cast
 
 from pydantic import TypeAdapter, ValidationError
-from typing_extensions import TypeAliasType, TypedDict, TypeVar
+from typing_extensions import TypedDict, TypeVar, get_args, get_origin
+from typing_inspection import typing_objects
+from typing_inspection.introspection import is_union_origin
 
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
@@ -248,23 +248,12 @@ def extract_str_from_union(response_type: Any) -> _utils.Option[Any]:
 
 
 def get_union_args(tp: Any) -> tuple[Any, ...]:
-    """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty union."""
-    if isinstance(tp, TypeAliasType):
+    """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty tuple."""
+    if typing_objects.is_typealiastype(tp):
         tp = tp.__value__
 
     origin = get_origin(tp)
-    if origin_is_union(origin):
+    if is_union_origin(origin):
         return get_args(tp)
     else:
         return ()
-
-
-if sys.version_info < (3, 10):
-
-    def origin_is_union(tp: type[Any] | None) -> bool:
-        return tp is Union
-
-else:
-
-    def origin_is_union(tp: type[Any] | None) -> bool:
-        return tp is Union or tp is types.UnionType

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/__init__.py

@@ -49,6 +49,8 @@ KnownModelName = Literal[
     'cohere:command-r-plus-04-2024',
     'cohere:command-r-plus-08-2024',
     'cohere:command-r7b-12-2024',
+    'deepseek:deepseek-chat',
+    'deepseek:deepseek-reasoner',
     'google-gla:gemini-1.0-pro',
     'google-gla:gemini-1.5-flash',
     'google-gla:gemini-1.5-flash-8b',
@@ -320,54 +322,52 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .test import TestModel
 
         return TestModel()
-    elif model.startswith('cohere:'):
-        from .cohere import CohereModel
 
-        return CohereModel(model[7:])
-    elif model.startswith('openai:'):
-        from .openai import OpenAIModel
+    try:
+        provider, model_name = model.split(':')
+    except ValueError:
+        model_name = model
+        # TODO(Marcelo): We should deprecate this way.
+        if model_name.startswith(('gpt', 'o1', 'o3')):
+            provider = 'openai'
+        elif model_name.startswith('claude'):
+            provider = 'anthropic'
+        elif model_name.startswith('gemini'):
+            provider = 'google-gla'
+        else:
+            raise UserError(f'Unknown model: {model}')
+
+    if provider == 'vertexai':
+        provider = 'google-vertex'
+
+    if provider == 'cohere':
+        from .cohere import CohereModel
 
-        return OpenAIModel(model[7:])
-    elif model.startswith(('gpt', 'o1', 'o3')):
+        # TODO(Marcelo): Missing provider API.
+        return CohereModel(model_name)
+    elif provider in ('deepseek', 'openai'):
         from .openai import OpenAIModel
 
-        return OpenAIModel(model)
-    elif model.startswith('google-gla'):
-        from .gemini import GeminiModel
-
-        return GeminiModel(model[11:])
-    # backwards compatibility with old model names (ex, gemini-1.5-flash -> google-gla:gemini-1.5-flash)
-    elif model.startswith('gemini'):
+        return OpenAIModel(model_name, provider=provider)
+    elif provider in ('google-gla', 'google-vertex'):
         from .gemini import GeminiModel
 
-        # noinspection PyTypeChecker
-        return GeminiModel(model)
-    elif model.startswith('groq:'):
+        return GeminiModel(model_name, provider=provider)
+    elif provider == 'groq':
         from .groq import GroqModel
 
-        return GroqModel(model[5:])
-    elif model.startswith('google-vertex'):
-        from .vertexai import VertexAIModel
-
-        return VertexAIModel(model[14:])
-    # backwards compatibility with old model names (ex, vertexai:gemini-1.5-flash -> google-vertex:gemini-1.5-flash)
-    elif model.startswith('vertexai:'):
-        from .vertexai import VertexAIModel
-
-        return VertexAIModel(model[9:])
-    elif model.startswith('mistral:'):
+        # TODO(Marcelo): Missing provider API.
+        return GroqModel(model_name)
+    elif provider == 'mistral':
         from .mistral import MistralModel
 
-        return MistralModel(model[8:])
-    elif model.startswith('anthropic'):
-        from .anthropic import AnthropicModel
-
-        return AnthropicModel(model[10:])
-    # backwards compatibility with old model names (ex, claude-3-5-sonnet-latest -> anthropic:claude-3-5-sonnet-latest)
-    elif model.startswith('claude'):
+        # TODO(Marcelo): Missing provider API.
+        return MistralModel(model_name)
+    elif provider == 'anthropic':
         from .anthropic import AnthropicModel
 
-        return AnthropicModel(model)
+        # TODO(Marcelo): Missing provider API.
+        return AnthropicModel(model_name)
     else:
         raise UserError(f'Unknown model: {model}')
 

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/gemini.py

@@ -8,12 +8,14 @@ from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Protocol, Union, cast
+from typing import Annotated, Any, Literal, Protocol, Union, cast, overload
 from uuid import uuid4
 
 import pydantic
 from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse
-from typing_extensions import NotRequired, TypedDict, assert_never
+from typing_extensions import NotRequired, TypedDict, assert_never, deprecated
+
+from pydantic_ai.providers import Provider, infer_provider
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, UserError, _utils, usage
 from ..messages import (
@@ -82,17 +84,39 @@ class GeminiModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    http_client: AsyncHTTPClient = field(repr=False)
+    client: AsyncHTTPClient = field(repr=False)
 
     _model_name: GeminiModelName = field(repr=False)
+    _provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] | None = field(repr=False)
     _auth: AuthProtocol | None = field(repr=False)
     _url: str | None = field(repr=False)
     _system: str | None = field(default='google-gla', repr=False)
 
+    @overload
+    def __init__(
+        self,
+        model_name: GeminiModelName,
+        *,
+        provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] = 'google-gla',
+    ) -> None: ...
+
+    @deprecated('Use the `provider` argument instead of the `api_key`, `http_client`, and `url_template` arguments.')
+    @overload
     def __init__(
         self,
         model_name: GeminiModelName,
         *,
+        provider: None = None,
+        api_key: str | None = None,
+        http_client: AsyncHTTPClient | None = None,
+        url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:',
+    ) -> None: ...
+
+    def __init__(
+        self,
+        model_name: GeminiModelName,
+        *,
+        provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] | None = None,
         api_key: str | None = None,
         http_client: AsyncHTTPClient | None = None,
         url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:',
@@ -101,6 +125,7 @@ class GeminiModel(Model):
 
         Args:
             model_name: The name of the model to use.
+            provider: The provider to use for the model.
             api_key: The API key to use for authentication, if not provided, the `GEMINI_API_KEY` environment variable
                 will be used if available.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
@@ -109,14 +134,24 @@ class GeminiModel(Model):
                 `model` is substituted with the model name, and `function` is added to the end of the URL.
         """
         self._model_name = model_name
-        if api_key is None:
-            if env_api_key := os.getenv('GEMINI_API_KEY'):
-                api_key = env_api_key
+        self._provider = provider
+
+        if provider is not None:
+            if isinstance(provider, str):
+                self._system = provider
+                self.client = infer_provider(provider).client
             else:
-                raise UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
-        self.http_client = http_client or cached_async_http_client()
-        self._auth = ApiKeyAuth(api_key)
-        self._url = url_template.format(model=model_name)
+                self._system = provider.name
+                self.client = provider.client
+        else:
+            if api_key is None:
+                if env_api_key := os.getenv('GEMINI_API_KEY'):
+                    api_key = env_api_key
+                else:
+                    raise UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
+            self.client = http_client or cached_async_http_client()
+            self._auth = ApiKeyAuth(api_key)
+            self._url = url_template.format(model=model_name)
 
     @property
     def auth(self) -> AuthProtocol:
@@ -217,17 +252,19 @@ class GeminiModel(Model):
         if generation_config:
             request_data['generation_config'] = generation_config
 
-        url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
-
         headers = {
             'Content-Type': 'application/json',
             'User-Agent': get_user_agent(),
-            **await self.auth.headers(),
         }
+        if self._provider is None:  # pragma: no cover
+            url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
+            headers.update(await self.auth.headers())
+        else:
+            url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'
 
         request_json = _gemini_request_ta.dump_json(request_data, by_alias=True)
 
-        async with self.http_client.stream(
+        async with self.client.stream(
             'POST',
             url,
             content=request_json,

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/openai.py

@@ -9,7 +9,9 @@ from datetime import datetime, timezone
 from typing import Literal, Union, cast, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
+
+from pydantic_ai.providers import Provider, infer_provider
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
@@ -98,10 +100,36 @@ class OpenAIModel(Model):
     _model_name: OpenAIModelName = field(repr=False)
     _system: str | None = field(repr=False)
 
+    @overload
     def __init__(
         self,
         model_name: OpenAIModelName,
         *,
+        provider: Literal['openai', 'deepseek'] | Provider[AsyncOpenAI] = 'openai',
+        system_prompt_role: OpenAISystemPromptRole | None = None,
+        system: str | None = 'openai',
+    ) -> None: ...
+
+    @deprecated('Use the `provider` parameter instead of `base_url`, `api_key`, `openai_client` and `http_client`.')
+    @overload
+    def __init__(
+        self,
+        model_name: OpenAIModelName,
+        *,
+        provider: None = None,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+        system_prompt_role: OpenAISystemPromptRole | None = None,
+        system: str | None = 'openai',
+    ) -> None: ...
+
+    def __init__(
+        self,
+        model_name: OpenAIModelName,
+        *,
+        provider: Literal['openai', 'deepseek'] | Provider[AsyncOpenAI] | None = None,
         base_url: str | None = None,
         api_key: str | None = None,
         openai_client: AsyncOpenAI | None = None,
@@ -115,6 +143,7 @@ class OpenAIModel(Model):
             model_name: The name of the OpenAI model to use. List of model names available
                 [here](https://github.com/openai/openai-python/blob/v1.54.3/src/openai/types/chat_model.py#L7)
                 (Unfortunately, despite being ask to do so, OpenAI do not provide `.inv` files for their API).
+            provider: The provider to use. Defaults to `'openai'`.
             base_url: The base url for the OpenAI requests. If not provided, the `OPENAI_BASE_URL` environment variable
                 will be used if available. Otherwise, defaults to OpenAI's base url.
             api_key: The API key to use for authentication, if not provided, the `OPENAI_API_KEY` environment variable
@@ -129,20 +158,32 @@ class OpenAIModel(Model):
                 customize the `base_url` and `api_key` to use a different provider.
         """
         self._model_name = model_name
-        # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
-        # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.
-        if api_key is None and 'OPENAI_API_KEY' not in os.environ and base_url is not None and openai_client is None:
-            api_key = 'api-key-not-set'
-
-        if openai_client is not None:
-            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
-            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
-            assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
-            self.client = openai_client
-        elif http_client is not None:
-            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
-        else:
-            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
+
+        if provider is not None:
+            if isinstance(provider, str):
+                self.client = infer_provider(provider).client
+            else:
+                self.client = provider.client
+        else:  # pragma: no cover
+            # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
+            # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.
+            if (
+                api_key is None
+                and 'OPENAI_API_KEY' not in os.environ
+                and base_url is not None
+                and openai_client is None
+            ):
+                api_key = 'api-key-not-set'
+
+            if openai_client is not None:
+                assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+                assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
+                assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
+                self.client = openai_client
+            elif http_client is not None:
+                self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
+            else:
+                self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
         self.system_prompt_role = system_prompt_role
         self._system = system
 

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/vertexai.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import warnings
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
@@ -8,6 +9,7 @@ from pathlib import Path
 from typing import Literal
 
 from httpx import AsyncClient as AsyncHTTPClient
+from typing_extensions import deprecated
 
 from .. import usage
 from .._utils import run_in_executor
@@ -55,6 +57,7 @@ The template is used thus:
 """
 
 
+@deprecated('Please use `GeminiModel(provider=GoogleVertexProvider(...))` instead.')
 @dataclass(init=False)
 class VertexAIModel(GeminiModel):
     """A model that uses Gemini via the `*-aiplatform.googleapis.com` VertexAI API."""
@@ -103,11 +106,16 @@ class VertexAIModel(GeminiModel):
         self.project_id = project_id
         self.region = region
         self.model_publisher = model_publisher
-        self.http_client = http_client or cached_async_http_client()
+        self.client = http_client or cached_async_http_client()
         self.url_template = url_template
 
         self._auth = None
         self._url = None
+        warnings.warn(
+            'VertexAIModel is deprecated, please use `GeminiModel(provider=GoogleVertexProvider(...))` instead.',
+            DeprecationWarning,
+        )
+        self._provider = None
 
     async def ainit(self) -> None:
         """Initialize the model, setting the URL and auth.

pydantic_ai_slim-0.0.33/pydantic_ai/providers/__init__.py

@@ -0,0 +1,64 @@
+"""Providers for the API clients.
+
+The providers are in charge of providing an authenticated client to the API.
+"""
+
+from __future__ import annotations as _annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, TypeVar
+
+InterfaceClient = TypeVar('InterfaceClient')
+
+
+class Provider(ABC, Generic[InterfaceClient]):
+    """Abstract class for a provider.
+
+    The provider is in charge of providing an authenticated client to the API.
+
+    Each provider only supports a specific interface. A interface can be supported by multiple providers.
+
+    For example, the OpenAIModel interface can be supported by the OpenAIProvider and the DeepSeekProvider.
+    """
+
+    _client: InterfaceClient
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """The provider name."""
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def base_url(self) -> str:
+        """The base URL for the provider API."""
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def client(self) -> InterfaceClient:
+        """The client for the provider."""
+        raise NotImplementedError()
+
+
+def infer_provider(provider: str) -> Provider[Any]:
+    """Infer the provider from the provider name."""
+    if provider == 'openai':
+        from .openai import OpenAIProvider
+
+        return OpenAIProvider()
+    elif provider == 'deepseek':
+        from .deepseek import DeepSeekProvider
+
+        return DeepSeekProvider()
+    elif provider == 'google-vertex':
+        from .google_vertex import GoogleVertexProvider
+
+        return GoogleVertexProvider()
+    elif provider == 'google-gla':
+        from .google_gla import GoogleGLAProvider
+
+        return GoogleGLAProvider()
+    else:  # pragma: no cover
+        raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim-0.0.33/pydantic_ai/providers/deepseek.py

@@ -0,0 +1,68 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.models import cached_async_http_client
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install `openai` to use the DeepSeek provider, '
+        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+    ) from _import_error
+
+from . import Provider
+
+
+class DeepSeekProvider(Provider[AsyncOpenAI]):
+    """Provider for DeepSeek API."""
+
+    @property
+    def name(self) -> str:
+        return 'deepseek'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://api.deepseek.com'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('DEEPSEEK_API_KEY')
+        if api_key is None and openai_client is None:
+            raise ValueError(
+                'Set the `DEEPSEEK_API_KEY` environment variable or pass it via `DeepSeekProvider(api_key=...)`'
+                'to use the DeepSeek provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=cached_async_http_client())

pydantic_ai_slim-0.0.33/pydantic_ai/providers/google_gla.py

@@ -0,0 +1,44 @@
+from __future__ import annotations as _annotations
+
+import os
+
+import httpx
+
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.providers import Provider
+
+
+class GoogleGLAProvider(Provider[httpx.AsyncClient]):
+    """Provider for Google Generative Language AI API."""
+
+    @property
+    def name(self):
+        return 'google-gla'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://generativelanguage.googleapis.com/v1beta/models/'
+
+    @property
+    def client(self) -> httpx.AsyncClient:
+        return self._client
+
+    def __init__(self, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None:
+        """Create a new Google GLA provider.
+
+        Args:
+            api_key: The API key to use for authentication, if not provided, the `GEMINI_API_KEY` environment variable
+                will be used if available.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        api_key = api_key or os.environ.get('GEMINI_API_KEY')
+        if api_key is None:
+            raise ValueError(
+                'Set the `GEMINI_API_KEY` environment variable or pass it via `GoogleGLAProvider(api_key=...)`'
+                'to use the Google GLA provider.'
+            )
+
+        self._client = http_client or cached_async_http_client()
+        self._client.base_url = self.base_url
+        # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
+        self._client.headers['X-Goog-Api-Key'] = api_key

pydantic_ai_slim-0.0.33/pydantic_ai/providers/google_vertex.py

@@ -0,0 +1,200 @@
+from __future__ import annotations as _annotations
+
+import functools
+from collections.abc import AsyncGenerator
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Literal
+
+import anyio.to_thread
+import httpx
+
+from pydantic_ai.exceptions import UserError
+
+from ..models import cached_async_http_client
+from . import Provider
+
+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 _import_error:
+    raise ImportError(
+        'Please install `google-auth` to use the Google Vertex AI provider, '
+        "you can use the `vertexai` optional group — `pip install 'pydantic-ai-slim[vertexai]'`"
+    ) from _import_error
+
+
+__all__ = ('GoogleVertexProvider',)
+
+# default expiry is 3600 seconds
+MAX_TOKEN_AGE = timedelta(seconds=3000)
+
+
+class GoogleVertexProvider(Provider[httpx.AsyncClient]):
+    """Provider for Vertex AI API."""
+
+    @property
+    def name(self) -> str:
+        return 'google-vertex'
+
+    @property
+    def base_url(self) -> str:
+        return (
+            f'https://{self.region}-aiplatform.googleapis.com/v1'
+            f'/projects/{self.project_id}'
+            f'/locations/{self.region}'
+            f'/publishers/{self.model_publisher}/models/'
+        )
+
+    @property
+    def client(self) -> httpx.AsyncClient:
+        return self._client
+
+    def __init__(
+        self,
+        service_account_file: Path | str | None = None,
+        project_id: str | None = None,
+        region: VertexAiRegion = 'us-central1',
+        model_publisher: str = 'google',
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        """Create a new Vertex AI provider.
+
+        Args:
+            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.
+        """
+        self._client = http_client or cached_async_http_client()
+        self.service_account_file = service_account_file
+        self.project_id = project_id
+        self.region = region
+        self.model_publisher = model_publisher
+
+        self._client.auth = _VertexAIAuth(service_account_file, project_id, region)
+        self._client.base_url = self.base_url
+
+
+class _VertexAIAuth(httpx.Auth):
+    """Auth class for Vertex AI API."""
+
+    credentials: BaseCredentials | ServiceAccountCredentials | None
+
+    def __init__(
+        self,
+        service_account_file: Path | str | None = None,
+        project_id: str | None = None,
+        region: VertexAiRegion = 'us-central1',
+    ) -> None:
+        self.service_account_file = service_account_file
+        self.project_id = project_id
+        self.region = region
+
+        self.credentials = None
+        self.token_created: datetime | None = None
+
+    async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        if self.credentials is None:
+            self.credentials = await self._get_credentials()
+        if self.credentials.token is None or self._token_expired():  # type: ignore[reportUnknownMemberType]
+            await anyio.to_thread.run_sync(self._refresh_token)
+            self.token_created = datetime.now()
+        request.headers['Authorization'] = f'Bearer {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
+
+        # NOTE: This workaround is in place because we might get the project_id from the credentials.
+        request.url = httpx.URL(str(request.url).replace('projects/None', f'projects/{self.project_id}'))
+        yield request
+
+    async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials:
+        if self.service_account_file is not None:
+            creds = await _creds_from_file(self.service_account_file)
+            assert creds.project_id is None or isinstance(creds.project_id, str)  # type: ignore[reportUnknownMemberType]
+            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}')
+            self.project_id = creds_project_id
+        return creds
+
+    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:  # pragma: no cover
+        assert self.credentials is not None
+        self.credentials.refresh(Request())  # type: ignore[reportUnknownMemberType]
+        assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
+        return self.credentials.token
+
+
+async def _async_google_auth() -> tuple[BaseCredentials, str | None]:
+    return await anyio.to_thread.run_sync(google.auth.default, ['https://www.googleapis.com/auth/cloud-platform'])  # type: ignore
+
+
+async def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredentials:
+    service_account_credentials_from_file = functools.partial(
+        ServiceAccountCredentials.from_service_account_file,  # type: ignore[reportUnknownMemberType]
+        scopes=['https://www.googleapis.com/auth/cloud-platform'],
+    )
+    return await anyio.to_thread.run_sync(service_account_credentials_from_file, str(service_account_file))
+
+
+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).
+"""

pydantic_ai_slim-0.0.33/pydantic_ai/providers/openai.py

@@ -0,0 +1,72 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import TypeVar
+
+import httpx
+
+from pydantic_ai.models import cached_async_http_client
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install `openai` to use the OpenAI provider, '
+        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+    ) from _import_error
+
+
+from . import Provider
+
+InterfaceClient = TypeVar('InterfaceClient')
+
+
+class OpenAIProvider(Provider[AsyncOpenAI]):
+    """Provider for OpenAI API."""
+
+    @property
+    def name(self) -> str:
+        return 'openai'  # pragma: no cover
+
+    @property
+    def base_url(self) -> str:
+        return self._base_url
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def __init__(
+        self,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        """Create a new OpenAI provider.
+
+        Args:
+            base_url: The base url for the OpenAI requests. If not provided, the `OPENAI_BASE_URL` environment variable
+                will be used if available. Otherwise, defaults to OpenAI's base url.
+            api_key: The API key to use for authentication, if not provided, the `OPENAI_API_KEY` environment variable
+                will be used if available.
+            openai_client: An existing
+                [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
+                client to use. If provided, `base_url`, `api_key`, and `http_client` must be `None`.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        self._base_url = base_url or 'https://api.openai.com/v1'
+        # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
+        # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.
+        if api_key is None and 'OPENAI_API_KEY' not in os.environ and openai_client is None:
+            api_key = 'api-key-not-set'
+
+        if openai_client is not None:
+            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
+            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=cached_async_http_client())

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.33}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.32"
+version = "0.0.33"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
 license = "MIT"
@@ -36,9 +36,10 @@ dependencies = [
     "griffe>=1.3.2",
     "httpx>=0.27",
     "pydantic>=2.10",
-    "pydantic-graph==0.0.32",
+    "pydantic-graph==0.0.33",
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
+    "typing-inspection>=0.4.0",
 ]
 
 [project.optional-dependencies]