pydantic-ai-slim 0.0.54__py3-none-any.whl → 0.1.0__py3-none-any.whl

pydantic_ai/models/anthropic.py

@@ -31,7 +31,14 @@ from ..messages import (
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client, check_allow_model_requests
+from . import (
+    Model,
+    ModelRequestParameters,
+    StreamedResponse,
+    cached_async_http_client,
+    check_allow_model_requests,
+    get_user_agent,
+)
 
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
@@ -207,7 +214,7 @@ class AnthropicModel(Model):
         if not tools:
             tool_choice = None
         else:
-            if not model_request_parameters.allow_text_result:
+            if not model_request_parameters.allow_text_output:
                 tool_choice = {'type': 'any'}
             else:
                 tool_choice = {'type': 'auto'}
@@ -231,6 +238,7 @@ class AnthropicModel(Model):
                 top_p=model_settings.get('top_p', NOT_GIVEN),
                 timeout=model_settings.get('timeout', NOT_GIVEN),
                 metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
+                extra_headers={'User-Agent': get_user_agent()},
             )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
@@ -269,8 +277,8 @@ class AnthropicModel(Model):
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolParam]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
     async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
@@ -324,6 +332,8 @@ class AnthropicModel(Model):
                 anthropic_messages.append(MessageParam(role='assistant', content=assistant_content_params))
             else:
                 assert_never(m)
+        if instructions := self._get_instructions(messages):
+            system_prompt = f'{instructions}\n\n{system_prompt}'
         return system_prompt, anthropic_messages
 
     @staticmethod

pydantic_ai/models/bedrock.py

@@ -2,17 +2,17 @@ from __future__ import annotations
 
 import functools
 import typing
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator, Iterable, Mapping
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import TYPE_CHECKING, Generic, Literal, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast, overload
 
 import anyio
 import anyio.to_thread
 from typing_extensions import ParamSpec, assert_never
 
-from pydantic_ai import _utils, result
+from pydantic_ai import _utils, usage
 from pydantic_ai.messages import (
     AudioUrl,
     BinaryContent,
@@ -29,6 +29,7 @@ from pydantic_ai.messages import (
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
+    VideoUrl,
 )
 from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client
 from pydantic_ai.providers import Provider, infer_provider
@@ -46,12 +47,16 @@ if TYPE_CHECKING:
         ConverseResponseTypeDef,
         ConverseStreamMetadataEventTypeDef,
         ConverseStreamOutputTypeDef,
+        GuardrailConfigurationTypeDef,
         ImageBlockTypeDef,
         InferenceConfigurationTypeDef,
         MessageUnionTypeDef,
+        PerformanceConfigurationTypeDef,
+        PromptVariableValuesTypeDef,
         SystemContentBlockTypeDef,
         ToolChoiceTypeDef,
         ToolTypeDef,
+        VideoBlockTypeDef,
     )
 
 
@@ -114,10 +119,49 @@ P = ParamSpec('P')
 T = typing.TypeVar('T')
 
 
-class BedrockModelSettings(ModelSettings):
+class BedrockModelSettings(ModelSettings, total=False):
     """Settings for Bedrock models.
 
     ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+
+    See [the Bedrock Converse API docs](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_RequestSyntax) for a full list.
+    See [the boto3 implementation](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html) of the Bedrock Converse API.
+    """
+
+    bedrock_guardrail_config: GuardrailConfigurationTypeDef
+    """Content moderation and safety settings for Bedrock API requests.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailConfiguration.html>.
+    """
+
+    bedrock_performance_configuration: PerformanceConfigurationTypeDef
+    """Performance optimization settings for model inference.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_PerformanceConfiguration.html>.
+    """
+
+    bedrock_request_metadata: dict[str, str]
+    """Additional metadata to attach to Bedrock API requests.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_RequestSyntax>.
+    """
+
+    bedrock_additional_model_response_fields_paths: list[str]
+    """JSON paths to extract additional fields from model responses.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html>.
+    """
+
+    bedrock_prompt_variables: Mapping[str, PromptVariableValuesTypeDef]
+    """Variables for substitution into prompt templates.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_PromptVariableValues.html>.
+    """
+
+    bedrock_additional_model_requests_fields: Mapping[str, Any]
+    """Additional model-specific parameters to include in requests.
+
+    See more about it on <https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html>.
     """
 
 
@@ -164,8 +208,8 @@ class BedrockConverseModel(Model):
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolTypeDef]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
     @staticmethod
@@ -187,8 +231,9 @@ class BedrockConverseModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> tuple[ModelResponse, result.Usage]:
-        response = await self._messages_create(messages, False, model_settings, model_request_parameters)
+    ) -> tuple[ModelResponse, usage.Usage]:
+        settings = cast(BedrockModelSettings, model_settings or {})
+        response = await self._messages_create(messages, False, settings, model_request_parameters)
         return await self._process_response(response)
 
     @asynccontextmanager
@@ -198,10 +243,11 @@ class BedrockConverseModel(Model):
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._messages_create(messages, True, model_settings, model_request_parameters)
+        settings = cast(BedrockModelSettings, model_settings or {})
+        response = await self._messages_create(messages, True, settings, model_request_parameters)
         yield BedrockStreamedResponse(_model_name=self.model_name, _event_stream=response)
 
-    async def _process_response(self, response: ConverseResponseTypeDef) -> tuple[ModelResponse, result.Usage]:
+    async def _process_response(self, response: ConverseResponseTypeDef) -> tuple[ModelResponse, usage.Usage]:
         items: list[ModelResponsePart] = []
         if message := response['output'].get('message'):
             for item in message['content']:
@@ -217,19 +263,19 @@ class BedrockConverseModel(Model):
                             tool_call_id=tool_use['toolUseId'],
                         ),
                     )
-        usage = result.Usage(
+        u = usage.Usage(
             request_tokens=response['usage']['inputTokens'],
             response_tokens=response['usage']['outputTokens'],
             total_tokens=response['usage']['totalTokens'],
         )
-        return ModelResponse(items, model_name=self.model_name), usage
+        return ModelResponse(items, model_name=self.model_name), u
 
     @overload
     async def _messages_create(
         self,
         messages: list[ModelMessage],
         stream: Literal[True],
-        model_settings: ModelSettings | None,
+        model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> EventStream[ConverseStreamOutputTypeDef]:
         pass
@@ -239,7 +285,7 @@ class BedrockConverseModel(Model):
         self,
         messages: list[ModelMessage],
         stream: Literal[False],
-        model_settings: ModelSettings | None,
+        model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ConverseResponseTypeDef:
         pass
@@ -248,14 +294,14 @@ class BedrockConverseModel(Model):
         self,
         messages: list[ModelMessage],
         stream: bool,
-        model_settings: ModelSettings | None,
+        model_settings: BedrockModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ConverseResponseTypeDef | EventStream[ConverseStreamOutputTypeDef]:
         tools = self._get_tools(model_request_parameters)
         support_tools_choice = self.model_name.startswith(('anthropic', 'us.anthropic'))
         if not tools or not support_tools_choice:
             tool_choice: ToolChoiceTypeDef = {}
-        elif not model_request_parameters.allow_text_result:
+        elif not model_request_parameters.allow_text_output:
             tool_choice = {'any': {}}
         else:
             tool_choice = {'auto': {}}
@@ -269,6 +315,24 @@ class BedrockConverseModel(Model):
             'system': system_prompt,
             'inferenceConfig': inference_config,
         }
+
+        # Bedrock supports a set of specific extra parameters
+        if model_settings:
+            if guardrail_config := model_settings.get('bedrock_guardrail_config', None):
+                params['guardrailConfig'] = guardrail_config
+            if performance_configuration := model_settings.get('bedrock_performance_configuration', None):
+                params['performanceConfig'] = performance_configuration
+            if request_metadata := model_settings.get('bedrock_request_metadata', None):
+                params['requestMetadata'] = request_metadata
+            if additional_model_response_fields_paths := model_settings.get(
+                'bedrock_additional_model_response_fields_paths', None
+            ):
+                params['additionalModelResponseFieldPaths'] = additional_model_response_fields_paths
+            if additional_model_requests_fields := model_settings.get('bedrock_additional_model_requests_fields', None):
+                params['additionalModelRequestFields'] = additional_model_requests_fields
+            if prompt_variables := model_settings.get('bedrock_prompt_variables', None):
+                params['promptVariables'] = prompt_variables
+
         if tools:
             params['toolConfig'] = {'tools': tools}
             if tool_choice:
@@ -359,6 +423,10 @@ class BedrockConverseModel(Model):
                 bedrock_messages.append({'role': 'assistant', 'content': content})
             else:
                 assert_never(m)
+
+        if instructions := self._get_instructions(messages):
+            system_prompt.insert(0, {'text': instructions})
+
         return system_prompt, bedrock_messages
 
     @staticmethod
@@ -381,9 +449,12 @@ class BedrockConverseModel(Model):
                     elif item.is_image:
                         assert format in ('jpeg', 'png', 'gif', 'webp')
                         content.append({'image': {'format': format, 'source': {'bytes': item.data}}})
+                    elif item.is_video:
+                        assert format in ('mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp')
+                        content.append({'video': {'format': format, 'source': {'bytes': item.data}}})
                     else:
                         raise NotImplementedError('Binary content is not supported yet.')
-                elif isinstance(item, (ImageUrl, DocumentUrl)):
+                elif isinstance(item, (ImageUrl, DocumentUrl, VideoUrl)):
                     response = await cached_async_http_client().get(item.url)
                     response.raise_for_status()
                     if item.kind == 'image-url':
@@ -391,11 +462,20 @@ class BedrockConverseModel(Model):
                         assert format in ('jpeg', 'png', 'gif', 'webp'), f'Unsupported image format: {format}'
                         image: ImageBlockTypeDef = {'format': format, 'source': {'bytes': response.content}}
                         content.append({'image': image})
+
                     elif item.kind == 'document-url':
                         document_count += 1
                         name = f'Document {document_count}'
                         data = response.content
                         content.append({'document': {'name': name, 'format': item.format, 'source': {'bytes': data}}})
+
+                    elif item.kind == 'video-url':
+                        format = item.media_type.split('/')[1]
+                        assert format in ('mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp'), (
+                            f'Unsupported video format: {format}'
+                        )
+                        video: VideoBlockTypeDef = {'format': format, 'source': {'bytes': response.content}}
+                        content.append({'video': video})
                 elif isinstance(item, AudioUrl):  # pragma: no cover
                     raise NotImplementedError('Audio is not supported yet.')
                 else:
@@ -475,8 +555,8 @@ class BedrockStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
-    def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> result.Usage:
-        return result.Usage(
+    def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> usage.Usage:
+        return usage.Usage(
             request_tokens=metadata['usage']['inputTokens'],
             response_tokens=metadata['usage']['outputTokens'],
             total_tokens=metadata['usage']['totalTokens'],
@@ -494,9 +574,7 @@ class _AsyncIteratorWrapper(Generic[T]):
 
     async def __anext__(self) -> T:
         try:
-            # Run the synchronous next() call in a thread pool
-            item = await anyio.to_thread.run_sync(next, self.sync_iterator)
-            return item
+            return await anyio.to_thread.run_sync(next, self.sync_iterator)
         except RuntimeError as e:
             if type(e.__cause__) is StopIteration:
                 raise StopAsyncIteration

pydantic_ai/models/cohere.py

@@ -2,12 +2,11 @@ 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 typing_extensions import assert_never
 
-from .. import ModelHTTPError, result
+from .. import ModelHTTPError, usage
 from .._utils import generate_tool_call_id as _generate_tool_call_id, guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
     ModelMessage,
@@ -134,7 +133,7 @@ class CohereModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> tuple[ModelResponse, result.Usage]:
+    ) -> tuple[ModelResponse, usage.Usage]:
         check_allow_model_requests()
         response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
         return self._process_response(response), _map_usage(response)
@@ -156,7 +155,7 @@ class CohereModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> ChatResponse:
         tools = self._get_tools(model_request_parameters)
-        cohere_messages = list(chain(*(self._map_message(m) for m in messages)))
+        cohere_messages = self._map_messages(messages)
         try:
             return await self.client.chat(
                 model=self._model_name,
@@ -194,33 +193,38 @@ class CohereModel(Model):
                 )
         return ModelResponse(parts=parts, model_name=self._model_name)
 
-    def _map_message(self, message: ModelMessage) -> Iterable[ChatMessageV2]:
+    def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:
         """Just maps a `pydantic_ai.Message` to a `cohere.ChatMessageV2`."""
-        if isinstance(message, ModelRequest):
-            yield from self._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(self._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)
+        cohere_messages: list[ChatMessageV2] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                cohere_messages.extend(self._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(self._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
+                cohere_messages.append(message_param)
+            else:
+                assert_never(message)
+        if instructions := self._get_instructions(messages):
+            cohere_messages.insert(0, SystemChatMessageV2(role='system', content=instructions))
+        return cohere_messages
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolV2]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
     @staticmethod
@@ -274,25 +278,25 @@ class CohereModel(Model):
                 assert_never(part)
 
 
-def _map_usage(response: ChatResponse) -> result.Usage:
-    usage = response.usage
-    if usage is None:
-        return result.Usage()
+def _map_usage(response: ChatResponse) -> usage.Usage:
+    u = response.usage
+    if u is None:
+        return usage.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(
+        if u.billed_units is not None:
+            if u.billed_units.input_tokens:
+                details['input_tokens'] = int(u.billed_units.input_tokens)
+            if u.billed_units.output_tokens:
+                details['output_tokens'] = int(u.billed_units.output_tokens)
+            if u.billed_units.search_units:  # pragma: no cover
+                details['search_units'] = int(u.billed_units.search_units)
+            if u.billed_units.classifications:  # pragma: no cover
+                details['classifications'] = int(u.billed_units.classifications)
+
+        request_tokens = int(u.tokens.input_tokens) if u.tokens and u.tokens.input_tokens else None
+        response_tokens = int(u.tokens.output_tokens) if u.tokens and u.tokens.output_tokens else None
+        return usage.Usage(
             request_tokens=request_tokens,
             response_tokens=response_tokens,
             total_tokens=(request_tokens or 0) + (response_tokens or 0),

pydantic_ai/models/fallback.py

@@ -63,8 +63,9 @@ class FallbackModel(Model):
         exceptions: list[Exception] = []
 
         for model in self.models:
+            customized_model_request_parameters = model.customize_request_parameters(model_request_parameters)
             try:
-                response, usage = await model.request(messages, model_settings, model_request_parameters)
+                response, usage = await model.request(messages, model_settings, customized_model_request_parameters)
             except Exception as exc:
                 if self._fallback_on(exc):
                     exceptions.append(exc)

pydantic_ai/models/function.py

@@ -91,8 +91,8 @@ class FunctionModel(Model):
     ) -> tuple[ModelResponse, usage.Usage]:
         agent_info = AgentInfo(
             model_request_parameters.function_tools,
-            model_request_parameters.allow_text_result,
-            model_request_parameters.result_tools,
+            model_request_parameters.allow_text_output,
+            model_request_parameters.output_tools,
             model_settings,
         )
 
@@ -117,8 +117,8 @@ class FunctionModel(Model):
     ) -> AsyncIterator[StreamedResponse]:
         agent_info = AgentInfo(
             model_request_parameters.function_tools,
-            model_request_parameters.allow_text_result,
-            model_request_parameters.result_tools,
+            model_request_parameters.allow_text_output,
+            model_request_parameters.output_tools,
             model_settings,
         )
 
@@ -158,10 +158,10 @@ class AgentInfo:
     These are the tools registered via the [`tool`][pydantic_ai.Agent.tool] and
     [`tool_plain`][pydantic_ai.Agent.tool_plain] decorators.
     """
-    allow_text_result: bool
-    """Whether a plain text result is allowed."""
-    result_tools: list[ToolDefinition]
-    """The tools that can called as the final result of the run."""
+    allow_text_output: bool
+    """Whether a plain text output is allowed."""
+    output_tools: list[ToolDefinition]
+    """The tools that can called to produce the final output of the run."""
     model_settings: ModelSettings | None
     """The model settings passed to the run call."""