pydantic-ai-slim 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

pydantic_ai/models/groq.py

@@ -23,6 +23,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -100,6 +101,14 @@ but allow any name in the type hints.
 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 @@ class GroqModel(Model):
                         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
@@ -298,16 +313,16 @@ class GroqModel(Model):
                 tool_call_id = generate_tool_call_id()
                 items.append(
                     BuiltinToolCallPart(
-                        tool_name=tool.type, args=tool.arguments, provider_name='groq', tool_call_id=tool_call_id
+                        tool_name=tool.type, args=tool.arguments, provider_name=self.system, tool_call_id=tool_call_id
                     )
                 )
                 items.append(
                     BuiltinToolReturnPart(
-                        provider_name='groq', tool_name=tool.type, content=tool.output, tool_call_id=tool_call_id
+                        provider_name=self.system, tool_name=tool.type, content=tool.output, tool_call_id=tool_call_id
                     )
                 )
-        # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`.
         if choice.message.reasoning is not None:
+            # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`.
             items.append(ThinkingPart(content=choice.message.reasoning))
         if choice.message.content is not None:
             # NOTE: The `<think>` tag is only present if `groq_reasoning_format` is set to `raw`.
@@ -315,6 +330,10 @@ class GroqModel(Model):
         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),
@@ -322,6 +341,8 @@ class GroqModel(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 @@ class GroqModel(Model):
         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,
@@ -376,8 +397,8 @@ class GroqModel(Model):
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
                     elif isinstance(item, ThinkingPart):
-                        # Skip thinking parts when mapping to Groq messages
-                        continue
+                        start_tag, end_tag = self.profile.thinking_tags
+                        texts.append('\n'.join([start_tag, item.content, end_tag]))
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from groq
                         pass
@@ -497,11 +518,18 @@ class GroqStreamedResponse(StreamedResponse):
             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/models/huggingface.py

@@ -20,6 +20,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -58,6 +59,7 @@ try:
         ChatCompletionOutput,
         ChatCompletionOutputMessage,
         ChatCompletionStreamOutput,
+        TextGenerationOutputFinishReason,
     )
     from huggingface_hub.errors import HfHubHTTPError
 
@@ -94,6 +96,12 @@ HuggingFaceModelName = str | LatestHuggingFaceModelNames
 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,6 +274,11 @@ class HuggingFaceModel(Model):
         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),
@@ -273,6 +286,8 @@ class HuggingFaceModel(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 @@ class HuggingFaceModel(Model):
 
         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),
@@ -316,10 +331,8 @@ class HuggingFaceModel(Model):
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
                     elif isinstance(item, ThinkingPart):
-                        # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
-                        # please open an issue. The below code is the code to send thinking to the provider.
-                        # texts.append(f'<think>\n{item.content}\n</think>')
-                        pass
+                        start_tag, end_tag = self.profile.thinking_tags
+                        texts.append('\n'.join([start_tag, item.content, end_tag]))
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from huggingface
                         pass
@@ -445,11 +458,20 @@ class HuggingFaceStreamedResponse(StreamedResponse):
         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:

pydantic_ai/models/instrumented.py

@@ -221,7 +221,10 @@ class InstrumentationSettings:
                         _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 @@ class InstrumentationSettings:
         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,
@@ -420,17 +421,25 @@ class InstrumentedModel(WrapperModel):
                         return
 
                     self.instrumentation_settings.handle_messages(messages, response, system, span)
+
+                    attributes_to_set = {
+                        **response.usage.opentelemetry_attributes(),
+                        'gen_ai.response.model': response_model,
+                    }
                     try:
-                        cost_attributes = {'operation.cost': float(response.cost().total_price)}
+                        attributes_to_set['operation.cost'] = float(response.cost().total_price)
                     except LookupError:
-                        cost_attributes = {}
-                    span.set_attributes(
-                        {
-                            **response.usage.opentelemetry_attributes(),
-                            'gen_ai.response.model': response_model,
-                            **cost_attributes,
-                        }
-                    )
+                        # The cost of this provider/model is unknown, which is common.
+                        pass
+                    except Exception as e:
+                        warnings.warn(
+                            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning
+                        )
+                    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}')
 
                 yield finish
@@ -478,3 +487,7 @@ class InstrumentedModel(WrapperModel):
                 return str(value)
             except Exception as e:
                 return f'Unable to serialize: {e}'
+
+
+class CostCalculationFailedWarning(Warning):
+    """Warning raised when cost calculation fails."""

pydantic_ai/models/mistral.py

@@ -13,7 +13,6 @@ from typing_extensions import assert_never
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
 from .._run_context import RunContext
-from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc, number_to_datetime
 from ..exceptions import UserError
 from ..messages import (
@@ -21,6 +20,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -61,12 +61,15 @@ try:
         ImageURLChunk as MistralImageURLChunk,
         Mistral,
         OptionalNullable as MistralOptionalNullable,
+        ReferenceChunk as MistralReferenceChunk,
         TextChunk as MistralTextChunk,
+        ThinkChunk as MistralThinkChunk,
         ToolChoiceEnum as MistralToolChoiceEnum,
     )
     from mistralai.models import (
         ChatCompletionResponse as MistralChatCompletionResponse,
         CompletionEvent as MistralCompletionEvent,
+        FinishReason as MistralFinishReason,
         Messages as MistralMessages,
         SDKError,
         Tool as MistralTool,
@@ -98,6 +101,14 @@ allow any name in the type hints.
 Since [the Mistral docs](https://docs.mistral.ai/getting-started/models/models_overview/) for a full list.
 """
 
+_FINISH_REASON_MAP: dict[MistralFinishReason, FinishReason] = {
+    'stop': 'stop',
+    'length': 'length',
+    'model_length': 'length',
+    'error': 'error',
+    'tool_calls': 'tool_call',
+}
+
 
 class MistralModelSettings(ModelSettings, total=False):
     """Settings used for a Mistral model request."""
@@ -339,14 +350,21 @@ class MistralModel(Model):
         tool_calls = choice.message.tool_calls
 
         parts: list[ModelResponsePart] = []
-        if text := _map_content(content):
-            parts.extend(split_content_into_text_and_thinking(text, self.profile.thinking_tags))
+        text, thinking = _map_content(content)
+        for thought in thinking:
+            parts.append(ThinkingPart(content=thought))
+        if text:
+            parts.append(TextPart(content=text))
 
         if isinstance(tool_calls, list):
             for tool_call in tool_calls:
                 tool = self._map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
+        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=parts,
             usage=_map_usage(response),
@@ -354,6 +372,8 @@ class MistralModel(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(
@@ -377,7 +397,7 @@ class MistralModel(Model):
         return MistralStreamedResponse(
             model_request_parameters=model_request_parameters,
             _response=peekable_response,
-            _model_name=self._model_name,
+            _model_name=first_chunk.data.model,
             _timestamp=timestamp,
             _provider_name=self._provider.name,
         )
@@ -503,16 +523,14 @@ class MistralModel(Model):
                 mistral_messages.extend(self._map_user_message(message))
             elif isinstance(message, ModelResponse):
                 content_chunks: list[MistralContentChunk] = []
+                thinking_chunks: list[MistralTextChunk | MistralReferenceChunk] = []
                 tool_calls: list[MistralToolCall] = []
 
                 for part in message.parts:
                     if isinstance(part, TextPart):
                         content_chunks.append(MistralTextChunk(text=part.content))
                     elif isinstance(part, ThinkingPart):
-                        # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
-                        # please open an issue. The below code is the code to send thinking to the provider.
-                        # content_chunks.append(MistralTextChunk(text=f'<think>{part.content}</think>'))
-                        pass
+                        thinking_chunks.append(MistralTextChunk(text=part.content))
                     elif isinstance(part, ToolCallPart):
                         tool_calls.append(self._map_tool_call(part))
                     elif isinstance(part, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
@@ -520,6 +538,8 @@ class MistralModel(Model):
                         pass
                     else:
                         assert_never(part)
+                if thinking_chunks:
+                    content_chunks.insert(0, MistralThinkChunk(thinking=thinking_chunks))
                 mistral_messages.append(MistralAssistantMessage(content=content_chunks, tool_calls=tool_calls))
             else:
                 assert_never(message)
@@ -595,14 +615,23 @@ class MistralStreamedResponse(StreamedResponse):
         async for chunk in self._response:
             self._usage += _map_usage(chunk.data)
 
+            if chunk.data.id:  # pragma: no branch
+                self.provider_response_id = chunk.data.id
+
             try:
                 choice = chunk.data.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
-            text = _map_content(content)
+            text, thinking = _map_content(content)
+            for thought in thinking:
+                self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=thought)
             if text:
                 # Attempt to produce an output tool call from the received text
                 output_tools = {c.name: c for c in self.model_request_parameters.output_tools}
@@ -715,32 +744,37 @@ def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk)
     """Maps a Mistral Completion Chunk or Chat Completion Response to a Usage."""
     if response.usage:
         return RequestUsage(
-            input_tokens=response.usage.prompt_tokens,
-            output_tokens=response.usage.completion_tokens,
+            input_tokens=response.usage.prompt_tokens or 0,
+            output_tokens=response.usage.completion_tokens or 0,
         )
     else:
-        return RequestUsage()  # pragma: no cover
+        return RequestUsage()
 
 
-def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None:
+def _map_content(content: MistralOptionalNullable[MistralContent]) -> tuple[str | None, list[str]]:
     """Maps the delta content from a Mistral Completion Chunk to a string or None."""
-    output: str | None = None
+    text: str | None = None
+    thinking: list[str] = []
 
     if isinstance(content, MistralUnset) or not content:
-        output = None
+        return None, []
     elif isinstance(content, list):
         for chunk in content:
             if isinstance(chunk, MistralTextChunk):
-                output = output or '' + chunk.text
+                text = text or '' + chunk.text
+            elif isinstance(chunk, MistralThinkChunk):
+                for thought in chunk.thinking:
+                    if thought.type == 'text':  # pragma: no branch
+                        thinking.append(thought.text)
             else:
                 assert False, (  # pragma: no cover
                     f'Other data types like (Image, Reference) are not yet supported,  got {type(chunk)}'
                 )
     elif isinstance(content, str):
-        output = content
+        text = content
 
     # Note: Check len to handle potential mismatch between function calls and responses from the API. (`msg: not the same number of function class and responses`)
-    if output and len(output) == 0:  # pragma: no cover
-        output = None
+    if text and len(text) == 0:  # pragma: no cover
+        text = None
 
-    return output
+    return text, thinking