pydantic-ai-slim 0.0.54__py3-none-any.whl → 0.1.0__py3-none-any.whl

pydantic_ai/models/gemini.py

@@ -1,10 +1,8 @@
 from __future__ import annotations as _annotations
 
 import base64
-import re
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
-from copy import deepcopy
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Annotated, Any, Literal, Protocol, Union, cast
@@ -34,6 +32,7 @@ from ..messages import (
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
+    VideoUrl,
 )
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -45,6 +44,7 @@ from . import (
     check_allow_model_requests,
     get_user_agent,
 )
+from ._json_schema import JsonSchema, WalkJsonSchema
 
 LatestGeminiModelNames = Literal[
     'gemini-1.5-flash',
@@ -58,6 +58,7 @@ LatestGeminiModelNames = Literal[
     'gemini-2.0-flash-lite-preview-02-05',
     'gemini-2.0-pro-exp-02-05',
     'gemini-2.5-pro-exp-03-25',
+    'gemini-2.5-pro-preview-03-25',
 ]
 """Latest Gemini models."""
 
@@ -154,12 +155,12 @@ class GeminiModel(Model):
 
     def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
         def _customize_tool_def(t: ToolDefinition):
-            return replace(t, parameters_json_schema=_GeminiJsonSchema(t.parameters_json_schema).simplify())
+            return replace(t, parameters_json_schema=_GeminiJsonSchema(t.parameters_json_schema).walk())
 
         return ModelRequestParameters(
             function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
-            allow_text_result=model_request_parameters.allow_text_result,
-            result_tools=[_customize_tool_def(tool) for tool in model_request_parameters.result_tools],
+            allow_text_output=model_request_parameters.allow_text_output,
+            output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
         )
 
     @property
@@ -174,14 +175,14 @@ class GeminiModel(Model):
 
     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]
+        if model_request_parameters.output_tools:
+            tools += [_function_from_abstract_tool(t) for t in model_request_parameters.output_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:
+        if model_request_parameters.allow_text_output:
             return None
         elif tools:
             return _tool_config([t['name'] for t in tools['function_declarations']])
@@ -202,11 +203,11 @@ class GeminiModel(Model):
 
         request_data = _GeminiRequest(contents=contents)
         if sys_prompt_parts:
-            request_data['system_instruction'] = _GeminiTextContent(role='user', parts=sys_prompt_parts)
+            request_data['systemInstruction'] = _GeminiTextContent(role='user', parts=sys_prompt_parts)
         if tools is not None:
             request_data['tools'] = tools
         if tool_config is not None:
-            request_data['tool_config'] = tool_config
+            request_data['toolConfig'] = tool_config
 
         generation_config: _GeminiGenerationConfig = {}
         if model_settings:
@@ -221,9 +222,9 @@ class GeminiModel(Model):
             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
+                request_data['safetySettings'] = gemini_safety_settings
         if generation_config:
-            request_data['generation_config'] = generation_config
+            request_data['generationConfig'] = generation_config
 
         headers = {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
         url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'
@@ -277,9 +278,8 @@ class GeminiModel(Model):
 
         return GeminiStreamedResponse(_model_name=self._model_name, _content=content, _stream=aiter_bytes)
 
-    @classmethod
     async def _message_to_gemini_content(
-        cls, messages: list[ModelMessage]
+        self, messages: list[ModelMessage]
     ) -> tuple[list[_GeminiTextPart], list[_GeminiContent]]:
         sys_prompt_parts: list[_GeminiTextPart] = []
         contents: list[_GeminiContent] = []
@@ -291,7 +291,7 @@ class GeminiModel(Model):
                     if isinstance(part, SystemPromptPart):
                         sys_prompt_parts.append(_GeminiTextPart(text=part.content))
                     elif isinstance(part, UserPromptPart):
-                        message_parts.extend(await cls._map_user_prompt(part))
+                        message_parts.extend(await self._map_user_prompt(part))
                     elif isinstance(part, ToolReturnPart):
                         message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object()))
                     elif isinstance(part, RetryPromptPart):
@@ -309,11 +309,11 @@ class GeminiModel(Model):
                 contents.append(_content_model_response(m))
             else:
                 assert_never(m)
-
+        if instructions := self._get_instructions(messages):
+            sys_prompt_parts.insert(0, _GeminiTextPart(text=instructions))
         return sys_prompt_parts, contents
 
-    @staticmethod
-    async def _map_user_prompt(part: UserPromptPart) -> list[_GeminiPartUnion]:
+    async def _map_user_prompt(self, part: UserPromptPart) -> list[_GeminiPartUnion]:
         if isinstance(part.content, str):
             return [{'text': part.content}]
         else:
@@ -335,6 +335,8 @@ class GeminiModel(Model):
                         inline_data={'data': base64.b64encode(response.content).decode('utf-8'), 'mime_type': mime_type}
                     )
                     content.append(inline_data)
+                elif isinstance(item, VideoUrl):  # pragma: no cover
+                    raise NotImplementedError('VideoUrl is not supported for Gemini.')
                 else:
                     assert_never(item)
         return content
@@ -448,17 +450,19 @@ class _GeminiRequest(TypedDict):
     See <https://ai.google.dev/api/generate-content#request-body> for API docs.
     """
 
+    # Note: Even though Google supposedly supports camelCase and snake_case, we've had user report misbehavior
+    # when using snake_case, which is why this typeddict now uses camelCase. And anyway, the plan is to replace this
+    # with an official google SDK in the near future anyway.
     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]
+    toolConfig: NotRequired[_GeminiToolConfig]
+    safetySettings: NotRequired[list[GeminiSafetySettings]]
+    systemInstruction: NotRequired[_GeminiTextContent]
     """
     Developer generated system instructions, see
     <https://ai.google.dev/gemini-api/docs/system-instructions?lang=rest>
     """
-    generation_config: NotRequired[_GeminiGenerationConfig]
+    generationConfig: NotRequired[_GeminiGenerationConfig]
 
 
 class GeminiSafetySettings(TypedDict):
@@ -757,7 +761,7 @@ _gemini_response_ta = pydantic.TypeAdapter(_GeminiResponse)
 _gemini_streamed_response_ta = pydantic.TypeAdapter(list[_GeminiResponse], config=pydantic.ConfigDict(defer_build=True))
 
 
-class _GeminiJsonSchema:
+class _GeminiJsonSchema(WalkJsonSchema):
     """Transforms the JSON Schema from Pydantic to be suitable for Gemini.
 
     Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
@@ -768,72 +772,58 @@ class _GeminiJsonSchema:
     * gemini doesn't allow `$defs` — we need to inline the definitions where possible
     """
 
-    def __init__(self, schema: _utils.ObjectJsonSchema):
-        self.schema = deepcopy(schema)
-        self.defs = self.schema.pop('$defs', {})
-
-    def simplify(self) -> dict[str, Any]:
-        self._simplify(self.schema, refs_stack=())
-        return self.schema
+    def __init__(self, schema: JsonSchema):
+        super().__init__(schema, prefer_inlined_defs=True, simplify_nullable_unions=True)
 
-    def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
+    def transform(self, schema: JsonSchema) -> JsonSchema:
         schema.pop('title', None)
         schema.pop('default', None)
         schema.pop('$schema', None)
+        if (const := schema.pop('const', None)) is not None:  # pragma: no cover
+            # Gemini doesn't support const, but it does support enum with a single value
+            schema['enum'] = [const]
+        schema.pop('discriminator', None)
+        schema.pop('examples', None)
+
+        # TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
+        #   where we add notes about these properties to the field description?
         schema.pop('exclusiveMaximum', None)
         schema.pop('exclusiveMinimum', None)
-        if ref := schema.pop('$ref', None):
-            # noinspection PyTypeChecker
-            key = re.sub(r'^#/\$defs/', '', ref)
-            if key in refs_stack:
-                raise UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini')
-            refs_stack += (key,)
-            schema_def = self.defs[key]
-            self._simplify(schema_def, refs_stack)
-            schema.update(schema_def)
-            return
-
-        if any_of := schema.get('anyOf'):
-            for item_schema in any_of:
-                self._simplify(item_schema, refs_stack)
-            if len(any_of) == 2 and {'type': 'null'} in any_of:
-                for item_schema in any_of:
-                    if item_schema != {'type': 'null'}:
-                        schema.clear()
-                        schema.update(item_schema)
-                        schema['nullable'] = True
-                        return
 
         type_ = schema.get('type')
+        if 'oneOf' in schema and 'type' not in schema:  # pragma: no cover
+            # This gets hit when we have a discriminated union
+            # Gemini returns an API error in this case even though it says in its error message it shouldn't...
+            # Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
+            schema['anyOf'] = schema.pop('oneOf')
 
-        if type_ == 'object':
-            self._object(schema, refs_stack)
-        elif type_ == 'array':
-            return self._array(schema, refs_stack)
-        elif type_ == 'string' and (fmt := schema.pop('format', None)):
+        if type_ == 'string' and (fmt := schema.pop('format', None)):
             description = schema.get('description')
             if description:
                 schema['description'] = f'{description} (format: {fmt})'
             else:
                 schema['description'] = f'Format: {fmt}'
 
-    def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
-        ad_props = schema.pop('additionalProperties', None)
-        if ad_props:
-            raise UserError('Additional properties in JSON Schema are not supported by Gemini')
-
-        if properties := schema.get('properties'):  # pragma: no branch
-            for value in properties.values():
-                self._simplify(value, refs_stack)
-
-    def _array(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
-        if prefix_items := schema.get('prefixItems'):
-            # TODO I think this not is supported by Gemini, maybe we should raise an error?
-            for prefix_item in prefix_items:
-                self._simplify(prefix_item, refs_stack)
-
-        if items_schema := schema.get('items'):  # pragma: no branch
-            self._simplify(items_schema, refs_stack)
+        if '$ref' in schema:
+            raise UserError(f'Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema["$ref"]}')
+
+        if 'prefixItems' in schema:
+            # prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
+            prefix_items = schema.pop('prefixItems')
+            items = schema.get('items')
+            unique_items = [items] if items is not None else []
+            for item in prefix_items:
+                if item not in unique_items:
+                    unique_items.append(item)
+            if len(unique_items) > 1:  # pragma: no cover
+                schema['items'] = {'anyOf': unique_items}
+            elif len(unique_items) == 1:
+                schema['items'] = unique_items[0]
+            schema.setdefault('minItems', len(prefix_items))
+            if items is None:
+                schema.setdefault('maxItems', len(prefix_items))
+
+        return schema
 
 
 def _ensure_decodeable(content: bytearray) -> bytearray:

pydantic_ai/models/groq.py

@@ -5,7 +5,6 @@ from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from itertools import chain
 from typing import Literal, Union, cast, overload
 
 from typing_extensions import assert_never
@@ -31,7 +30,7 @@ from ..messages import (
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests
+from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, get_user_agent
 
 try:
     from groq import NOT_GIVEN, APIStatusError, AsyncGroq, AsyncStream
@@ -193,12 +192,12 @@ class GroqModel(Model):
         # standalone function to make it easier to override
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
-        elif not model_request_parameters.allow_text_result:
+        elif not model_request_parameters.allow_text_output:
             tool_choice = 'required'
         else:
             tool_choice = 'auto'
 
-        groq_messages = list(chain(*(self._map_message(m) for m in messages)))
+        groq_messages = self._map_messages(messages)
 
         try:
             return await self.client.chat.completions.create(
@@ -218,6 +217,7 @@ class GroqModel(Model):
                 presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN),
                 frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
                 logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
+                extra_headers={'User-Agent': get_user_agent()},
             )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
@@ -251,34 +251,39 @@ class GroqModel(Model):
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
-    def _map_message(self, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
+    def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `groq.types.ChatCompletionMessageParam`."""
-        if isinstance(message, ModelRequest):
-            yield from self._map_user_message(message)
-        elif isinstance(message, ModelResponse):
-            texts: list[str] = []
-            tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
-            for item in message.parts:
-                if isinstance(item, TextPart):
-                    texts.append(item.content)
-                elif isinstance(item, ToolCallPart):
-                    tool_calls.append(self._map_tool_call(item))
-                else:
-                    assert_never(item)
-            message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
-            if texts:
-                # Note: model responses from this model should only have one text item, so the following
-                # shouldn't merge multiple texts into one unless you switch models between runs:
-                message_param['content'] = '\n\n'.join(texts)
-            if tool_calls:
-                message_param['tool_calls'] = tool_calls
-            yield message_param
-        else:
-            assert_never(message)
+        groq_messages: list[chat.ChatCompletionMessageParam] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                groq_messages.extend(self._map_user_message(message))
+            elif isinstance(message, ModelResponse):
+                texts: list[str] = []
+                tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
+                for item in message.parts:
+                    if isinstance(item, TextPart):
+                        texts.append(item.content)
+                    elif isinstance(item, ToolCallPart):
+                        tool_calls.append(self._map_tool_call(item))
+                    else:
+                        assert_never(item)
+                message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
+                if texts:
+                    # Note: model responses from this model should only have one text item, so the following
+                    # shouldn't merge multiple texts into one unless you switch models between runs:
+                    message_param['content'] = '\n\n'.join(texts)
+                if tool_calls:
+                    message_param['tool_calls'] = tool_calls
+                groq_messages.append(message_param)
+            else:
+                assert_never(message)
+        if instructions := self._get_instructions(messages):
+            groq_messages.insert(0, chat.ChatCompletionSystemMessageParam(role='system', content=instructions))
+        return groq_messages
 
     @staticmethod
     def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:

pydantic_ai/models/instrumented.py

@@ -260,7 +260,7 @@ class InstrumentedModel(WrapperModel):
 
     @staticmethod
     def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]:
-        result: list[Event] = []
+        events: list[Event] = []
         for message_index, message in enumerate(messages):
             message_events: list[Event] = []
             if isinstance(message, ModelRequest):
@@ -274,10 +274,10 @@ class InstrumentedModel(WrapperModel):
                     'gen_ai.message.index': message_index,
                     **(event.attributes or {}),
                 }
-            result.extend(message_events)
-        for event in result:
+            events.extend(message_events)
+        for event in events:
             event.body = InstrumentedModel.serialize_any(event.body)
-        return result
+        return events
 
     @staticmethod
     def serialize_any(value: Any) -> str: