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

pydantic_ai/models/gemini.py

@@ -2,29 +2,32 @@ 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 +40,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 +169,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.Usage]:
+        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)
+        return self._process_response(response), _metadata_as_usage(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 +197,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 +218,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 +259,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
@@ -277,7 +300,7 @@ class GeminiStreamTextResponse(StreamTextResponse):
     _stream: AsyncIterator[bytes]
     _position: int = 0
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _cost: result.Cost = field(default_factory=result.Cost, init=False)
+    _usage: result.Usage = field(default_factory=result.Usage, init=False)
 
     async def __anext__(self) -> None:
         chunk = await self._stream.__anext__()
@@ -297,7 +320,7 @@ class GeminiStreamTextResponse(StreamTextResponse):
                 new_items, experimental_allow_partial='trailing-strings'
             )
         for r in new_responses:
-            self._cost += _metadata_as_cost(r)
+            self._usage += _metadata_as_usage(r)
             parts = r['candidates'][0]['content']['parts']
             if _all_text_parts(parts):
                 for part in parts:
@@ -307,8 +330,8 @@ class GeminiStreamTextResponse(StreamTextResponse):
                     'Streamed response with unexpected content, expected all parts to be text'
                 )
 
-    def cost(self) -> result.Cost:
-        return self._cost
+    def usage(self) -> result.Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -321,14 +344,14 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
     _content: bytearray
     _stream: AsyncIterator[bytes]
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _cost: result.Cost = field(default_factory=result.Cost, init=False)
+    _usage: result.Usage = field(default_factory=result.Usage, init=False)
 
     async def __anext__(self) -> None:
         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,23 +363,16 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
             self._content,
             experimental_allow_partial='off' if final else 'trailing-strings',
         )
-        combined_parts: list[_GeminiFunctionCallPart] = []
-        self._cost = result.Cost()
+        combined_parts: list[_GeminiPartUnion] = []
+        self._usage = result.Usage()
         for r in responses:
-            self._cost += _metadata_as_cost(r)
+            self._usage += _metadata_as_usage(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
+    def usage(self) -> result.Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -367,6 +383,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 +399,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 +438,42 @@ 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:
-    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 _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart:
+    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_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_raw_args(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 +484,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 +514,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 +528,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 +569,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 +579,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 +616,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,20 +632,20 @@ 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:
+def _metadata_as_usage(response: _GeminiResponse) -> result.Usage:
     metadata = response.get('usage_metadata')
     if metadata is None:
-        return result.Cost()
+        return result.Usage()
     details: dict[str, int] = {}
     if cached_content_token_count := metadata.get('cached_content_token_count'):
         details['cached_content_token_count'] = cached_content_token_count
-    return result.Cost(
+    return result.Usage(
         request_tokens=metadata.get('prompt_token_count', 0),
         response_tokens=metadata.get('candidates_token_count', 0),
         total_tokens=metadata.get('total_token_count', 0),
@@ -629,15 +669,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: