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

pydantic_ai/models/function.py

@@ -4,7 +4,7 @@ import inspect
 import re
 from collections.abc import AsyncIterator, Awaitable, Iterable
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from itertools import chain
 from typing import Callable, Union, cast
@@ -13,7 +13,20 @@ import pydantic_core
 from typing_extensions import TypeAlias, assert_never, overload
 
 from .. import _utils, result
-from ..messages import ArgsJson, Message, ModelAnyResponse, ModelStructuredResponse, ToolCall
+from ..messages import (
+    ArgsJson,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import AgentModel, EitherStreamedResponse, Model, StreamStructuredResponse, StreamTextResponse
 
@@ -59,7 +72,7 @@ class FunctionModel(Model):
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         return FunctionAgentModel(
-            self.function, self.stream_function, AgentInfo(function_tools, allow_text_result, result_tools)
+            self.function, self.stream_function, AgentInfo(function_tools, allow_text_result, result_tools, None)
         )
 
     def name(self) -> str:
@@ -88,6 +101,8 @@ class AgentInfo:
     """Whether a plain text result is allowed."""
     result_tools: list[ToolDefinition]
     """The tools that can called as the final result of the run."""
+    model_settings: ModelSettings | None
+    """The model settings passed to the run call."""
 
 
 @dataclass
@@ -106,10 +121,10 @@ class DeltaToolCall:
 DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
 """A mapping of tool call IDs to incremental changes."""
 
-FunctionDef: TypeAlias = Callable[[list[Message], AgentInfo], Union[ModelAnyResponse, Awaitable[ModelAnyResponse]]]
+FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], Union[ModelResponse, Awaitable[ModelResponse]]]
 """A function used to generate a non-streamed response."""
 
-StreamFunctionDef: TypeAlias = Callable[[list[Message], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls]]]
+StreamFunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls]]]
 """A function used to generate a streamed response.
 
 While this is defined as having return type of `AsyncIterator[Union[str, DeltaToolCalls]]`, it should
@@ -127,18 +142,25 @@ class FunctionAgentModel(AgentModel):
     stream_function: StreamFunctionDef | None
     agent_info: AgentInfo
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, result.Cost]:
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Cost]:
+        agent_info = replace(self.agent_info, model_settings=model_settings)
+
         assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'
         if inspect.iscoroutinefunction(self.function):
-            response = await self.function(messages, self.agent_info)
+            response = await self.function(messages, agent_info)
         else:
-            response_ = await _utils.run_in_executor(self.function, messages, self.agent_info)
-            response = cast(ModelAnyResponse, response_)
+            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
+            assert isinstance(response_, ModelResponse), response_
+            response = response_
         # TODO is `messages` right here? Should it just be new messages?
         return response, _estimate_cost(chain(messages, [response]))
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
         assert (
             self.stream_function is not None
         ), 'FunctionModel must receive a `stream_function` to support streamed requests'
@@ -206,13 +228,13 @@ class FunctionStreamStructuredResponse(StreamStructuredResponse):
             else:
                 self._delta_tool_calls[key] = new
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        calls: list[ToolCall] = []
+    def get(self, *, final: bool = False) -> ModelResponse:
+        calls: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if c.name is not None and c.json_args is not None:
-                calls.append(ToolCall.from_json(c.name, c.json_args))
+                calls.append(ToolCallPart.from_json(c.name, c.json_args))
 
-        return ModelStructuredResponse(calls, timestamp=self._timestamp)
+        return ModelResponse(calls, timestamp=self._timestamp)
 
     def cost(self) -> result.Cost:
         return result.Cost()
@@ -221,32 +243,38 @@ class FunctionStreamStructuredResponse(StreamStructuredResponse):
         return self._timestamp
 
 
-def _estimate_cost(messages: Iterable[Message]) -> result.Cost:
+def _estimate_cost(messages: Iterable[ModelMessage]) -> result.Cost:
     """Very rough guesstimate of the number of tokens associate with a series of messages.
 
     This is designed to be used solely to give plausible numbers for testing!
     """
     # there seem to be about 50 tokens of overhead for both Gemini and OpenAI calls, so add that here ¯\_(ツ)_/¯
-
     request_tokens = 50
     response_tokens = 0
     for message in messages:
-        if message.role == 'system' or message.role == 'user':
-            request_tokens += _string_cost(message.content)
-        elif message.role == 'tool-return':
-            request_tokens += _string_cost(message.model_response_str())
-        elif message.role == 'retry-prompt':
-            request_tokens += _string_cost(message.model_response())
-        elif message.role == 'model-text-response':
-            response_tokens += _string_cost(message.content)
-        elif message.role == 'model-structured-response':
-            for call in message.calls:
-                if isinstance(call.args, ArgsJson):
-                    args_str = call.args.args_json
+        if isinstance(message, ModelRequest):
+            for part in message.parts:
+                if isinstance(part, (SystemPromptPart, UserPromptPart)):
+                    request_tokens += _string_cost(part.content)
+                elif isinstance(part, ToolReturnPart):
+                    request_tokens += _string_cost(part.model_response_str())
+                elif isinstance(part, RetryPromptPart):
+                    request_tokens += _string_cost(part.model_response())
                 else:
-                    args_str = pydantic_core.to_json(call.args.args_dict).decode()
-
-                response_tokens += 1 + _string_cost(args_str)
+                    assert_never(part)
+        elif isinstance(message, ModelResponse):
+            for part in message.parts:
+                if isinstance(part, TextPart):
+                    response_tokens += _string_cost(part.content)
+                elif isinstance(part, ToolCallPart):
+                    call = part
+                    if isinstance(call.args, ArgsJson):
+                        args_str = call.args.args_json
+                    else:
+                        args_str = pydantic_core.to_json(call.args.args_dict).decode()
+                    response_tokens += 1 + _string_cost(args_str)
+                else:
+                    assert_never(part)
         else:
             assert_never(message)
     return result.Cost(

pydantic_ai/models/gemini.py

@@ -2,29 +2,33 @@ from __future__ import annotations as _annotations
 
 import os
 import re
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator, Iterable, Sequence
 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
 
+import pydantic
 import pydantic_core
-from httpx import AsyncClient as AsyncHTTPClient, Response as HTTPResponse
-from pydantic import Discriminator, Field, Tag
+from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse
 from typing_extensions import NotRequired, TypedDict, TypeGuard, assert_never
 
-from .. import UnexpectedModelBehavior, _pydantic, _utils, exceptions, result
+from .. import UnexpectedModelBehavior, _utils, exceptions, result
 from ..messages import (
     ArgsDict,
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
 )
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
@@ -37,7 +41,9 @@ from . import (
     get_user_agent,
 )
 
-GeminiModelName = Literal['gemini-1.5-flash', 'gemini-1.5-flash-8b', 'gemini-1.5-pro', 'gemini-1.0-pro']
+GeminiModelName = Literal[
+    'gemini-1.5-flash', 'gemini-1.5-flash-8b', 'gemini-1.5-pro', 'gemini-1.0-pro', 'gemini-2.0-flash-exp'
+]
 """Named Gemini models.
 
 See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list.
@@ -164,26 +170,25 @@ class GeminiAgentModel(AgentModel):
         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:
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Cost]:
+        async with self._make_request(messages, False, model_settings) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
         return self._process_response(response), _metadata_as_cost(response)
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        async with self._make_request(messages, True) as http_response:
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        async with self._make_request(messages, True, model_settings) as http_response:
             yield await self._process_streamed_response(http_response)
 
     @asynccontextmanager
-    async def _make_request(self, messages: list[Message], streamed: bool) -> AsyncIterator[HTTPResponse]:
-        contents: list[_GeminiContent] = []
-        sys_prompt_parts: list[_GeminiTextPart] = []
-        for m in messages:
-            either_content = self._message_to_gemini(m)
-            if left := either_content.left:
-                sys_prompt_parts.append(left.value)
-            else:
-                contents.append(either_content.right)
+    async def _make_request(
+        self, messages: list[ModelMessage], streamed: bool, model_settings: ModelSettings | None
+    ) -> AsyncIterator[HTTPResponse]:
+        sys_prompt_parts, contents = self._message_to_gemini_content(messages)
 
         request_data = _GeminiRequest(contents=contents)
         if sys_prompt_parts:
@@ -193,6 +198,17 @@ class GeminiAgentModel(AgentModel):
         if self.tool_config is not None:
             request_data['tool_config'] = self.tool_config
 
+        generation_config: _GeminiGenerationConfig = {}
+        if model_settings:
+            if (max_tokens := model_settings.get('max_tokens')) is not None:
+                generation_config['max_output_tokens'] = max_tokens
+            if (temperature := model_settings.get('temperature')) is not None:
+                generation_config['temperature'] = temperature
+            if (top_p := model_settings.get('top_p')) is not None:
+                generation_config['top_p'] = top_p
+        if generation_config:
+            request_data['generation_config'] = generation_config
+
         url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
 
         headers = {
@@ -203,19 +219,24 @@ class GeminiAgentModel(AgentModel):
 
         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:
+        async with self.http_client.stream(
+            'POST',
+            url,
+            content=request_json,
+            headers=headers,
+            timeout=(model_settings or {}).get('timeout', USE_CLIENT_DEFAULT),
+        ) as r:
             if r.status_code != 200:
                 await r.aread()
                 raise exceptions.UnexpectedModelBehavior(f'Unexpected response from gemini {r.status_code}', r.text)
             yield r
 
     @staticmethod
-    def _process_response(response: _GeminiResponse) -> ModelAnyResponse:
-        either = _extract_response_parts(response)
-        if left := either.left:
-            return _structured_response_from_parts(left.value)
-        else:
-            return ModelTextResponse(content=''.join(part['text'] for part in either.right))
+    def _process_response(response: _GeminiResponse) -> ModelResponse:
+        if len(response['candidates']) != 1:
+            raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')
+        parts = response['candidates'][0]['content']['parts']
+        return _process_response_from_parts(parts)
 
     @staticmethod
     async def _process_streamed_response(http_response: HTTPResponse) -> EitherStreamedResponse:
@@ -239,34 +260,37 @@ class GeminiAgentModel(AgentModel):
         if start_response is None:
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
+        # TODO: Update this once we rework stream responses to be more flexible
         if _extract_response_parts(start_response).is_left():
             return GeminiStreamStructuredResponse(_content=content, _stream=aiter_bytes)
         else:
             return GeminiStreamTextResponse(_json_content=content, _stream=aiter_bytes)
 
-    @staticmethod
-    def _message_to_gemini(m: Message) -> _utils.Either[_GeminiTextPart, _GeminiContent]:
-        """Convert a message to a _GeminiTextPart for "system_instructions" or _GeminiContent for "contents"."""
-        if m.role == 'system':
-            # SystemPrompt ->
-            return _utils.Either(left=_GeminiTextPart(text=m.content))
-        elif m.role == 'user':
-            # UserPrompt ->
-            return _utils.Either(right=_content_user_text(m.content))
-        elif m.role == 'tool-return':
-            # ToolReturn ->
-            return _utils.Either(right=_content_function_return(m))
-        elif m.role == 'retry-prompt':
-            # RetryPrompt ->
-            return _utils.Either(right=_content_function_retry(m))
-        elif m.role == 'model-text-response':
-            # ModelTextResponse ->
-            return _utils.Either(right=_content_model_text(m.content))
-        elif m.role == 'model-structured-response':
-            # ModelStructuredResponse ->
-            return _utils.Either(right=_content_function_call(m))
-        else:
-            assert_never(m)
+    @classmethod
+    def _message_to_gemini_content(
+        cls, messages: list[ModelMessage]
+    ) -> tuple[list[_GeminiTextPart], list[_GeminiContent]]:
+        sys_prompt_parts: list[_GeminiTextPart] = []
+        contents: list[_GeminiContent] = []
+        for m in messages:
+            if isinstance(m, ModelRequest):
+                for part in m.parts:
+                    if isinstance(part, SystemPromptPart):
+                        sys_prompt_parts.append(_GeminiTextPart(text=part.content))
+                    elif isinstance(part, UserPromptPart):
+                        contents.append(_content_user_prompt(part))
+                    elif isinstance(part, ToolReturnPart):
+                        contents.append(_content_tool_return(part))
+                    elif isinstance(part, RetryPromptPart):
+                        contents.append(_content_retry_prompt(part))
+                    else:
+                        assert_never(part)
+            elif isinstance(m, ModelResponse):
+                contents.append(_content_model_response(m))
+            else:
+                assert_never(m)
+
+        return sys_prompt_parts, contents
 
 
 @dataclass
@@ -327,8 +351,8 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
         chunk = await self._stream.__anext__()
         self._content.extend(chunk)
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        """Get the `ModelStructuredResponse` at this point.
+    def get(self, *, final: bool = False) -> ModelResponse:
+        """Get the `ModelResponse` at this point.
 
         NOTE: It's not clear how the stream of responses should be combined because Gemini seems to always
         reply with a single response, when returning a structured data.
@@ -340,20 +364,13 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
             self._content,
             experimental_allow_partial='off' if final else 'trailing-strings',
         )
-        combined_parts: list[_GeminiFunctionCallPart] = []
+        combined_parts: list[_GeminiPartUnion] = []
         self._cost = result.Cost()
         for r in responses:
             self._cost += _metadata_as_cost(r)
             candidate = r['candidates'][0]
-            parts = candidate['content']['parts']
-            if _all_function_call_parts(parts):
-                combined_parts.extend(parts)
-            elif not candidate.get('finish_reason'):
-                # you can get an empty text part along with the finish_reason, so we ignore that case
-                raise UnexpectedModelBehavior(
-                    'Streamed response with unexpected content, expected all parts to be function calls'
-                )
-        return _structured_response_from_parts(combined_parts, timestamp=self._timestamp)
+            combined_parts.extend(candidate['content']['parts'])
+        return _process_response_from_parts(combined_parts, timestamp=self._timestamp)
 
     def cost(self) -> result.Cost:
         return self._cost
@@ -367,6 +384,7 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
 # TypeAdapters take care of validation and serialization
 
 
+@pydantic.with_config(pydantic.ConfigDict(defer_build=True))
 class _GeminiRequest(TypedDict):
     """Schema for an API request to the Gemini API.
 
@@ -382,32 +400,37 @@ class _GeminiRequest(TypedDict):
     Developer generated system instructions, see
     <https://ai.google.dev/gemini-api/docs/system-instructions?lang=rest>
     """
+    generation_config: NotRequired[_GeminiGenerationConfig]
 
 
-class _GeminiContent(TypedDict):
-    role: Literal['user', 'model']
-    parts: list[_GeminiPartUnion]
+class _GeminiGenerationConfig(TypedDict, total=False):
+    """Schema for an API request to the Gemini API.
 
+    Note there are many additional fields available that have not been added yet.
 
-def _content_user_text(text: str) -> _GeminiContent:
-    return _GeminiContent(role='user', parts=[_GeminiTextPart(text=text)])
+    See <https://ai.google.dev/api/generate-content#generationconfig> for API docs.
+    """
+
+    max_output_tokens: int
+    temperature: float
+    top_p: float
 
 
-def _content_model_text(text: str) -> _GeminiContent:
-    return _GeminiContent(role='model', parts=[_GeminiTextPart(text=text)])
+class _GeminiContent(TypedDict):
+    role: Literal['user', 'model']
+    parts: list[_GeminiPartUnion]
 
 
-def _content_function_call(m: ModelStructuredResponse) -> _GeminiContent:
-    parts: list[_GeminiPartUnion] = [_function_call_part_from_call(t) for t in m.calls]
-    return _GeminiContent(role='model', parts=parts)
+def _content_user_prompt(m: UserPromptPart) -> _GeminiContent:
+    return _GeminiContent(role='user', parts=[_GeminiTextPart(text=m.content)])
 
 
-def _content_function_return(m: ToolReturn) -> _GeminiContent:
+def _content_tool_return(m: ToolReturnPart) -> _GeminiContent:
     f_response = _response_part_from_response(m.tool_name, m.model_response_object())
     return _GeminiContent(role='user', parts=[f_response])
 
 
-def _content_function_retry(m: RetryPrompt) -> _GeminiContent:
+def _content_retry_prompt(m: RetryPromptPart) -> _GeminiContent:
     if m.tool_name is None:
         part = _GeminiTextPart(text=m.model_response())
     else:
@@ -416,26 +439,43 @@ def _content_function_retry(m: RetryPrompt) -> _GeminiContent:
     return _GeminiContent(role='user', parts=[part])
 
 
+def _content_model_response(m: ModelResponse) -> _GeminiContent:
+    parts: list[_GeminiPartUnion] = []
+    for item in m.parts:
+        if isinstance(item, ToolCallPart):
+            parts.append(_function_call_part_from_call(item))
+        elif isinstance(item, TextPart):
+            parts.append(_GeminiTextPart(text=item.content))
+        else:
+            assert_never(item)
+    return _GeminiContent(role='model', parts=parts)
+
+
 class _GeminiTextPart(TypedDict):
     text: str
 
 
 class _GeminiFunctionCallPart(TypedDict):
-    function_call: Annotated[_GeminiFunctionCall, Field(alias='functionCall')]
+    function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')]
 
 
-def _function_call_part_from_call(tool: ToolCall) -> _GeminiFunctionCallPart:
+def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart:
     assert isinstance(tool.args, ArgsDict), f'Expected ArgsObject, got {tool.args}'
     return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args.args_dict))
 
 
-def _structured_response_from_parts(
-    parts: list[_GeminiFunctionCallPart], timestamp: datetime | None = None
-) -> ModelStructuredResponse:
-    return ModelStructuredResponse(
-        calls=[ToolCall.from_dict(part['function_call']['name'], part['function_call']['args']) for part in parts],
-        timestamp=timestamp or _utils.now_utc(),
-    )
+def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: datetime | None = None) -> ModelResponse:
+    items: list[ModelResponsePart] = []
+    for part in parts:
+        if 'text' in part:
+            items.append(TextPart(part['text']))
+        elif 'function_call' in part:
+            items.append(ToolCallPart.from_dict(part['function_call']['name'], part['function_call']['args']))
+        elif 'function_response' in part:
+            raise exceptions.UnexpectedModelBehavior(
+                f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
+            )
+    return ModelResponse(items, timestamp=timestamp or _utils.now_utc())
 
 
 class _GeminiFunctionCall(TypedDict):
@@ -446,7 +486,7 @@ class _GeminiFunctionCall(TypedDict):
 
 
 class _GeminiFunctionResponsePart(TypedDict):
-    function_response: Annotated[_GeminiFunctionResponse, Field(alias='functionResponse')]
+    function_response: Annotated[_GeminiFunctionResponse, pydantic.Field(alias='functionResponse')]
 
 
 def _response_part_from_response(name: str, response: dict[str, Any]) -> _GeminiFunctionResponsePart:
@@ -476,11 +516,11 @@ def _part_discriminator(v: Any) -> str:
 # TODO discriminator
 _GeminiPartUnion = Annotated[
     Union[
-        Annotated[_GeminiTextPart, Tag('text')],
-        Annotated[_GeminiFunctionCallPart, Tag('function_call')],
-        Annotated[_GeminiFunctionResponsePart, Tag('function_response')],
+        Annotated[_GeminiTextPart, pydantic.Tag('text')],
+        Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')],
+        Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')],
     ],
-    Discriminator(_part_discriminator),
+    pydantic.Discriminator(_part_discriminator),
 ]
 
 
@@ -490,7 +530,7 @@ class _GeminiTextContent(TypedDict):
 
 
 class _GeminiTools(TypedDict):
-    function_declarations: list[Annotated[_GeminiFunction, Field(alias='functionDeclarations')]]
+    function_declarations: list[Annotated[_GeminiFunction, pydantic.Field(alias='functionDeclarations')]]
 
 
 class _GeminiFunction(TypedDict):
@@ -531,6 +571,7 @@ class _GeminiFunctionCallingConfig(TypedDict):
     allowed_function_names: list[str]
 
 
+@pydantic.with_config(pydantic.ConfigDict(defer_build=True))
 class _GeminiResponse(TypedDict):
     """Schema for the response from the Gemini API.
 
@@ -540,10 +581,11 @@ class _GeminiResponse(TypedDict):
 
     candidates: list[_GeminiCandidates]
     # 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')]]
+    usage_metadata: NotRequired[Annotated[_GeminiUsageMetaData, pydantic.Field(alias='usageMetadata')]]
+    prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, pydantic.Field(alias='promptFeedback')]]
 
 
+# TODO: Delete the next three functions once we've reworked streams to be more flexible
 def _extract_response_parts(
     response: _GeminiResponse,
 ) -> _utils.Either[list[_GeminiFunctionCallPart], list[_GeminiTextPart]]:
@@ -576,14 +618,14 @@ class _GeminiCandidates(TypedDict):
     """See <https://ai.google.dev/api/generate-content#v1beta.Candidate>."""
 
     content: _GeminiContent
-    finish_reason: NotRequired[Annotated[Literal['STOP'], Field(alias='finishReason')]]
+    finish_reason: NotRequired[Annotated[Literal['STOP', 'MAX_TOKENS'], pydantic.Field(alias='finishReason')]]
     """
     See <https://ai.google.dev/api/generate-content#FinishReason>, lots of other values are possible,
     but let's wait until we see them and know what they mean to add them here.
     """
-    avg_log_probs: NotRequired[Annotated[float, Field(alias='avgLogProbs')]]
+    avg_log_probs: NotRequired[Annotated[float, pydantic.Field(alias='avgLogProbs')]]
     index: NotRequired[int]
-    safety_ratings: NotRequired[Annotated[list[_GeminiSafetyRating], Field(alias='safetyRatings')]]
+    safety_ratings: NotRequired[Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')]]
 
 
 class _GeminiUsageMetaData(TypedDict, total=False):
@@ -592,10 +634,10 @@ class _GeminiUsageMetaData(TypedDict, total=False):
     The docs suggest all fields are required, but some are actually not required, so we assume they are all optional.
     """
 
-    prompt_token_count: Annotated[int, Field(alias='promptTokenCount')]
-    candidates_token_count: NotRequired[Annotated[int, Field(alias='candidatesTokenCount')]]
-    total_token_count: Annotated[int, Field(alias='totalTokenCount')]
-    cached_content_token_count: NotRequired[Annotated[int, Field(alias='cachedContentTokenCount')]]
+    prompt_token_count: Annotated[int, pydantic.Field(alias='promptTokenCount')]
+    candidates_token_count: NotRequired[Annotated[int, pydantic.Field(alias='candidatesTokenCount')]]
+    total_token_count: Annotated[int, pydantic.Field(alias='totalTokenCount')]
+    cached_content_token_count: NotRequired[Annotated[int, pydantic.Field(alias='cachedContentTokenCount')]]
 
 
 def _metadata_as_cost(response: _GeminiResponse) -> result.Cost:
@@ -629,15 +671,15 @@ class _GeminiSafetyRating(TypedDict):
 class _GeminiPromptFeedback(TypedDict):
     """See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse>."""
 
-    block_reason: Annotated[str, Field(alias='blockReason')]
-    safety_ratings: Annotated[list[_GeminiSafetyRating], Field(alias='safetyRatings')]
+    block_reason: Annotated[str, pydantic.Field(alias='blockReason')]
+    safety_ratings: Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')]
 
 
-_gemini_request_ta = _pydantic.LazyTypeAdapter(_GeminiRequest)
-_gemini_response_ta = _pydantic.LazyTypeAdapter(_GeminiResponse)
+_gemini_request_ta = pydantic.TypeAdapter(_GeminiRequest)
+_gemini_response_ta = pydantic.TypeAdapter(_GeminiResponse)
 
 # steam requests return a list of https://ai.google.dev/api/generate-content#method:-models.streamgeneratecontent
-_gemini_streamed_response_ta = _pydantic.LazyTypeAdapter(list[_GeminiResponse])
+_gemini_streamed_response_ta = pydantic.TypeAdapter(list[_GeminiResponse], config=pydantic.ConfigDict(defer_build=True))
 
 
 class _GeminiJsonSchema: