pydantic-ai-slim 0.0.22__py3-none-any.whl → 0.0.23__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, replace
+from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import chain
 from typing import Callable, Union
@@ -27,7 +27,7 @@ from ..messages import (
 )
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import AgentModel, Model, StreamedResponse
+from . import Model, ModelRequestParameters, StreamedResponse
 
 
 @dataclass(init=False)
@@ -40,6 +40,9 @@ class FunctionModel(Model):
     function: FunctionDef | None = None
     stream_function: StreamFunctionDef | None = None
 
+    _model_name: str = field(repr=False)
+    _system: str | None = field(default=None, repr=False)
+
     @overload
     def __init__(self, function: FunctionDef) -> None: ...
 
@@ -63,23 +66,60 @@ class FunctionModel(Model):
         self.function = function
         self.stream_function = stream_function
 
-    async def agent_model(
+        function_name = self.function.__name__ if self.function is not None else ''
+        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
+        self._model_name = f'function:{function_name}:{stream_function_name}'
+
+    async def request(
         self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        return FunctionAgentModel(
-            self.function,
-            self.stream_function,
-            AgentInfo(function_tools, allow_text_result, result_tools, None),
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, usage.Usage]:
+        agent_info = AgentInfo(
+            model_request_parameters.function_tools,
+            model_request_parameters.allow_text_result,
+            model_request_parameters.result_tools,
+            model_settings,
         )
 
-    def name(self) -> str:
-        function_name = self.function.__name__ if self.function is not None else ''
-        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
-        return f'function:{function_name}:{stream_function_name}'
+        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, agent_info)
+        else:
+            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
+            assert isinstance(response_, ModelResponse), response_
+            response = response_
+        response.model_name = f'function:{self.function.__name__}'
+        # TODO is `messages` right here? Should it just be new messages?
+        return response, _estimate_usage(chain(messages, [response]))
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        agent_info = AgentInfo(
+            model_request_parameters.function_tools,
+            model_request_parameters.allow_text_result,
+            model_request_parameters.result_tools,
+            model_settings,
+        )
+
+        assert (
+            self.stream_function is not None
+        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
+
+        response_stream = PeekableAsyncStream(self.stream_function(messages, agent_info))
+
+        first = await response_stream.peek()
+        if isinstance(first, _utils.Unset):
+            raise ValueError('Stream function must return at least one item')
+
+        yield FunctionStreamedResponse(_model_name=f'function:{self.stream_function.__name__}', _iter=response_stream)
 
 
 @dataclass(frozen=True)
@@ -119,9 +159,11 @@ class DeltaToolCall:
 DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
 """A mapping of tool call IDs to incremental changes."""
 
+# TODO: Change the signature to Callable[[list[ModelMessage], ModelSettings, ModelRequestParameters], ...]
 FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], Union[ModelResponse, Awaitable[ModelResponse]]]
 """A function used to generate a non-streamed response."""
 
+# TODO: Change signature as indicated above
 StreamFunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls]]]
 """A function used to generate a streamed response.
 
@@ -132,50 +174,6 @@ E.g. you need to yield all text or all `DeltaToolCalls`, not mix them.
 """
 
 
-@dataclass
-class FunctionAgentModel(AgentModel):
-    """Implementation of `AgentModel` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
-
-    function: FunctionDef | None
-    stream_function: StreamFunctionDef | None
-    agent_info: AgentInfo
-
-    async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, usage.Usage]:
-        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'
-        model_name = f'function:{self.function.__name__}'
-
-        if inspect.iscoroutinefunction(self.function):
-            response = await self.function(messages, agent_info)
-        else:
-            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
-            assert isinstance(response_, ModelResponse), response_
-            response = response_
-        response.model_name = model_name
-        # TODO is `messages` right here? Should it just be new messages?
-        return response, _estimate_usage(chain(messages, [response]))
-
-    @asynccontextmanager
-    async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[StreamedResponse]:
-        assert (
-            self.stream_function is not None
-        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
-        model_name = f'function:{self.stream_function.__name__}'
-
-        response_stream = PeekableAsyncStream(self.stream_function(messages, self.agent_info))
-
-        first = await response_stream.peek()
-        if isinstance(first, _utils.Unset):
-            raise ValueError('Stream function must return at least one item')
-
-        yield FunctionStreamedResponse(_model_name=model_name, _iter=response_stream)
-
-
 @dataclass
 class FunctionStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""

pydantic_ai/models/gemini.py

@@ -31,15 +31,15 @@ from ..messages import (
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
     get_user_agent,
 )
 
-GeminiModelName = Literal[
+LatestGeminiModelNames = Literal[
     'gemini-1.5-flash',
     'gemini-1.5-flash-8b',
     'gemini-1.5-pro',
@@ -48,8 +48,13 @@ GeminiModelName = Literal[
     'gemini-2.0-flash-thinking-exp-01-21',
     'gemini-exp-1206',
 ]
-"""Named Gemini models.
+"""Latest Gemini models."""
 
+GeminiModelName = Union[str, LatestGeminiModelNames]
+"""Possible Gemini model names.
+
+Since Gemini supports a variety of date-stamped models, we explicitly list the latest models but
+allow any name in the type hints.
 See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list.
 """
 
@@ -57,7 +62,7 @@ See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#mo
 class GeminiModelSettings(ModelSettings):
     """Settings used for a Gemini model request."""
 
-    # This class is a placeholder for any future gemini-specific settings
+    gemini_safety_settings: list[GeminiSafetySettings]
 
 
 @dataclass(init=False)
@@ -70,10 +75,12 @@ class GeminiModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    model_name: GeminiModelName
-    auth: AuthProtocol
-    http_client: AsyncHTTPClient
-    url: str
+    http_client: AsyncHTTPClient = field(repr=False)
+
+    _model_name: GeminiModelName = field(repr=False)
+    _auth: AuthProtocol | None = field(repr=False)
+    _url: str | None = field(repr=False)
+    _system: str | None = field(default='google-gla', repr=False)
 
     def __init__(
         self,
@@ -94,121 +101,87 @@ class GeminiModel(Model):
                 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
+        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.auth = ApiKeyAuth(api_key)
         self.http_client = http_client or cached_async_http_client()
-        self.url = url_template.format(model=model_name)
-
-    async def agent_model(
-        self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> GeminiAgentModel:
-        check_allow_model_requests()
-        return GeminiAgentModel(
-            http_client=self.http_client,
-            model_name=self.model_name,
-            auth=self.auth,
-            url=self.url,
-            function_tools=function_tools,
-            allow_text_result=allow_text_result,
-            result_tools=result_tools,
-        )
-
-    def name(self) -> str:
-        return f'google-gla:{self.model_name}'
-
-
-class AuthProtocol(Protocol):
-    """Abstract definition for Gemini authentication."""
-
-    async def headers(self) -> dict[str, str]: ...
-
-
-@dataclass
-class ApiKeyAuth:
-    """Authentication using an API key for the `X-Goog-Api-Key` header."""
-
-    api_key: str
+        self._auth = ApiKeyAuth(api_key)
+        self._url = url_template.format(model=model_name)
 
-    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
-    auth: AuthProtocol
-    tools: _GeminiTools | None
-    tool_config: _GeminiToolConfig | None
-    url: str
-
-    def __init__(
-        self,
-        http_client: AsyncHTTPClient,
-        model_name: GeminiModelName,
-        auth: AuthProtocol,
-        url: str,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ):
-        tools = [_function_from_abstract_tool(t) for t in function_tools]
-        if result_tools:
-            tools += [_function_from_abstract_tool(t) for t in result_tools]
+    @property
+    def auth(self) -> AuthProtocol:
+        assert self._auth is not None, 'Auth not initialized'
+        return self._auth
 
-        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
+    @property
+    def url(self) -> str:
+        assert self._url is not None, 'URL not initialized'
+        return self._url
 
     async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, usage.Usage]:
+        check_allow_model_requests()
         async with self._make_request(
-            messages, False, cast(GeminiModelSettings, model_settings or {})
+            messages, False, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
         ) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
         return self._process_response(response), _metadata_as_usage(response)
 
     @asynccontextmanager
     async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[StreamedResponse]:
-        async with self._make_request(messages, True, cast(GeminiModelSettings, model_settings or {})) as http_response:
+        check_allow_model_requests()
+        async with self._make_request(
+            messages, True, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
+        ) as http_response:
             yield await self._process_streamed_response(http_response)
 
+    def _get_tools(self, model_request_parameters: ModelRequestParameters) -> _GeminiTools | None:
+        tools = [_function_from_abstract_tool(t) for t in model_request_parameters.function_tools]
+        if model_request_parameters.result_tools:
+            tools += [_function_from_abstract_tool(t) for t in model_request_parameters.result_tools]
+        return _GeminiTools(function_declarations=tools) if tools else None
+
+    def _get_tool_config(
+        self, model_request_parameters: ModelRequestParameters, tools: _GeminiTools | None
+    ) -> _GeminiToolConfig | None:
+        if model_request_parameters.allow_text_result:
+            return None
+        elif tools:
+            return _tool_config([t['name'] for t in tools['function_declarations']])
+        else:
+            return _tool_config([])
+
     @asynccontextmanager
     async def _make_request(
-        self, messages: list[ModelMessage], streamed: bool, model_settings: GeminiModelSettings
+        self,
+        messages: list[ModelMessage],
+        streamed: bool,
+        model_settings: GeminiModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[HTTPResponse]:
+        tools = self._get_tools(model_request_parameters)
+        tool_config = self._get_tool_config(model_request_parameters, tools)
         sys_prompt_parts, contents = self._message_to_gemini_content(messages)
 
         request_data = _GeminiRequest(contents=contents)
         if sys_prompt_parts:
             request_data['system_instruction'] = _GeminiTextContent(role='user', parts=sys_prompt_parts)
-        if self.tools is not None:
-            request_data['tools'] = self.tools
-        if self.tool_config is not None:
-            request_data['tool_config'] = self.tool_config
+        if tools is not None:
+            request_data['tools'] = tools
+        if tool_config is not None:
+            request_data['tool_config'] = tool_config
 
         generation_config: _GeminiGenerationConfig = {}
         if model_settings:
@@ -222,6 +195,8 @@ class GeminiAgentModel(AgentModel):
                 generation_config['presence_penalty'] = presence_penalty
             if (frequency_penalty := model_settings.get('frequency_penalty')) is not None:
                 generation_config['frequency_penalty'] = frequency_penalty
+            if (gemini_safety_settings := model_settings.get('gemini_safety_settings')) != []:
+                request_data['safety_settings'] = gemini_safety_settings
         if generation_config:
             request_data['generation_config'] = generation_config
 
@@ -250,8 +225,13 @@ class GeminiAgentModel(AgentModel):
     def _process_response(self, response: _GeminiResponse) -> ModelResponse:
         if len(response['candidates']) != 1:
             raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')
+        if 'content' not in response['candidates'][0]:
+            if response['candidates'][0].get('finish_reason') == 'SAFETY':
+                raise UnexpectedModelBehavior('Safety settings triggered', str(response))
+            else:
+                raise UnexpectedModelBehavior('Content field missing from Gemini response', str(response))
         parts = response['candidates'][0]['content']['parts']
-        return _process_response_from_parts(parts, model_name=self.model_name)
+        return _process_response_from_parts(parts, model_name=self._model_name)
 
     async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
@@ -267,14 +247,14 @@ class GeminiAgentModel(AgentModel):
             )
             if responses:
                 last = responses[-1]
-                if last['candidates'] and last['candidates'][0]['content']['parts']:
+                if last['candidates'] and last['candidates'][0].get('content', {}).get('parts'):
                     start_response = last
                     break
 
         if start_response is None:
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return GeminiStreamedResponse(_model_name=self.model_name, _content=content, _stream=aiter_bytes)
+        return GeminiStreamedResponse(_model_name=self._model_name, _content=content, _stream=aiter_bytes)
 
     @classmethod
     def _message_to_gemini_content(
@@ -312,6 +292,23 @@ class GeminiAgentModel(AgentModel):
         return sys_prompt_parts, contents
 
 
+class AuthProtocol(Protocol):
+    """Abstract definition for Gemini authentication."""
+
+    async def headers(self) -> dict[str, str]: ...
+
+
+@dataclass
+class ApiKeyAuth:
+    """Authentication using an API key for the `X-Goog-Api-Key` header."""
+
+    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
 class GeminiStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for the Gemini model."""
@@ -323,6 +320,8 @@ class GeminiStreamedResponse(StreamedResponse):
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for gemini_response in self._get_gemini_responses():
             candidate = gemini_response['candidates'][0]
+            if 'content' not in candidate:
+                raise UnexpectedModelBehavior('Streamed response has no content field')
             gemini_part: _GeminiPartUnion
             for gemini_part in candidate['content']['parts']:
                 if 'text' in gemini_part:
@@ -396,6 +395,7 @@ class _GeminiRequest(TypedDict):
     contents: list[_GeminiContent]
     tools: NotRequired[_GeminiTools]
     tool_config: NotRequired[_GeminiToolConfig]
+    safety_settings: NotRequired[list[GeminiSafetySettings]]
     # we don't implement `generationConfig`, instead we use a named tool for the response
     system_instruction: NotRequired[_GeminiTextContent]
     """
@@ -405,6 +405,38 @@ class _GeminiRequest(TypedDict):
     generation_config: NotRequired[_GeminiGenerationConfig]
 
 
+class GeminiSafetySettings(TypedDict):
+    """Safety settings options for Gemini model request.
+
+    See [Gemini API docs](https://ai.google.dev/gemini-api/docs/safety-settings) for safety category and threshold descriptions.
+    For an example on how to use `GeminiSafetySettings`, see [here](../../agents.md#model-specific-settings).
+    """
+
+    category: Literal[
+        'HARM_CATEGORY_UNSPECIFIED',
+        'HARM_CATEGORY_HARASSMENT',
+        'HARM_CATEGORY_HATE_SPEECH',
+        'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+        'HARM_CATEGORY_DANGEROUS_CONTENT',
+        'HARM_CATEGORY_CIVIC_INTEGRITY',
+    ]
+    """
+    Safety settings category.
+    """
+
+    threshold: Literal[
+        'HARM_BLOCK_THRESHOLD_UNSPECIFIED',
+        'BLOCK_LOW_AND_ABOVE',
+        'BLOCK_MEDIUM_AND_ABOVE',
+        'BLOCK_ONLY_HIGH',
+        'BLOCK_NONE',
+        'OFF',
+    ]
+    """
+    Safety settings threshold.
+    """
+
+
 class _GeminiGenerationConfig(TypedDict, total=False):
     """Schema for an API request to the Gemini API.
 
@@ -581,8 +613,8 @@ class _GeminiResponse(TypedDict):
 class _GeminiCandidates(TypedDict):
     """See <https://ai.google.dev/api/generate-content#v1beta.Candidate>."""
 
-    content: _GeminiContent
-    finish_reason: NotRequired[Annotated[Literal['STOP', 'MAX_TOKENS'], pydantic.Field(alias='finishReason')]]
+    content: NotRequired[_GeminiContent]
+    finish_reason: NotRequired[Annotated[Literal['STOP', 'MAX_TOKENS', 'SAFETY'], 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.
@@ -630,6 +662,7 @@ class _GeminiSafetyRating(TypedDict):
         'HARM_CATEGORY_CIVIC_INTEGRITY',
     ]
     probability: Literal['NEGLIGIBLE', 'LOW', 'MEDIUM', 'HIGH']
+    blocked: NotRequired[bool]
 
 
 class _GeminiPromptFeedback(TypedDict):