pydantic-ai-slim 0.0.19__py3-none-any.whl → 0.0.20__py3-none-any.whl

pydantic_ai/models/__init__.py

@@ -61,6 +61,7 @@ KnownModelName = Literal[
     'mistral:codestral-latest',
     'mistral:mistral-moderation-latest',
     'ollama:codellama',
+    'ollama:deepseek-r1',
     'ollama:gemma',
     'ollama:gemma2',
     'ollama:llama3',
@@ -81,6 +82,22 @@ KnownModelName = Literal[
     'anthropic:claude-3-5-haiku-latest',
     'anthropic:claude-3-5-sonnet-latest',
     'anthropic:claude-3-opus-latest',
+    'claude-3-5-haiku-latest',
+    'claude-3-5-sonnet-latest',
+    'claude-3-opus-latest',
+    'cohere:c4ai-aya-expanse-32b',
+    'cohere:c4ai-aya-expanse-8b',
+    'cohere:command',
+    'cohere:command-light',
+    'cohere:command-light-nightly',
+    'cohere:command-nightly',
+    'cohere:command-r',
+    'cohere:command-r-03-2024',
+    'cohere:command-r-08-2024',
+    'cohere:command-r-plus',
+    'cohere:command-r-plus-04-2024',
+    'cohere:command-r-plus-08-2024',
+    'cohere:command-r7b-12-2024',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -145,6 +162,7 @@ class AgentModel(ABC):
 class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
+    _model_name: str
     _usage: Usage = field(default_factory=Usage, init=False)
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
@@ -168,7 +186,13 @@ class StreamedResponse(ABC):
 
     def get(self) -> ModelResponse:
         """Build a [`ModelResponse`][pydantic_ai.messages.ModelResponse] from the data received from the stream so far."""
-        return ModelResponse(parts=self._parts_manager.get_parts(), timestamp=self.timestamp())
+        return ModelResponse(
+            parts=self._parts_manager.get_parts(), model_name=self._model_name, timestamp=self.timestamp()
+        )
+
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
 
     def usage(self) -> Usage:
         """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""
@@ -228,6 +252,10 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .test import TestModel
 
         return TestModel()
+    elif model.startswith('cohere:'):
+        from .cohere import CohereModel
+
+        return CohereModel(model[7:])
     elif model.startswith('openai:'):
         from .openai import OpenAIModel
 

pydantic_ai/models/anthropic.py

@@ -1,14 +1,16 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from json import JSONDecodeError, loads as json_loads
 from typing import Any, Literal, Union, cast, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
-from .. import usage
+from .. import UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
     ArgsDict,
@@ -16,6 +18,7 @@ from ..messages import (
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -38,11 +41,16 @@ try:
     from anthropic.types import (
         Message as AnthropicMessage,
         MessageParam,
+        RawContentBlockDeltaEvent,
+        RawContentBlockStartEvent,
+        RawContentBlockStopEvent,
         RawMessageDeltaEvent,
         RawMessageStartEvent,
+        RawMessageStopEvent,
         RawMessageStreamEvent,
         TextBlock,
         TextBlockParam,
+        TextDelta,
         ToolChoiceParam,
         ToolParam,
         ToolResultBlockParam,
@@ -152,7 +160,7 @@ class AnthropicAgentModel(AgentModel):
     """Implementation of `AgentModel` for Anthropic models."""
 
     client: AsyncAnthropic
-    model_name: str
+    model_name: AnthropicModelName
     allow_text_result: bool
     tools: list[ToolParam]
 
@@ -186,16 +194,22 @@ class AnthropicAgentModel(AgentModel):
         self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
     ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]:
         # standalone function to make it easier to override
+        model_settings = model_settings or {}
+
+        tool_choice: ToolChoiceParam | None
+
         if not self.tools:
-            tool_choice: ToolChoiceParam | None = None
-        elif not self.allow_text_result:
-            tool_choice = {'type': 'any'}
+            tool_choice = None
         else:
-            tool_choice = {'type': 'auto'}
+            if not self.allow_text_result:
+                tool_choice = {'type': 'any'}
+            else:
+                tool_choice = {'type': 'auto'}
 
-        system_prompt, anthropic_messages = self._map_message(messages)
+            if (allow_parallel_tool_calls := model_settings.get('parallel_tool_calls')) is not None:
+                tool_choice['disable_parallel_tool_use'] = not allow_parallel_tool_calls
 
-        model_settings = model_settings or {}
+        system_prompt, anthropic_messages = self._map_message(messages)
 
         return await self.client.messages.create(
             max_tokens=model_settings.get('max_tokens', 1024),
@@ -210,8 +224,7 @@ class AnthropicAgentModel(AgentModel):
             timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: AnthropicMessage) -> ModelResponse:
+    def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
         for item in response.content:
@@ -227,26 +240,17 @@ class AnthropicAgentModel(AgentModel):
                     )
                 )
 
-        return ModelResponse(items)
+        return ModelResponse(items, model_name=self.model_name)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
-        """TODO: Process a streamed response, and prepare a streaming response to return."""
-        # We don't yet support streamed responses from Anthropic, so we raise an error here for now.
-        # Streamed responses will be supported in a future release.
-
-        raise RuntimeError('Streamed responses are not yet supported for Anthropic models.')
-
-        # Should be returning some sort of AnthropicStreamTextResponse or AnthropicStreamedResponse
-        # depending on the type of chunk we get, but we need to establish how we handle (and when we get) the following:
-        # RawMessageStartEvent
-        # RawMessageDeltaEvent
-        # RawMessageStopEvent
-        # RawContentBlockStartEvent
-        # RawContentBlockDeltaEvent
-        # RawContentBlockDeltaEvent
-        #
-        # We might refactor streaming internally before we implement this...
+    async def _process_streamed_response(self, response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
+        peekable_response = _utils.PeekableAsyncStream(response)
+        first_chunk = await peekable_response.peek()
+        if isinstance(first_chunk, _utils.Unset):
+            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
+
+        # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time
+        timestamp = datetime.now(tz=timezone.utc)
+        return AnthropicStreamedResponse(_model_name=self.model_name, _response=peekable_response, _timestamp=timestamp)
 
     @staticmethod
     def _map_message(messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
@@ -342,3 +346,63 @@ def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage
         response_tokens=response_usage.output_tokens,
         total_tokens=(request_tokens or 0) + response_usage.output_tokens,
     )
+
+
+@dataclass
+class AnthropicStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Anthropic models."""
+
+    _response: AsyncIterable[RawMessageStreamEvent]
+    _timestamp: datetime
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        current_block: TextBlock | ToolUseBlock | None = None
+        current_json: str = ''
+
+        async for event in self._response:
+            self._usage += _map_usage(event)
+
+            if isinstance(event, RawContentBlockStartEvent):
+                current_block = event.content_block
+                if isinstance(current_block, TextBlock) and current_block.text:
+                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=current_block.text)
+                elif isinstance(current_block, ToolUseBlock):
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=current_block.id,
+                        tool_name=current_block.name,
+                        args=cast(dict[str, Any], current_block.input),
+                        tool_call_id=current_block.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+
+            elif isinstance(event, RawContentBlockDeltaEvent):
+                if isinstance(event.delta, TextDelta):
+                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=event.delta.text)
+                elif (
+                    current_block and event.delta.type == 'input_json_delta' and isinstance(current_block, ToolUseBlock)
+                ):
+                    # Try to parse the JSON immediately, otherwise cache the value for later. This handles
+                    # cases where the JSON is not currently valid but will be valid once we stream more tokens.
+                    try:
+                        parsed_args = json_loads(current_json + event.delta.partial_json)
+                        current_json = ''
+                    except JSONDecodeError:
+                        current_json += event.delta.partial_json
+                        continue
+
+                    # For tool calls, we need to handle partial JSON updates
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=current_block.id,
+                        tool_name='',
+                        args=parsed_args,
+                        tool_call_id=current_block.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+
+            elif isinstance(event, (RawContentBlockStopEvent, RawMessageStopEvent)):
+                current_block = None
+
+    def timestamp(self) -> datetime:
+        return self._timestamp

pydantic_ai/models/cohere.py

@@ -0,0 +1,278 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from itertools import chain
+from typing import Literal, TypeAlias, Union
+
+from cohere import TextAssistantMessageContentItem
+from typing_extensions import assert_never
+
+from .. import result
+from .._utils import guard_tool_call_id as _guard_tool_call_id
+from ..messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from ..settings import ModelSettings
+from ..tools import ToolDefinition
+from . import (
+    AgentModel,
+    Model,
+    check_allow_model_requests,
+)
+
+try:
+    from cohere import (
+        AssistantChatMessageV2,
+        AsyncClientV2,
+        ChatMessageV2,
+        ChatResponse,
+        SystemChatMessageV2,
+        ToolCallV2,
+        ToolCallV2Function,
+        ToolChatMessageV2,
+        ToolV2,
+        ToolV2Function,
+        UserChatMessageV2,
+    )
+    from cohere.v2.client import OMIT
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `cohere` to use the Cohere model, '
+        "you can use the `cohere` optional group — `pip install 'pydantic-ai-slim[cohere]'`"
+    ) from _import_error
+
+CohereModelName: TypeAlias = Union[
+    str,
+    Literal[
+        'c4ai-aya-expanse-32b',
+        'c4ai-aya-expanse-8b',
+        'command',
+        'command-light',
+        'command-light-nightly',
+        'command-nightly',
+        'command-r',
+        'command-r-03-2024',
+        'command-r-08-2024',
+        'command-r-plus',
+        'command-r-plus-04-2024',
+        'command-r-plus-08-2024',
+        'command-r7b-12-2024',
+    ],
+]
+
+
+@dataclass(init=False)
+class CohereModel(Model):
+    """A model that uses the Cohere API.
+
+    Internally, this uses the [Cohere Python client](
+    https://github.com/cohere-ai/cohere-python) to interact with the API.
+
+    Apart from `__init__`, all methods are private or match those of the base class.
+    """
+
+    model_name: CohereModelName
+    client: AsyncClientV2 = field(repr=False)
+
+    def __init__(
+        self,
+        model_name: CohereModelName,
+        *,
+        api_key: str | None = None,
+        cohere_client: AsyncClientV2 | None = None,
+    ):
+        """Initialize an Cohere model.
+
+        Args:
+            model_name: The name of the Cohere model to use. List of model names
+                available [here](https://docs.cohere.com/docs/models#command).
+            api_key: The API key to use for authentication, if not provided, the
+                `COHERE_API_KEY` environment variable will be used if available.
+            cohere_client: An existing Cohere async client to use. If provided,
+                `api_key` must be `None`.
+        """
+        self.model_name: CohereModelName = model_name
+        if cohere_client is not None:
+            assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
+            self.client = cohere_client
+        else:
+            self.client = AsyncClientV2(api_key=api_key)  # type: ignore
+
+    async def agent_model(
+        self,
+        *,
+        function_tools: list[ToolDefinition],
+        allow_text_result: bool,
+        result_tools: list[ToolDefinition],
+    ) -> AgentModel:
+        check_allow_model_requests()
+        tools = [self._map_tool_definition(r) for r in function_tools]
+        if result_tools:
+            tools += [self._map_tool_definition(r) for r in result_tools]
+        return CohereAgentModel(
+            self.client,
+            self.model_name,
+            allow_text_result,
+            tools,
+        )
+
+    def name(self) -> str:
+        return f'cohere:{self.model_name}'
+
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> ToolV2:
+        return ToolV2(
+            type='function',
+            function=ToolV2Function(
+                name=f.name,
+                description=f.description,
+                parameters=f.parameters_json_schema,
+            ),
+        )
+
+
+@dataclass
+class CohereAgentModel(AgentModel):
+    """Implementation of `AgentModel` for Cohere models."""
+
+    client: AsyncClientV2
+    model_name: CohereModelName
+    allow_text_result: bool
+    tools: list[ToolV2]
+
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Usage]:
+        response = await self._chat(messages, model_settings)
+        return self._process_response(response), _map_usage(response)
+
+    async def _chat(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+    ) -> ChatResponse:
+        cohere_messages = list(chain(*(self._map_message(m) for m in messages)))
+        model_settings = model_settings or {}
+        return await self.client.chat(
+            model=self.model_name,
+            messages=cohere_messages,
+            tools=self.tools or OMIT,
+            max_tokens=model_settings.get('max_tokens', OMIT),
+            temperature=model_settings.get('temperature', OMIT),
+            p=model_settings.get('top_p', OMIT),
+        )
+
+    def _process_response(self, response: ChatResponse) -> ModelResponse:
+        """Process a non-streamed response, and prepare a message to return."""
+        parts: list[ModelResponsePart] = []
+        if response.message.content is not None and len(response.message.content) > 0:
+            # While Cohere's API returns a list, it only does that for future proofing
+            # and currently only one item is being returned.
+            choice = response.message.content[0]
+            parts.append(TextPart(choice.text))
+        for c in response.message.tool_calls or []:
+            if c.function and c.function.name and c.function.arguments:
+                parts.append(
+                    ToolCallPart.from_raw_args(
+                        tool_name=c.function.name,
+                        args=c.function.arguments,
+                        tool_call_id=c.id,
+                    )
+                )
+        return ModelResponse(parts=parts, model_name=self.model_name)
+
+    @classmethod
+    def _map_message(cls, message: ModelMessage) -> Iterable[ChatMessageV2]:
+        """Just maps a `pydantic_ai.Message` to a `cohere.ChatMessageV2`."""
+        if isinstance(message, ModelRequest):
+            yield from cls._map_user_message(message)
+        elif isinstance(message, ModelResponse):
+            texts: list[str] = []
+            tool_calls: list[ToolCallV2] = []
+            for item in message.parts:
+                if isinstance(item, TextPart):
+                    texts.append(item.content)
+                elif isinstance(item, ToolCallPart):
+                    tool_calls.append(_map_tool_call(item))
+                else:
+                    assert_never(item)
+            message_param = AssistantChatMessageV2(role='assistant')
+            if texts:
+                message_param.content = [TextAssistantMessageContentItem(text='\n\n'.join(texts))]
+            if tool_calls:
+                message_param.tool_calls = tool_calls
+            yield message_param
+        else:
+            assert_never(message)
+
+    @classmethod
+    def _map_user_message(cls, message: ModelRequest) -> Iterable[ChatMessageV2]:
+        for part in message.parts:
+            if isinstance(part, SystemPromptPart):
+                yield SystemChatMessageV2(role='system', content=part.content)
+            elif isinstance(part, UserPromptPart):
+                yield UserChatMessageV2(role='user', content=part.content)
+            elif isinstance(part, ToolReturnPart):
+                yield ToolChatMessageV2(
+                    role='tool',
+                    tool_call_id=_guard_tool_call_id(t=part, model_source='Cohere'),
+                    content=part.model_response_str(),
+                )
+            elif isinstance(part, RetryPromptPart):
+                if part.tool_name is None:
+                    yield UserChatMessageV2(role='user', content=part.model_response())
+                else:
+                    yield ToolChatMessageV2(
+                        role='tool',
+                        tool_call_id=_guard_tool_call_id(t=part, model_source='Cohere'),
+                        content=part.model_response(),
+                    )
+            else:
+                assert_never(part)
+
+
+def _map_tool_call(t: ToolCallPart) -> ToolCallV2:
+    return ToolCallV2(
+        id=_guard_tool_call_id(t=t, model_source='Cohere'),
+        type='function',
+        function=ToolCallV2Function(
+            name=t.tool_name,
+            arguments=t.args_as_json_str(),
+        ),
+    )
+
+
+def _map_usage(response: ChatResponse) -> result.Usage:
+    usage = response.usage
+    if usage is None:
+        return result.Usage()
+    else:
+        details: dict[str, int] = {}
+        if usage.billed_units is not None:
+            if usage.billed_units.input_tokens:
+                details['input_tokens'] = int(usage.billed_units.input_tokens)
+            if usage.billed_units.output_tokens:
+                details['output_tokens'] = int(usage.billed_units.output_tokens)
+            if usage.billed_units.search_units:
+                details['search_units'] = int(usage.billed_units.search_units)
+            if usage.billed_units.classifications:
+                details['classifications'] = int(usage.billed_units.classifications)
+
+        request_tokens = int(usage.tokens.input_tokens) if usage.tokens and usage.tokens.input_tokens else None
+        response_tokens = int(usage.tokens.output_tokens) if usage.tokens and usage.tokens.output_tokens else None
+        return result.Usage(
+            request_tokens=request_tokens,
+            response_tokens=response_tokens,
+            total_tokens=(request_tokens or 0) + (response_tokens or 0),
+            details=details,
+        )

pydantic_ai/models/function.py

@@ -71,16 +71,15 @@ class FunctionModel(Model):
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         return FunctionAgentModel(
-            self.function, self.stream_function, AgentInfo(function_tools, allow_text_result, result_tools, None)
+            self.function,
+            self.stream_function,
+            AgentInfo(function_tools, allow_text_result, result_tools, None),
         )
 
     def name(self) -> str:
-        labels: list[str] = []
-        if self.function is not None:
-            labels.append(self.function.__name__)
-        if self.stream_function is not None:
-            labels.append(f'stream-{self.stream_function.__name__}')
-        return f'function:{",".join(labels)}'
+        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}'
 
 
 @dataclass(frozen=True)
@@ -147,12 +146,15 @@ class FunctionAgentModel(AgentModel):
         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]))
 
@@ -163,13 +165,15 @@ class FunctionAgentModel(AgentModel):
         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(response_stream)
+        yield FunctionStreamedResponse(_model_name=model_name, _iter=response_stream)
 
 
 @dataclass

pydantic_ai/models/gemini.py

@@ -99,6 +99,7 @@ class GeminiModel(Model):
         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,
@@ -151,7 +152,6 @@ class GeminiAgentModel(AgentModel):
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ):
-        check_allow_model_requests()
         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]
@@ -229,15 +229,13 @@ class GeminiAgentModel(AgentModel):
                 raise exceptions.UnexpectedModelBehavior(f'Unexpected response from gemini {r.status_code}', r.text)
             yield r
 
-    @staticmethod
-    def _process_response(response: _GeminiResponse) -> ModelResponse:
+    def _process_response(self, 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)
+        return _process_response_from_parts(parts, model_name=self.model_name)
 
-    @staticmethod
-    async def _process_streamed_response(http_response: HTTPResponse) -> StreamedResponse:
+    async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         aiter_bytes = http_response.aiter_bytes()
         start_response: _GeminiResponse | None = None
@@ -258,7 +256,7 @@ class GeminiAgentModel(AgentModel):
         if start_response is None:
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return GeminiStreamedResponse(_content=content, _stream=aiter_bytes)
+        return GeminiStreamedResponse(_model_name=self.model_name, _content=content, _stream=aiter_bytes)
 
     @classmethod
     def _message_to_gemini_content(
@@ -432,7 +430,9 @@ def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart
     return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict()))
 
 
-def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: datetime | None = None) -> ModelResponse:
+def _process_response_from_parts(
+    parts: Sequence[_GeminiPartUnion], model_name: GeminiModelName, timestamp: datetime | None = None
+) -> ModelResponse:
     items: list[ModelResponsePart] = []
     for part in parts:
         if 'text' in part:
@@ -448,7 +448,7 @@ def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: d
             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())
+    return ModelResponse(parts=items, model_name=model_name, timestamp=timestamp or _utils.now_utc())
 
 
 class _GeminiFunctionCall(TypedDict):