pydantic-ai-slim 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

pydantic_ai/models/bedrock.py

@@ -22,6 +22,7 @@ from pydantic_ai.messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -48,6 +49,7 @@ if TYPE_CHECKING:
     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,
@@ -55,6 +57,7 @@ if TYPE_CHECKING:
         ConverseResponseTypeDef,
         ConverseStreamMetadataEventTypeDef,
         ConverseStreamOutputTypeDef,
+        ConverseStreamResponseTypeDef,
         DocumentBlockTypeDef,
         GuardrailConfigurationTypeDef,
         ImageBlockTypeDef,
@@ -63,7 +66,6 @@ if TYPE_CHECKING:
         PerformanceConfigurationTypeDef,
         PromptVariableValuesTypeDef,
         ReasoningContentBlockOutputTypeDef,
-        ReasoningTextBlockTypeDef,
         SystemContentBlockTypeDef,
         ToolChoiceTypeDef,
         ToolConfigurationTypeDef,
@@ -135,6 +137,15 @@ See [the Bedrock docs](https://docs.aws.amazon.com/bedrock/latest/userguide/mode
 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 +281,9 @@ class BedrockConverseModel(Model):
         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:
@@ -279,13 +291,24 @@ class BedrockConverseModel(Model):
         if message := response['output'].get('message'):  # pragma: no branch
             for item in message['content']:
                 if reasoning_content := item.get('reasoningContent'):
-                    reasoning_text = reasoning_content.get('reasoningText')
-                    if reasoning_text:  # pragma: no branch
-                        thinking_part = ThinkingPart(
-                            content=reasoning_text['text'],
-                            signature=reasoning_text.get('signature'),
+                    if redacted_content := reasoning_content.get('redactedContent'):
+                        items.append(
+                            ThinkingPart(
+                                id='redacted_content',
+                                content='',
+                                signature=redacted_content.decode('utf-8'),
+                                provider_name=self.system,
+                            )
+                        )
+                    elif reasoning_text := reasoning_content.get('reasoningText'):  # pragma: no branch
+                        signature = reasoning_text.get('signature')
+                        items.append(
+                            ThinkingPart(
+                                content=reasoning_text['text'],
+                                signature=signature,
+                                provider_name=self.system if signature else None,
+                            )
                         )
-                        items.append(thinking_part)
                 if text := item.get('text'):
                     items.append(TextPart(content=text))
                 elif tool_use := item.get('toolUse'):
@@ -301,12 +324,18 @@ class BedrockConverseModel(Model):
             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 +345,7 @@ class BedrockConverseModel(Model):
         stream: Literal[True],
         model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> EventStream[ConverseStreamOutputTypeDef]:
+    ) -> ConverseStreamResponseTypeDef:
         pass
 
     @overload
@@ -335,7 +364,7 @@ class BedrockConverseModel(Model):
         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 +401,6 @@ class BedrockConverseModel(Model):
 
         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
@@ -476,19 +504,26 @@ class BedrockConverseModel(Model):
                     if isinstance(item, TextPart):
                         content.append({'text': item.content})
                     elif isinstance(item, ThinkingPart):
-                        if BedrockModelProfile.from_profile(self.profile).bedrock_send_back_thinking_parts:
-                            reasoning_text: ReasoningTextBlockTypeDef = {
-                                'text': item.content,
-                            }
-                            if item.signature:
-                                reasoning_text['signature'] = item.signature
-                            reasoning_content: ReasoningContentBlockOutputTypeDef = {
-                                'reasoningText': reasoning_text,
-                            }
+                        if (
+                            item.provider_name == self.system
+                            and item.signature
+                            and BedrockModelProfile.from_profile(self.profile).bedrock_send_back_thinking_parts
+                        ):
+                            if item.id == 'redacted_content':
+                                reasoning_content: ReasoningContentBlockOutputTypeDef = {
+                                    'redactedContent': item.signature.encode('utf-8'),
+                                }
+                            else:
+                                reasoning_content: ReasoningContentBlockOutputTypeDef = {
+                                    'reasoningText': {
+                                        'text': item.content,
+                                        'signature': item.signature,
+                                    }
+                                }
                             content.append({'reasoningContent': reasoning_content})
                         else:
-                            # NOTE: We don't pass the thinking part to Bedrock for models other than Claude since it raises an error.
-                            pass
+                            start_tag, end_tag = self.profile.thinking_tags
+                            content.append({'text': '\n'.join([start_tag, item.content, end_tag])})
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):
                         pass
                     else:
@@ -599,25 +634,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']
@@ -637,11 +677,21 @@ class BedrockStreamedResponse(StreamedResponse):
                     index = content_block_delta['contentBlockIndex']
                     delta = content_block_delta['delta']
                     if 'reasoningContent' in delta:
-                        yield self._parts_manager.handle_thinking_delta(
-                            vendor_part_id=index,
-                            content=delta['reasoningContent'].get('text'),
-                            signature=delta['reasoningContent'].get('signature'),
-                        )
+                        if redacted_content := delta['reasoningContent'].get('redactedContent'):
+                            yield self._parts_manager.handle_thinking_delta(
+                                vendor_part_id=index,
+                                id='redacted_content',
+                                signature=redacted_content.decode('utf-8'),
+                                provider_name=self.provider_name,
+                            )
+                        else:
+                            signature = delta['reasoningContent'].get('signature')
+                            yield self._parts_manager.handle_thinking_delta(
+                                vendor_part_id=index,
+                                content=delta['reasoningContent'].get('text'),
+                                signature=signature,
+                                provider_name=self.provider_name if signature else None,
+                            )
                     if 'text' in delta:
                         maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
                         if maybe_event is not None:  # pragma: no branch

pydantic_ai/models/cohere.py

@@ -6,7 +6,6 @@ from typing import Literal, cast
 
 from typing_extensions import assert_never
 
-from pydantic_ai._thinking_part import split_content_into_text_and_thinking
 from pydantic_ai.exceptions import UserError
 
 from .. import ModelHTTPError, usage
@@ -14,6 +13,7 @@ from .._utils import generate_tool_call_id as _generate_tool_call_id, guard_tool
 from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -35,10 +35,13 @@ from . import Model, ModelRequestParameters, check_allow_model_requests
 try:
     from cohere import (
         AssistantChatMessageV2,
+        AssistantMessageV2ContentItem,
         AsyncClientV2,
+        ChatFinishReason,
         ChatMessageV2,
         SystemChatMessageV2,
         TextAssistantMessageV2ContentItem,
+        ThinkingAssistantMessageV2ContentItem,
         ToolCallV2,
         ToolCallV2Function,
         ToolChatMessageV2,
@@ -80,6 +83,14 @@ allow any name in the type hints.
 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."""
@@ -191,11 +202,12 @@ class CohereModel(Model):
     def _process_response(self, response: V2ChatResponse) -> 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.extend(split_content_into_text_and_thinking(choice.text, self.profile.thinking_tags))
+        if response.message.content is not None:
+            for content in response.message.content:
+                if content.type == 'text':
+                    parts.append(TextPart(content=content.text))
+                elif content.type == 'thinking':  # pragma: no branch
+                    parts.append(ThinkingPart(content=cast(str, content.thinking)))  # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue] - https://github.com/cohere-ai/cohere-python/issues/692
         for c in response.message.tool_calls or []:
             if c.function and c.function.name and c.function.arguments:  # pragma: no branch
                 parts.append(
@@ -205,8 +217,18 @@ class CohereModel(Model):
                         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]:
@@ -217,15 +239,13 @@ class CohereModel(Model):
                 cohere_messages.extend(self._map_user_message(message))
             elif isinstance(message, ModelResponse):
                 texts: list[str] = []
+                thinking: list[str] = []
                 tool_calls: list[ToolCallV2] = []
                 for item in message.parts:
                     if isinstance(item, TextPart):
                         texts.append(item.content)
                     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
+                        thinking.append(item.content)
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
@@ -233,9 +253,15 @@ class CohereModel(Model):
                         pass
                     else:
                         assert_never(item)
+
                 message_param = AssistantChatMessageV2(role='assistant')
-                if texts:
-                    message_param.content = [TextAssistantMessageV2ContentItem(text='\n\n'.join(texts))]
+                if texts or thinking:
+                    contents: list[AssistantMessageV2ContentItem] = []
+                    if thinking:
+                        contents.append(ThinkingAssistantMessageV2ContentItem(thinking='\n\n'.join(thinking)))  # pyright: ignore[reportCallIssue] - https://github.com/cohere-ai/cohere-python/issues/692
+                    if texts:  # pragma: no branch
+                        contents.append(TextAssistantMessageV2ContentItem(text='\n\n'.join(texts)))
+                    message_param.content = contents
                 if tool_calls:
                     message_param.tool_calls = tool_calls
                 cohere_messages.append(message_param)

pydantic_ai/models/function.py

@@ -31,7 +31,7 @@ from ..messages import (
     UserContent,
     UserPromptPart,
 )
-from ..profiles import ModelProfileSpec
+from ..profiles import ModelProfile, ModelProfileSpec
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse
@@ -111,6 +111,12 @@ class FunctionModel(Model):
         stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
         self._model_name = model_name or f'function:{function_name}:{stream_function_name}'
 
+        # Use a default profile that supports JSON schema and object output if none provided
+        if profile is None:
+            profile = ModelProfile(
+                supports_json_schema_output=True,
+                supports_json_object_output=True,
+            )
         super().__init__(settings=settings, profile=profile)
 
     async def request(
@@ -285,6 +291,7 @@ class FunctionStreamedResponse(StreamedResponse):
                             vendor_part_id=dtc_index,
                             content=delta.content,
                             signature=delta.signature,
+                            provider_name='function' if delta.signature else None,
                         )
                     elif isinstance(delta, DeltaToolCall):
                         if delta.json_args:

pydantic_ai/models/google.py

@@ -254,6 +254,7 @@ class GoogleModel(Model):
                     stop_sequences=generation_config.get('stop_sequences'),
                     presence_penalty=generation_config.get('presence_penalty'),
                     frequency_penalty=generation_config.get('frequency_penalty'),
+                    seed=generation_config.get('seed'),
                     thinking_config=generation_config.get('thinking_config'),
                     media_resolution=generation_config.get('media_resolution'),
                     response_mime_type=generation_config.get('response_mime_type'),
@@ -397,6 +398,7 @@ class GoogleModel(Model):
             stop_sequences=model_settings.get('stop_sequences'),
             presence_penalty=model_settings.get('presence_penalty'),
             frequency_penalty=model_settings.get('frequency_penalty'),
+            seed=model_settings.get('seed'),
             safety_settings=model_settings.get('google_safety_settings'),
             thinking_config=model_settings.get('google_thinking_config'),
             labels=model_settings.get('google_labels'),
@@ -451,7 +453,7 @@ class GoogleModel(Model):
 
         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,
@@ -501,7 +503,7 @@ class GoogleModel(Model):
                     message_parts = [{'text': ''}]
                 contents.append({'role': 'user', 'parts': message_parts})
             elif isinstance(m, ModelResponse):
-                contents.append(_content_model_response(m))
+                contents.append(_content_model_response(m, self.system))
             else:
                 assert_never(m)
         if instructions := self._get_instructions(messages):
@@ -566,7 +568,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _timestamp: datetime
     _provider_name: str
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         async for chunk in self._response:
             self._usage = _metadata_as_usage(chunk)
 
@@ -590,6 +592,14 @@ class GeminiStreamedResponse(StreamedResponse):
                     raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
             parts = candidate.content.parts or []
             for part in parts:
+                if part.thought_signature:
+                    signature = base64.b64encode(part.thought_signature).decode('utf-8')
+                    yield self._parts_manager.handle_thinking_delta(
+                        vendor_part_id='thinking',
+                        signature=signature,
+                        provider_name=self.provider_name,
+                    )
+
                 if part.text is not None:
                     if part.thought:
                         yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
@@ -629,29 +639,41 @@ class GeminiStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _content_model_response(m: ModelResponse) -> ContentDict:
+def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict:
     parts: list[PartDict] = []
+    thought_signature: bytes | None = None
     for item in m.parts:
+        part: PartDict = {}
+        if thought_signature:
+            part['thought_signature'] = thought_signature
+            thought_signature = None
+
         if isinstance(item, ToolCallPart):
             function_call = FunctionCallDict(name=item.tool_name, args=item.args_as_dict(), id=item.tool_call_id)
-            parts.append({'function_call': function_call})
+            part['function_call'] = function_call
         elif isinstance(item, TextPart):
-            parts.append({'text': item.content})
-        elif isinstance(item, ThinkingPart):  # pragma: no cover
-            # 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.
-            # parts.append({'text': item.content, 'thought': True})
-            pass
+            part['text'] = item.content
+        elif isinstance(item, ThinkingPart):
+            if item.provider_name == provider_name and item.signature:
+                # The thought signature is to be included on the _next_ part, not the thought part itself
+                thought_signature = base64.b64decode(item.signature)
+
+            if item.content:
+                part['text'] = item.content
+                part['thought'] = True
         elif isinstance(item, BuiltinToolCallPart):
-            if item.provider_name == 'google':
+            if item.provider_name == provider_name:
                 if item.tool_name == 'code_execution':  # pragma: no branch
-                    parts.append({'executable_code': cast(ExecutableCodeDict, item.args)})
+                    part['executable_code'] = cast(ExecutableCodeDict, item.args)
         elif isinstance(item, BuiltinToolReturnPart):
-            if item.provider_name == 'google':
+            if item.provider_name == provider_name:
                 if item.tool_name == 'code_execution':  # pragma: no branch
-                    parts.append({'code_execution_result': item.content})
+                    part['code_execution_result'] = item.content
         else:
             assert_never(item)
+
+        if part:
+            parts.append(part)
     return ContentDict(role='model', parts=parts)
 
 
@@ -665,37 +687,43 @@ def _process_response_from_parts(
     finish_reason: FinishReason | None = None,
 ) -> ModelResponse:
     items: list[ModelResponsePart] = []
+    item: ModelResponsePart | None = None
     for part in parts:
+        if part.thought_signature:
+            signature = base64.b64encode(part.thought_signature).decode('utf-8')
+            if not isinstance(item, ThinkingPart):
+                item = ThinkingPart(content='')
+                items.append(item)
+            item.signature = signature
+            item.provider_name = provider_name
+
         if part.executable_code is not None:
-            items.append(
-                BuiltinToolCallPart(
-                    provider_name='google', args=part.executable_code.model_dump(), tool_name='code_execution'
-                )
+            item = BuiltinToolCallPart(
+                provider_name=provider_name, args=part.executable_code.model_dump(), tool_name='code_execution'
             )
         elif part.code_execution_result is not None:
-            items.append(
-                BuiltinToolReturnPart(
-                    provider_name='google',
-                    tool_name='code_execution',
-                    content=part.code_execution_result,
-                    tool_call_id='not_provided',
-                )
+            item = BuiltinToolReturnPart(
+                provider_name=provider_name,
+                tool_name='code_execution',
+                content=part.code_execution_result,
+                tool_call_id='not_provided',
             )
         elif part.text is not None:
             if part.thought:
-                items.append(ThinkingPart(content=part.text))
+                item = ThinkingPart(content=part.text)
             else:
-                items.append(TextPart(content=part.text))
+                item = TextPart(content=part.text)
         elif part.function_call:
             assert part.function_call.name is not None
-            tool_call_part = ToolCallPart(tool_name=part.function_call.name, args=part.function_call.args)
+            item = ToolCallPart(tool_name=part.function_call.name, args=part.function_call.args)
             if part.function_call.id is not None:
-                tool_call_part.tool_call_id = part.function_call.id  # pragma: no cover
-            items.append(tool_call_part)
-        elif part.function_response:  # pragma: no cover
+                item.tool_call_id = part.function_call.id  # pragma: no cover
+        else:  # pragma: no cover
             raise UnexpectedModelBehavior(
-                f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
+                f'Unsupported response from Gemini, expected all parts to be function calls, text, or thoughts, got: {part!r}'
             )
+
+        items.append(item)
     return ModelResponse(
         parts=items,
         model_name=model_name,