Ensure ModelResponse fields are set from actual model response when streaming (pydantic#2848)

Author: DouweM

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -5,7 +5,7 @@
 from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Any, Literal, cast, overload
 
 from typing_extensions import assert_never
@@ -362,13 +362,13 @@ async def _process_streamed_response(
         if isinstance(first_chunk, _utils.Unset):
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')  # pragma: no cover
 
-        # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time
-        timestamp = datetime.now(tz=timezone.utc)
+        assert isinstance(first_chunk, BetaRawMessageStartEvent)
+
         return AnthropicStreamedResponse(
             model_request_parameters=model_request_parameters,
-            _model_name=self._model_name,
+            _model_name=first_chunk.message.model,
             _response=peekable_response,
-            _timestamp=timestamp,
+            _timestamp=_utils.now_utc(),
             _provider_name=self._provider.name,
         )
 

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -22,6 +22,7 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -48,13 +49,15 @@
     from botocore.client import BaseClient
     from botocore.eventstream import EventStream
     from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
+    from mypy_boto3_bedrock_runtime.literals import StopReasonType
     from mypy_boto3_bedrock_runtime.type_defs import (
         ContentBlockOutputTypeDef,
         ContentBlockUnionTypeDef,
         ConverseRequestTypeDef,
         ConverseResponseTypeDef,
         ConverseStreamMetadataEventTypeDef,
         ConverseStreamOutputTypeDef,
+        ConverseStreamResponseTypeDef,
         DocumentBlockTypeDef,
         GuardrailConfigurationTypeDef,
         ImageBlockTypeDef,
@@ -135,6 +138,15 @@
 P = ParamSpec('P')
 T = typing.TypeVar('T')
 
+_FINISH_REASON_MAP: dict[StopReasonType, FinishReason] = {
+    'content_filtered': 'content_filter',
+    'end_turn': 'stop',
+    'guardrail_intervened': 'content_filter',
+    'max_tokens': 'length',
+    'stop_sequence': 'stop',
+    'tool_use': 'tool_call',
+}
+
 
 class BedrockModelSettings(ModelSettings, total=False):
     """Settings for Bedrock models.
@@ -270,8 +282,9 @@ async def request_stream(
         yield BedrockStreamedResponse(
             model_request_parameters=model_request_parameters,
             _model_name=self.model_name,
-            _event_stream=response,
+            _event_stream=response['stream'],
             _provider_name=self._provider.name,
+            _provider_response_id=response.get('ResponseMetadata', {}).get('RequestId', None),
         )
 
     async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse:
@@ -301,12 +314,18 @@ async def _process_response(self, response: ConverseResponseTypeDef) -> ModelRes
             output_tokens=response['usage']['outputTokens'],
         )
         response_id = response.get('ResponseMetadata', {}).get('RequestId', None)
+        raw_finish_reason = response['stopReason']
+        provider_details = {'finish_reason': raw_finish_reason}
+        finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
         return ModelResponse(
             parts=items,
             usage=u,
             model_name=self.model_name,
             provider_response_id=response_id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     @overload
@@ -316,7 +335,7 @@ async def _messages_create(
         stream: Literal[True],
         model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> EventStream[ConverseStreamOutputTypeDef]:
+    ) -> ConverseStreamResponseTypeDef:
         pass
 
     @overload
@@ -335,7 +354,7 @@ async def _messages_create(
         stream: bool,
         model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> ConverseResponseTypeDef | EventStream[ConverseStreamOutputTypeDef]:
+    ) -> ConverseResponseTypeDef | ConverseStreamResponseTypeDef:
         system_prompt, bedrock_messages = await self._map_messages(messages)
         inference_config = self._map_inference_config(model_settings)
 
@@ -372,7 +391,6 @@ async def _messages_create(
 
         if stream:
             model_response = await anyio.to_thread.run_sync(functools.partial(self.client.converse_stream, **params))
-            model_response = model_response['stream']
         else:
             model_response = await anyio.to_thread.run_sync(functools.partial(self.client.converse, **params))
         return model_response
@@ -599,25 +617,30 @@ class BedrockStreamedResponse(StreamedResponse):
     _event_stream: EventStream[ConverseStreamOutputTypeDef]
     _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc)
+    _provider_response_id: str | None = None
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         """Return an async iterator of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
         This method should be implemented by subclasses to translate the vendor-specific stream of events into
         pydantic_ai-format events.
         """
+        if self._provider_response_id is not None:  # pragma: no cover
+            self.provider_response_id = self._provider_response_id
+
         chunk: ConverseStreamOutputTypeDef
         tool_id: str | None = None
         async for chunk in _AsyncIteratorWrapper(self._event_stream):
             match chunk:
                 case {'messageStart': _}:
                     continue
-                case {'messageStop': _}:
-                    continue
+                case {'messageStop': message_stop}:
+                    raw_finish_reason = message_stop['stopReason']
+                    self.provider_details = {'finish_reason': raw_finish_reason}
+                    self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
                 case {'metadata': metadata}:
                     if 'usage' in metadata:  # pragma: no branch
                         self._usage += self._map_usage(metadata)
-                    continue
                 case {'contentBlockStart': content_block_start}:
                     index = content_block_start['contentBlockIndex']
                     start = content_block_start['start']

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -14,6 +14,7 @@
 from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -36,6 +37,7 @@
     from cohere import (
         AssistantChatMessageV2,
         AsyncClientV2,
+        ChatFinishReason,
         ChatMessageV2,
         SystemChatMessageV2,
         TextAssistantMessageV2ContentItem,
@@ -80,6 +82,14 @@
 See [Cohere's docs](https://docs.cohere.com/v2/docs/models) for a list of all available models.
 """
 
+_FINISH_REASON_MAP: dict[ChatFinishReason, FinishReason] = {
+    'COMPLETE': 'stop',
+    'STOP_SEQUENCE': 'stop',
+    'MAX_TOKENS': 'length',
+    'TOOL_CALL': 'tool_call',
+    'ERROR': 'error',
+}
+
 
 class CohereModelSettings(ModelSettings, total=False):
     """Settings used for a Cohere model request."""
@@ -205,8 +215,18 @@ def _process_response(self, response: V2ChatResponse) -> ModelResponse:
                         tool_call_id=c.id or _generate_tool_call_id(),
                     )
                 )
+
+        raw_finish_reason = response.finish_reason
+        provider_details = {'finish_reason': raw_finish_reason}
+        finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
         return ModelResponse(
-            parts=parts, usage=_map_usage(response), model_name=self._model_name, provider_name=self._provider.name
+            parts=parts,
+            usage=_map_usage(response),
+            model_name=self._model_name,
+            provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -453,7 +453,7 @@ async def _process_streamed_response(
 
         return GeminiStreamedResponse(
             model_request_parameters=model_request_parameters,
-            _model_name=self._model_name,
+            _model_name=first_chunk.model_version or self._model_name,
             _response=peekable_response,
             _timestamp=first_chunk.create_time or _utils.now_utc(),
             _provider_name=self._provider.name,

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -23,6 +23,7 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -100,6 +101,14 @@
 See <https://console.groq.com/docs/models> for an up to date date list of models and more details.
 """
 
+_FINISH_REASON_MAP: dict[Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'], FinishReason] = {
+    'stop': 'stop',
+    'length': 'length',
+    'tool_calls': 'tool_call',
+    'content_filter': 'content_filter',
+    'function_call': 'tool_call',
+}
+
 
 class GroqModelSettings(ModelSettings, total=False):
     """Settings used for a Groq model request."""
@@ -186,7 +195,13 @@ async def request(
                         tool_name=error.error.failed_generation.name,
                         args=error.error.failed_generation.arguments,
                     )
-                    return ModelResponse(parts=[tool_call_part])
+                    return ModelResponse(
+                        parts=[tool_call_part],
+                        model_name=e.model_name,
+                        timestamp=_utils.now_utc(),
+                        provider_name=self._provider.name,
+                        finish_reason='error',
+                    )
                 except ValidationError:
                     pass
             raise
@@ -315,13 +330,19 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
+
+        raw_finish_reason = choice.finish_reason
+        provider_details = {'finish_reason': raw_finish_reason}
+        finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
         return ModelResponse(
             parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     async def _process_streamed_response(
@@ -338,7 +359,7 @@ async def _process_streamed_response(
         return GroqStreamedResponse(
             model_request_parameters=model_request_parameters,
             _response=peekable_response,
-            _model_name=self._model_name,
+            _model_name=first_chunk.model,
             _model_profile=self.profile,
             _timestamp=number_to_datetime(first_chunk.created),
             _provider_name=self._provider.name,
@@ -497,11 +518,18 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             async for chunk in self._response:
                 self._usage += _map_usage(chunk)
 
+                if chunk.id:  # pragma: no branch
+                    self.provider_response_id = chunk.id
+
                 try:
                     choice = chunk.choices[0]
                 except IndexError:
                     continue
 
+                if raw_finish_reason := choice.finish_reason:
+                    self.provider_details = {'finish_reason': raw_finish_reason}
+                    self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
                 # Handle the text part of the response
                 content = choice.delta.content
                 if content is not None:

pydantic_ai_slim/pydantic_ai/models/huggingface.py

@@ -20,6 +20,7 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -58,6 +59,7 @@
         ChatCompletionOutput,
         ChatCompletionOutputMessage,
         ChatCompletionStreamOutput,
+        TextGenerationOutputFinishReason,
     )
     from huggingface_hub.errors import HfHubHTTPError
 
@@ -94,6 +96,12 @@
 You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending).
 """
 
+_FINISH_REASON_MAP: dict[TextGenerationOutputFinishReason, FinishReason] = {
+    'length': 'length',
+    'eos_token': 'stop',
+    'stop_sequence': 'stop',
+}
+
 
 class HuggingFaceModelSettings(ModelSettings, total=False):
     """Settings used for a Hugging Face model request."""
@@ -266,13 +274,20 @@ def _process_response(self, response: ChatCompletionOutput) -> ModelResponse:
         if tool_calls is not None:
             for c in tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id))
+
+        raw_finish_reason = choice.finish_reason
+        provider_details = {'finish_reason': raw_finish_reason}
+        finish_reason = _FINISH_REASON_MAP.get(cast(TextGenerationOutputFinishReason, raw_finish_reason), None)
+
         return ModelResponse(
             parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     async def _process_streamed_response(
@@ -288,7 +303,7 @@ async def _process_streamed_response(
 
         return HuggingFaceStreamedResponse(
             model_request_parameters=model_request_parameters,
-            _model_name=self._model_name,
+            _model_name=first_chunk.model,
             _model_profile=self.profile,
             _response=peekable_response,
             _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
@@ -445,11 +460,20 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
             self._usage += _map_usage(chunk)
 
+            if chunk.id:  # pragma: no branch
+                self.provider_response_id = chunk.id
+
             try:
                 choice = chunk.choices[0]
             except IndexError:
                 continue
 
+            if raw_finish_reason := choice.finish_reason:
+                self.provider_details = {'finish_reason': raw_finish_reason}
+                self.finish_reason = _FINISH_REASON_MAP.get(
+                    cast(TextGenerationOutputFinishReason, raw_finish_reason), None
+                )
+
             # Handle the text part of the response
             content = choice.delta.content
             if content is not None: