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

pydantic_ai/models/cohere.py

@@ -0,0 +1,290 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from itertools import chain
+from typing import Literal, Union, cast
+
+from cohere import TextAssistantMessageContentItem
+from httpx import AsyncClient as AsyncHTTPClient
+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
+
+NamedCohereModels = 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',
+]
+"""Latest / most popular named Cohere models."""
+
+CohereModelName = Union[NamedCohereModels, str]
+
+
+class CohereModelSettings(ModelSettings):
+    """Settings used for a Cohere model request."""
+
+    # This class is a placeholder for any future cohere-specific settings
+
+
+@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,
+        http_client: AsyncHTTPClient | 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
+                `CO_API_KEY` environment variable will be used if available.
+            cohere_client: An existing Cohere async client to use. If provided,
+                `api_key` and `http_client` must be `None`.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        self.model_name: CohereModelName = model_name
+        if cohere_client is not None:
+            assert http_client is None, 'Cannot provide both `cohere_client` and `http_client`'
+            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, httpx_client=http_client)  # 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, cast(CohereModelSettings, model_settings or {}))
+        return self._process_response(response), _map_usage(response)
+
+    async def _chat(
+        self,
+        messages: list[ModelMessage],
+        model_settings: CohereModelSettings,
+    ) -> ChatResponse:
+        cohere_messages = list(chain(*(self._map_message(m) for m in messages)))
+        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),
+            seed=model_settings.get('seed', OMIT),
+            presence_penalty=model_settings.get('presence_penalty', OMIT),
+            frequency_penalty=model_settings.get('frequency_penalty', 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(
+                        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

@@ -7,7 +7,7 @@ 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
+from typing import Annotated, Any, Literal, Protocol, Union, cast
 from uuid import uuid4
 
 import pydantic
@@ -48,6 +48,12 @@ 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
+
+
 @dataclass(init=False)
 class GeminiModel(Model):
     """A model that uses Gemini via `generativelanguage.googleapis.com` API.
@@ -99,6 +105,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 +158,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]
@@ -171,7 +177,9 @@ class GeminiAgentModel(AgentModel):
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> tuple[ModelResponse, usage.Usage]:
-        async with self._make_request(messages, False, model_settings) as http_response:
+        async with self._make_request(
+            messages, False, cast(GeminiModelSettings, model_settings or {})
+        ) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
         return self._process_response(response), _metadata_as_usage(response)
 
@@ -179,12 +187,12 @@ class GeminiAgentModel(AgentModel):
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
-        async with self._make_request(messages, True, model_settings) as http_response:
+        async with self._make_request(messages, True, cast(GeminiModelSettings, model_settings or {})) as http_response:
             yield await self._process_streamed_response(http_response)
 
     @asynccontextmanager
     async def _make_request(
-        self, messages: list[ModelMessage], streamed: bool, model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], streamed: bool, model_settings: GeminiModelSettings
     ) -> AsyncIterator[HTTPResponse]:
         sys_prompt_parts, contents = self._message_to_gemini_content(messages)
 
@@ -204,6 +212,10 @@ class GeminiAgentModel(AgentModel):
                 generation_config['temperature'] = temperature
             if (top_p := model_settings.get('top_p')) is not None:
                 generation_config['top_p'] = top_p
+            if (presence_penalty := model_settings.get('presence_penalty')) is not None:
+                generation_config['presence_penalty'] = presence_penalty
+            if (frequency_penalty := model_settings.get('frequency_penalty')) is not None:
+                generation_config['frequency_penalty'] = frequency_penalty
         if generation_config:
             request_data['generation_config'] = generation_config
 
@@ -222,22 +234,20 @@ class GeminiAgentModel(AgentModel):
             url,
             content=request_json,
             headers=headers,
-            timeout=(model_settings or {}).get('timeout', USE_CLIENT_DEFAULT),
+            timeout=model_settings.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) -> 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 +268,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(
@@ -400,6 +410,8 @@ class _GeminiGenerationConfig(TypedDict, total=False):
     max_output_tokens: int
     temperature: float
     top_p: float
+    presence_penalty: float
+    frequency_penalty: float
 
 
 class _GeminiContent(TypedDict):
@@ -432,14 +444,16 @@ 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:
             items.append(TextPart(content=part['text']))
         elif 'function_call' in part:
             items.append(
-                ToolCallPart.from_raw_args(
+                ToolCallPart(
                     tool_name=part['function_call']['name'],
                     args=part['function_call']['args'],
                 )
@@ -448,7 +462,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):

pydantic_ai/models/groq.py

@@ -5,7 +5,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import chain
-from typing import Literal, overload
+from typing import Literal, cast, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
@@ -47,10 +47,7 @@ except ImportError as _import_error:
 
 GroqModelName = Literal[
     'llama-3.3-70b-versatile',
-    'llama-3.1-70b-versatile',
-    'llama3-groq-70b-8192-tool-use-preview',
-    'llama3-groq-8b-8192-tool-use-preview',
-    'llama-3.1-70b-specdec',
+    'llama-3.3-70b-specdec',
     'llama-3.1-8b-instant',
     'llama-3.2-1b-preview',
     'llama-3.2-3b-preview',
@@ -60,7 +57,6 @@ GroqModelName = Literal[
     'llama3-8b-8192',
     'mixtral-8x7b-32768',
     'gemma2-9b-it',
-    'gemma-7b-it',
 ]
 """Named Groq models.
 
@@ -68,6 +64,12 @@ See [the Groq docs](https://console.groq.com/docs/models) for a full list.
 """
 
 
+class GroqModelSettings(ModelSettings):
+    """Settings used for a Groq model request."""
+
+    # This class is a placeholder for any future groq-specific settings
+
+
 @dataclass(init=False)
 class GroqModel(Model):
     """A model that uses the Groq API.
@@ -155,31 +157,31 @@ class GroqAgentModel(AgentModel):
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._completions_create(messages, False, model_settings)
+        response = await self._completions_create(messages, False, cast(GroqModelSettings, model_settings or {}))
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._completions_create(messages, True, model_settings)
+        response = await self._completions_create(messages, True, cast(GroqModelSettings, model_settings or {}))
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: GroqModelSettings
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: GroqModelSettings
     ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: bool, model_settings: GroqModelSettings
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
         # standalone function to make it easier to override
         if not self.tools:
@@ -191,13 +193,11 @@ class GroqAgentModel(AgentModel):
 
         groq_messages = list(chain(*(self._map_message(m) for m in messages)))
 
-        model_settings = model_settings or {}
-
         return await self.client.chat.completions.create(
             model=str(self.model_name),
             messages=groq_messages,
             n=1,
-            parallel_tool_calls=True if self.tools else NOT_GIVEN,
+            parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
             tools=self.tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
@@ -205,10 +205,13 @@ class GroqAgentModel(AgentModel):
             temperature=model_settings.get('temperature', NOT_GIVEN),
             top_p=model_settings.get('top_p', NOT_GIVEN),
             timeout=model_settings.get('timeout', NOT_GIVEN),
+            seed=model_settings.get('seed', NOT_GIVEN),
+            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),
         )
 
-    @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
+    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
@@ -217,20 +220,21 @@ class GroqAgentModel(AgentModel):
             items.append(TextPart(content=choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(
-                    ToolCallPart.from_raw_args(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id)
-                )
-        return ModelResponse(items, timestamp=timestamp)
+                items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
+        return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         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')
 
-        return GroqStreamedResponse(peekable_response, datetime.fromtimestamp(first_chunk.created, tz=timezone.utc))
+        return GroqStreamedResponse(
+            _response=peekable_response,
+            _model_name=self.model_name,
+            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+        )
 
     @classmethod
     def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]: