Add ModelResponse.finish_reason and set provider_response_id while streaming (#2590)

Co-authored-by: Douwe Maan <[email protected]>

Author: fatelei

pydantic_ai_slim/pydantic_ai/messages.py

@@ -52,6 +52,15 @@
 DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
 VideoFormat: TypeAlias = Literal['mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp']
 
+FinishReason: TypeAlias = Literal[
+    'stop',
+    'length',
+    'content_filter',
+    'tool_call',
+    'error',
+]
+"""Reason the model finished generating the response, normalized to OpenTelemetry values."""
+
 
 @dataclass(repr=False)
 class SystemPromptPart:
@@ -1032,6 +1041,9 @@ class ModelResponse:
     ] = None
     """request ID as specified by the model provider. This can be used to track the specific request to the model."""
 
+    finish_reason: FinishReason | None = None
+    """Reason the model finished generating the response, normalized to OpenTelemetry values."""
+
     @deprecated('`price` is deprecated, use `cost` instead')
     def price(self) -> genai_types.PriceCalculation:  # pragma: no cover
         return self.cost()

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -28,6 +28,7 @@
 from ..messages import (
     FileUrl,
     FinalResultEvent,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -555,6 +556,10 @@ class StreamedResponse(ABC):
 
     final_result_event: FinalResultEvent | None = field(default=None, init=False)
 
+    provider_response_id: str | None = field(default=None, init=False)
+    provider_details: dict[str, Any] | None = field(default=None, init=False)
+    finish_reason: FinishReason | None = field(default=None, init=False)
+
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
     _usage: RequestUsage = field(default_factory=RequestUsage, init=False)
@@ -609,6 +614,9 @@ def get(self) -> ModelResponse:
             timestamp=self.timestamp,
             usage=self.usage(),
             provider_name=self.provider_name,
+            provider_response_id=self.provider_response_id,
+            provider_details=self.provider_details,
+            finish_reason=self.finish_reason,
         )
 
     def usage(self) -> RequestUsage:

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -21,6 +21,7 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -42,6 +43,16 @@
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
+_FINISH_REASON_MAP: dict[BetaStopReason, FinishReason] = {
+    'end_turn': 'stop',
+    'max_tokens': 'length',
+    'stop_sequence': 'stop',
+    'tool_use': 'tool_call',
+    'pause_turn': 'stop',
+    'refusal': 'content_filter',
+}
+
+
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncStream
     from anthropic.types.beta import (
@@ -70,6 +81,7 @@
         BetaServerToolUseBlock,
         BetaServerToolUseBlockParam,
         BetaSignatureDelta,
+        BetaStopReason,
         BetaTextBlock,
         BetaTextBlockParam,
         BetaTextDelta,
@@ -326,12 +338,20 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
                     )
                 )
 
+        finish_reason: FinishReason | None = None
+        provider_details: dict[str, Any] | None = None
+        if raw_finish_reason := response.stop_reason:  # pragma: no branch
+            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,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     async def _process_streamed_response(
@@ -583,6 +603,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for event in self._response:
             if isinstance(event, BetaRawMessageStartEvent):
                 self._usage = _map_usage(event)
+                self.provider_response_id = event.message.id
 
             elif isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
@@ -646,6 +667,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
             elif isinstance(event, BetaRawMessageDeltaEvent):
                 self._usage = _map_usage(event)
+                if raw_finish_reason := event.delta.stop_reason:  # pragma: no branch
+                    self.provider_details = {'finish_reason': raw_finish_reason}
+                    self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
 
             elif isinstance(event, BetaRawContentBlockStopEvent | BetaRawMessageStopEvent):  # pragma: no branch
                 current_block = None

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -20,6 +20,7 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     FileUrl,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -54,6 +55,7 @@
         ContentUnionDict,
         CountTokensConfigDict,
         ExecutableCodeDict,
+        FinishReason as GoogleFinishReason,
         FunctionCallDict,
         FunctionCallingConfigDict,
         FunctionCallingConfigMode,
@@ -99,6 +101,22 @@
 See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list.
 """
 
+_FINISH_REASON_MAP: dict[GoogleFinishReason, FinishReason | None] = {
+    GoogleFinishReason.FINISH_REASON_UNSPECIFIED: None,
+    GoogleFinishReason.STOP: 'stop',
+    GoogleFinishReason.MAX_TOKENS: 'length',
+    GoogleFinishReason.SAFETY: 'content_filter',
+    GoogleFinishReason.RECITATION: 'content_filter',
+    GoogleFinishReason.LANGUAGE: 'error',
+    GoogleFinishReason.OTHER: None,
+    GoogleFinishReason.BLOCKLIST: 'content_filter',
+    GoogleFinishReason.PROHIBITED_CONTENT: 'content_filter',
+    GoogleFinishReason.SPII: 'content_filter',
+    GoogleFinishReason.MALFORMED_FUNCTION_CALL: 'error',
+    GoogleFinishReason.IMAGE_SAFETY: 'content_filter',
+    GoogleFinishReason.UNEXPECTED_TOOL_CALL: 'error',
+}
+
 
 class GoogleModelSettings(ModelSettings, total=False):
     """Settings used for a Gemini model request."""
@@ -403,11 +421,14 @@ def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
                     'Content field missing from Gemini response', str(response)
                 )  # pragma: no cover
         parts = candidate.content.parts or []
-        vendor_id = response.response_id or None
+
+        vendor_id = response.response_id
         vendor_details: dict[str, Any] | None = None
-        finish_reason = candidate.finish_reason
-        if finish_reason:  # pragma: no branch
-            vendor_details = {'finish_reason': finish_reason.value}
+        finish_reason: FinishReason | None = None
+        if raw_finish_reason := candidate.finish_reason:  # pragma: no branch
+            vendor_details = {'finish_reason': raw_finish_reason.value}
+            finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
         usage = _metadata_as_usage(response)
         return _process_response_from_parts(
             parts,
@@ -416,6 +437,7 @@ def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
+            finish_reason=finish_reason,
         )
 
     async def _process_streamed_response(
@@ -550,6 +572,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
             assert chunk.candidates is not None
             candidate = chunk.candidates[0]
+
+            if chunk.response_id:  # pragma: no branch
+                self.provider_response_id = chunk.response_id
+
+            if raw_finish_reason := candidate.finish_reason:
+                self.provider_details = {'finish_reason': raw_finish_reason.value}
+                self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
             if candidate.content is None or candidate.content.parts is None:
                 if candidate.finish_reason == 'STOP':  # pragma: no cover
                     # Normal completion - skip this chunk
@@ -632,6 +662,7 @@ def _process_response_from_parts(
     usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
+    finish_reason: FinishReason | None = None,
 ) -> ModelResponse:
     items: list[ModelResponsePart] = []
     for part in parts:
@@ -672,6 +703,7 @@ def _process_response_from_parts(
         provider_response_id=vendor_id,
         provider_details=vendor_details,
         provider_name=provider_name,
+        finish_reason=finish_reason,
     )
 
 

pydantic_ai_slim/pydantic_ai/models/instrumented.py

@@ -221,7 +221,10 @@ def messages_to_otel_messages(self, messages: list[ModelMessage]) -> list[_otel_
                         _otel_messages.ChatMessage(role='system' if is_system else 'user', parts=message_parts)
                     )
             elif isinstance(message, ModelResponse):  # pragma: no branch
-                result.append(_otel_messages.ChatMessage(role='assistant', parts=message.otel_message_parts(self)))
+                otel_message = _otel_messages.OutputMessage(role='assistant', parts=message.otel_message_parts(self))
+                if message.finish_reason is not None:
+                    otel_message['finish_reason'] = message.finish_reason
+                result.append(otel_message)
         return result
 
     def handle_messages(self, input_messages: list[ModelMessage], response: ModelResponse, system: str, span: Span):
@@ -246,12 +249,10 @@ def handle_messages(self, input_messages: list[ModelMessage], response: ModelRes
         else:
             output_messages = self.messages_to_otel_messages([response])
             assert len(output_messages) == 1
-            output_message = cast(_otel_messages.OutputMessage, output_messages[0])
-            if response.provider_details and 'finish_reason' in response.provider_details:
-                output_message['finish_reason'] = response.provider_details['finish_reason']
+            output_message = output_messages[0]
             instructions = InstrumentedModel._get_instructions(input_messages)  # pyright: ignore [reportPrivateUsage]
             system_instructions_attributes = self.system_instructions_attributes(instructions)
-            attributes = {
+            attributes: dict[str, AttributeValue] = {
                 'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
                 'gen_ai.output.messages': json.dumps([output_message]),
                 **system_instructions_attributes,
@@ -436,6 +437,8 @@ def _record_metrics():
                         )
                     if response.provider_response_id is not None:
                         attributes_to_set['gen_ai.response.id'] = response.provider_response_id
+                    if response.finish_reason is not None:
+                        attributes_to_set['gen_ai.response.finish_reasons'] = [response.finish_reason]
                     span.set_attributes(attributes_to_set)
                     span.update_name(f'{operation} {request_model}')