Add vertex AI (#85)

Author: samuelcolvin

.github/workflows/ci.yml

@@ -10,7 +10,7 @@ on:
 
 env:
   CI: true
-  RICH_COLUMNS: 120
+  COLUMNS: 120
 
 permissions:
   contents: read
@@ -80,11 +80,12 @@ jobs:
         with:
           enable-cache: true
 
-      - run: uv run --python 3.12 --frozen pytest tests/test_live.py -v --durations=100
+      - run: uv run --python 3.12 --frozen --extra vertexai pytest tests/test_live.py -v --durations=100
         env:
           PYDANTIC_AI_LIVE_TEST_DANGEROUS: 'CHARGE-ME!'
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
+          GOOGLE_SERVICE_ACCOUNT_CONTENT: ${{ secrets.GOOGLE_SERVICE_ACCOUNT_CONTENT }}
 
   test:
     name: test on ${{ matrix.python-version }}

docs/api/models/vertexai.md

@@ -0,0 +1,3 @@
+# `pydantic_ai.models.vertexai`
+
+::: pydantic_ai.models.vertexai

mkdocs.yml

@@ -38,6 +38,7 @@ nav:
     - api/models/base.md
     - api/models/openai.md
     - api/models/gemini.md
+    - api/models/vertexai.md
     - api/models/test.md
     - api/models/function.md
 

pydantic_ai/agent.py

@@ -601,7 +601,8 @@ async def _get_agent_model(
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
         result_tools = list(self._result_schema.tools.values()) if self._result_schema else None
-        return model_, custom_model, model_.agent_model(self._retrievers, self._allow_text_result, result_tools)
+        agent_model = await model_.agent_model(self._retrievers, self._allow_text_result, result_tools)
+        return model_, custom_model, agent_model
 
     async def _prepare_messages(
         self, deps: AgentDeps, user_prompt: str, message_history: list[_messages.Message] | None

pydantic_ai/models/__init__.py

@@ -44,14 +44,16 @@ class Model(ABC):
     """Abstract class for a model."""
 
     @abstractmethod
-    def agent_model(
+    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.

pydantic_ai/models/function.py

@@ -65,7 +65,7 @@ def __init__(self, function: FunctionDef | None = None, *, stream_function: Stre
         self.function = function
         self.stream_function = stream_function
 
-    def agent_model(
+    async def agent_model(
         self,
         retrievers: Mapping[str, AbstractToolDefinition],
         allow_text_result: bool,

pydantic_ai/models/gemini.py

@@ -29,7 +29,7 @@
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Protocol, Union
 
 import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Response as HTTPResponse
@@ -77,17 +77,17 @@ class GeminiModel(Model):
     """
 
     model_name: GeminiModelName
-    api_key: str
+    auth: AuthProtocol
     http_client: AsyncHTTPClient
-    url_template: str
+    url: str
 
     def __init__(
         self,
         model_name: GeminiModelName,
         *,
         api_key: str | None = None,
         http_client: AsyncHTTPClient | None = None,
-        url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:{function}',
+        url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:',
     ):
         """Initialize a Gemini model.
 
@@ -97,62 +97,94 @@ def __init__(
                 will be used if available.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
             url_template: The URL template to use for making requests, you shouldn't need to change this,
-                docs [here](https://ai.google.dev/gemini-api/docs/quickstart?lang=rest#make-first-request).
+                docs [here](https://ai.google.dev/gemini-api/docs/quickstart?lang=rest#make-first-request),
+                `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
             else:
                 raise exceptions.UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
-        self.api_key = api_key
+        self.auth = ApiKeyAuth(api_key)
         self.http_client = http_client or cached_async_http_client()
-        self.url_template = url_template
+        self.url = url_template.format(model=model_name)
 
-    def agent_model(
+    async def agent_model(
         self,
         retrievers: Mapping[str, AbstractToolDefinition],
         allow_text_result: bool,
         result_tools: Sequence[AbstractToolDefinition] | None,
     ) -> GeminiAgentModel:
-        check_allow_model_requests()
-        tools = [_function_from_abstract_tool(t) for t in retrievers.values()]
-        if result_tools is not None:
-            tools += [_function_from_abstract_tool(t) for t in result_tools]
-
-        if allow_text_result:
-            tool_config = None
-        else:
-            tool_config = _tool_config([t['name'] for t in tools])
-
         return GeminiAgentModel(
             http_client=self.http_client,
             model_name=self.model_name,
-            api_key=self.api_key,
-            tools=_GeminiTools(function_declarations=tools) if tools else None,
-            tool_config=tool_config,
-            url_template=self.url_template,
+            auth=self.auth,
+            url=self.url,
+            retrievers=retrievers,
+            allow_text_result=allow_text_result,
+            result_tools=result_tools,
         )
 
     def name(self) -> str:
         return self.model_name
 
 
+class AuthProtocol(Protocol):
+    async def headers(self) -> dict[str, str]: ...
+
+
 @dataclass
+class ApiKeyAuth:
+    api_key: str
+
+    async def headers(self) -> dict[str, str]:
+        # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
+        return {'X-Goog-Api-Key': self.api_key}
+
+
+@dataclass(init=False)
 class GeminiAgentModel(AgentModel):
     """Implementation of `AgentModel` for Gemini models."""
 
     http_client: AsyncHTTPClient
     model_name: GeminiModelName
-    api_key: str
+    auth: AuthProtocol
     tools: _GeminiTools | None
     tool_config: _GeminiToolConfig | None
-    url_template: str
+    url: str
+
+    def __init__(
+        self,
+        http_client: AsyncHTTPClient,
+        model_name: GeminiModelName,
+        auth: AuthProtocol,
+        url: str,
+        retrievers: Mapping[str, AbstractToolDefinition],
+        allow_text_result: bool,
+        result_tools: Sequence[AbstractToolDefinition] | None,
+    ):
+        check_allow_model_requests()
+        tools = [_function_from_abstract_tool(t) for t in retrievers.values()]
+        if result_tools is not None:
+            tools += [_function_from_abstract_tool(t) for t in result_tools]
+
+        if allow_text_result:
+            tool_config = None
+        else:
+            tool_config = _tool_config([t['name'] for t in tools])
+
+        self.http_client = http_client
+        self.model_name = model_name
+        self.auth = auth
+        self.tools = _GeminiTools(function_declarations=tools) if tools else None
+        self.tool_config = tool_config
+        self.url = url
 
     async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, result.Cost]:
         async with self._make_request(messages, False) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
-        return self._process_response(response), _metadata_as_cost(response['usage_metadata'])
+        return self._process_response(response), _metadata_as_cost(response)
 
     @asynccontextmanager
     async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
@@ -178,16 +210,15 @@ async def _make_request(self, messages: list[Message], streamed: bool) -> AsyncI
         if self.tool_config is not None:
             request_data['tool_config'] = self.tool_config
 
-        request_json = _gemini_request_ta.dump_json(request_data, by_alias=True)
-        # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
+        url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
+
         headers = {
-            'X-Goog-Api-Key': self.api_key,
             'Content-Type': 'application/json',
             'User-Agent': get_user_agent(),
+            **await self.auth.headers(),
         }
-        url = self.url_template.format(
-            model=self.model_name, function='streamGenerateContent' if streamed else 'generateContent'
-        )
+
+        request_json = _gemini_request_ta.dump_json(request_data, by_alias=True)
 
         async with self.http_client.stream('POST', url, content=request_json, headers=headers) as r:
             if r.status_code != 200:
@@ -283,7 +314,7 @@ def get(self, *, final: bool = False) -> Iterable[str]:
                 new_items, experimental_allow_partial='trailing-strings'
             )
         for r in new_responses:
-            self._cost += _metadata_as_cost(r['usage_metadata'])
+            self._cost += _metadata_as_cost(r)
             parts = r['candidates'][0]['content']['parts']
             if _all_text_parts(parts):
                 for part in parts:
@@ -329,7 +360,7 @@ def get(self, *, final: bool = False) -> ModelStructuredResponse:
         combined_parts: list[_GeminiFunctionCallPart] = []
         self._cost = result.Cost()
         for r in responses:
-            self._cost += _metadata_as_cost(r['usage_metadata'])
+            self._cost += _metadata_as_cost(r)
             candidate = r['candidates'][0]
             parts = candidate['content']['parts']
             if _all_function_call_parts(parts):
@@ -521,10 +552,12 @@ class _GeminiResponse(TypedDict):
     """Schema for the response from the Gemini API.
 
     See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse>
+    and <https://cloud.google.com/vertex-ai/docs/reference/rest/v1/GenerateContentResponse>
     """
 
     candidates: list[_GeminiCandidates]
-    usage_metadata: Annotated[_GeminiUsageMetaData, Field(alias='usageMetadata')]
+    # usageMetadata appears to be required by both APIs but is omitted when streaming responses until the last response
+    usage_metadata: NotRequired[Annotated[_GeminiUsageMetaData, Field(alias='usageMetadata')]]
     prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, Field(alias='promptFeedback')]]
 
 
@@ -582,7 +615,10 @@ class _GeminiUsageMetaData(TypedDict, total=False):
     cached_content_token_count: NotRequired[Annotated[int, Field(alias='cachedContentTokenCount')]]
 
 
-def _metadata_as_cost(metadata: _GeminiUsageMetaData) -> result.Cost:
+def _metadata_as_cost(response: _GeminiResponse) -> result.Cost:
+    metadata = response.get('usage_metadata')
+    if metadata is None:
+        return result.Cost()
     details: dict[str, int] = {}
     if cached_content_token_count := metadata.get('cached_content_token_count'):
         details['cached_content_token_count'] = cached_content_token_count

pydantic_ai/models/openai.py

@@ -80,7 +80,7 @@ def __init__(
         else:
             self.client = AsyncOpenAI(api_key=api_key, http_client=cached_async_http_client())
 
-    def agent_model(
+    async def agent_model(
         self,
         retrievers: Mapping[str, AbstractToolDefinition],
         allow_text_result: bool,

pydantic_ai/models/test.py

@@ -70,7 +70,7 @@ class TestModel(Model):
     agent_model_allow_text_result: bool | None = field(default=None, init=False)
     agent_model_result_tools: list[AbstractToolDefinition] | None = field(default=None, init=False)
 
-    def agent_model(
+    async def agent_model(
         self,
         retrievers: Mapping[str, AbstractToolDefinition],
         allow_text_result: bool,