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

pydantic_ai/models/mistral.py

@@ -5,7 +5,6 @@ from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from itertools import chain
 from typing import Any, Literal, Union, cast
 
 import pydantic_core
@@ -29,16 +28,18 @@ from ..messages import (
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
+    VideoUrl,
 )
 from ..providers import Provider, infer_provider
-from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
+from ..usage import Usage
 from . import (
     Model,
     ModelRequestParameters,
     StreamedResponse,
     check_allow_model_requests,
+    get_user_agent,
 )
 
 try:
@@ -167,7 +168,7 @@ class MistralModel(Model):
             messages, cast(MistralModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(model_request_parameters.result_tools, response)
+            yield await self._process_streamed_response(model_request_parameters.output_tools, response)
 
     @property
     def model_name(self) -> MistralModelName:
@@ -189,9 +190,9 @@ class MistralModel(Model):
         try:
             response = await self.client.chat.complete_async(
                 model=str(self._model_name),
-                messages=list(chain(*(self._map_message(m) for m in messages))),
+                messages=self._map_messages(messages),
                 n=1,
-                tools=self._map_function_and_result_tools_definition(model_request_parameters) or UNSET,
+                tools=self._map_function_and_output_tools_definition(model_request_parameters) or UNSET,
                 tool_choice=self._get_tool_choice(model_request_parameters),
                 stream=False,
                 max_tokens=model_settings.get('max_tokens', UNSET),
@@ -200,6 +201,7 @@ class MistralModel(Model):
                 timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
                 random_seed=model_settings.get('seed', UNSET),
                 stop=model_settings.get('stop_sequences', None),
+                http_headers={'User-Agent': get_user_agent()},
             )
         except SDKError as e:
             if (status_code := e.status_code) >= 400:
@@ -217,10 +219,10 @@ class MistralModel(Model):
     ) -> MistralEventStreamAsync[MistralCompletionEvent]:
         """Create a streaming completion request to the Mistral model."""
         response: MistralEventStreamAsync[MistralCompletionEvent] | None
-        mistral_messages = list(chain(*(self._map_message(m) for m in messages)))
+        mistral_messages = self._map_messages(messages)
 
         if (
-            model_request_parameters.result_tools
+            model_request_parameters.output_tools
             and model_request_parameters.function_tools
             or model_request_parameters.function_tools
         ):
@@ -229,7 +231,7 @@ class MistralModel(Model):
                 model=str(self._model_name),
                 messages=mistral_messages,
                 n=1,
-                tools=self._map_function_and_result_tools_definition(model_request_parameters) or UNSET,
+                tools=self._map_function_and_output_tools_definition(model_request_parameters) or UNSET,
                 tool_choice=self._get_tool_choice(model_request_parameters),
                 temperature=model_settings.get('temperature', UNSET),
                 top_p=model_settings.get('top_p', 1),
@@ -238,11 +240,12 @@ class MistralModel(Model):
                 presence_penalty=model_settings.get('presence_penalty'),
                 frequency_penalty=model_settings.get('frequency_penalty'),
                 stop=model_settings.get('stop_sequences', None),
+                http_headers={'User-Agent': get_user_agent()},
             )
 
-        elif model_request_parameters.result_tools:
+        elif model_request_parameters.output_tools:
             # Json Mode
-            parameters_json_schemas = [tool.parameters_json_schema for tool in model_request_parameters.result_tools]
+            parameters_json_schemas = [tool.parameters_json_schema for tool in model_request_parameters.output_tools]
             user_output_format_message = self._generate_user_output_format(parameters_json_schemas)
             mistral_messages.append(user_output_format_message)
 
@@ -251,6 +254,7 @@ class MistralModel(Model):
                 messages=mistral_messages,
                 response_format={'type': 'json_object'},
                 stream=True,
+                http_headers={'User-Agent': get_user_agent()},
             )
 
         else:
@@ -259,6 +263,7 @@ class MistralModel(Model):
                 model=str(self._model_name),
                 messages=mistral_messages,
                 stream=True,
+                http_headers={'User-Agent': get_user_agent()},
             )
         assert response, 'A unexpected empty response from Mistral.'
         return response
@@ -271,22 +276,22 @@ class MistralModel(Model):
         - "none": Prevents tool use.
         - "required": Forces tool use.
         """
-        if not model_request_parameters.function_tools and not model_request_parameters.result_tools:
+        if not model_request_parameters.function_tools and not model_request_parameters.output_tools:
             return None
-        elif not model_request_parameters.allow_text_result:
+        elif not model_request_parameters.allow_text_output:
             return 'required'
         else:
             return 'auto'
 
-    def _map_function_and_result_tools_definition(
+    def _map_function_and_output_tools_definition(
         self, model_request_parameters: ModelRequestParameters
     ) -> list[MistralTool] | None:
-        """Map function and result tools to MistralTool format.
+        """Map function and output tools to MistralTool format.
 
-        Returns None if both function_tools and result_tools are empty.
+        Returns None if both function_tools and output_tools are empty.
         """
         all_tools: list[ToolDefinition] = (
-            model_request_parameters.function_tools + model_request_parameters.result_tools
+            model_request_parameters.function_tools + model_request_parameters.output_tools
         )
         tools = [
             MistralTool(
@@ -322,7 +327,7 @@ class MistralModel(Model):
 
     async def _process_streamed_response(
         self,
-        result_tools: list[ToolDefinition],
+        output_tools: list[ToolDefinition],
         response: MistralEventStreamAsync[MistralCompletionEvent],
     ) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
@@ -340,7 +345,7 @@ class MistralModel(Model):
             _response=peekable_response,
             _model_name=self._model_name,
             _timestamp=timestamp,
-            _result_tools={c.name: c for c in result_tools},
+            _output_tools={c.name: c for c in output_tools},
         )
 
     @staticmethod
@@ -434,13 +439,12 @@ class MistralModel(Model):
             return int(1000 * timeout)
         raise NotImplementedError('Timeout object is not yet supported for MistralModel.')
 
-    @classmethod
-    def _map_user_message(cls, message: ModelRequest) -> Iterable[MistralMessages]:
+    def _map_user_message(self, message: ModelRequest) -> Iterable[MistralMessages]:
         for part in message.parts:
             if isinstance(part, SystemPromptPart):
                 yield MistralSystemMessage(content=part.content)
             elif isinstance(part, UserPromptPart):
-                yield cls._map_user_prompt(part)
+                yield self._map_user_prompt(part)
             elif isinstance(part, ToolReturnPart):
                 yield MistralToolMessage(
                     tool_call_id=part.tool_call_id,
@@ -457,28 +461,31 @@ class MistralModel(Model):
             else:
                 assert_never(part)
 
-    @classmethod
-    def _map_message(cls, message: ModelMessage) -> Iterable[MistralMessages]:
+    def _map_messages(self, messages: list[ModelMessage]) -> list[MistralMessages]:
         """Just maps a `pydantic_ai.Message` to a `MistralMessage`."""
-        if isinstance(message, ModelRequest):
-            yield from cls._map_user_message(message)
-        elif isinstance(message, ModelResponse):
-            content_chunks: list[MistralContentChunk] = []
-            tool_calls: list[MistralToolCall] = []
-
-            for part in message.parts:
-                if isinstance(part, TextPart):
-                    content_chunks.append(MistralTextChunk(text=part.content))
-                elif isinstance(part, ToolCallPart):
-                    tool_calls.append(cls._map_tool_call(part))
-                else:
-                    assert_never(part)
-            yield MistralAssistantMessage(content=content_chunks, tool_calls=tool_calls)
-        else:
-            assert_never(message)
+        mistral_messages: list[MistralMessages] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                mistral_messages.extend(self._map_user_message(message))
+            elif isinstance(message, ModelResponse):
+                content_chunks: list[MistralContentChunk] = []
+                tool_calls: list[MistralToolCall] = []
+
+                for part in message.parts:
+                    if isinstance(part, TextPart):
+                        content_chunks.append(MistralTextChunk(text=part.content))
+                    elif isinstance(part, ToolCallPart):
+                        tool_calls.append(self._map_tool_call(part))
+                    else:
+                        assert_never(part)
+                mistral_messages.append(MistralAssistantMessage(content=content_chunks, tool_calls=tool_calls))
+            else:
+                assert_never(message)
+        if instructions := self._get_instructions(messages):
+            mistral_messages.insert(0, MistralSystemMessage(content=instructions))
+        return mistral_messages
 
-    @staticmethod
-    def _map_user_prompt(part: UserPromptPart) -> MistralUserMessage:
+    def _map_user_prompt(self, part: UserPromptPart) -> MistralUserMessage:
         content: str | list[MistralContentChunk]
         if isinstance(part.content, str):
             content = part.content
@@ -498,6 +505,8 @@ class MistralModel(Model):
                         raise RuntimeError('Only image binary content is supported for Mistral.')
                 elif isinstance(item, DocumentUrl):
                     raise RuntimeError('DocumentUrl is not supported in Mistral.')
+                elif isinstance(item, VideoUrl):
+                    raise RuntimeError('VideoUrl is not supported in Mistral.')
                 else:  # pragma: no cover
                     raise RuntimeError(f'Unsupported content type: {type(item)}')
         return MistralUserMessage(content=content)
@@ -513,7 +522,7 @@ class MistralStreamedResponse(StreamedResponse):
     _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
-    _result_tools: dict[str, ToolDefinition]
+    _output_tools: dict[str, ToolDefinition]
 
     _delta_content: str = field(default='', init=False)
 
@@ -531,13 +540,13 @@ class MistralStreamedResponse(StreamedResponse):
             content = choice.delta.content
             text = _map_content(content)
             if text:
-                # Attempt to produce a result tool call from the received text
-                if self._result_tools:
+                # Attempt to produce an output tool call from the received text
+                if self._output_tools:
                     self._delta_content += text
-                    maybe_tool_call_part = self._try_get_result_tool_from_text(self._delta_content, self._result_tools)
+                    maybe_tool_call_part = self._try_get_output_tool_from_text(self._delta_content, self._output_tools)
                     if maybe_tool_call_part:
                         yield self._parts_manager.handle_tool_call_part(
-                            vendor_part_id='result',
+                            vendor_part_id='output',
                             tool_name=maybe_tool_call_part.tool_name,
                             args=maybe_tool_call_part.args_as_dict(),
                             tool_call_id=maybe_tool_call_part.tool_call_id,
@@ -563,20 +572,20 @@ class MistralStreamedResponse(StreamedResponse):
         return self._timestamp
 
     @staticmethod
-    def _try_get_result_tool_from_text(text: str, result_tools: dict[str, ToolDefinition]) -> ToolCallPart | None:
+    def _try_get_output_tool_from_text(text: str, output_tools: dict[str, ToolDefinition]) -> ToolCallPart | None:
         output_json: dict[str, Any] | None = pydantic_core.from_json(text, allow_partial='trailing-strings')
         if output_json:
-            for result_tool in result_tools.values():
-                # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
+            for output_tool in output_tools.values():
+                # NOTE: Additional verification to prevent JSON validation to crash
                 # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
                 # Example with BaseModel and required fields.
                 if not MistralStreamedResponse._validate_required_json_schema(
-                    output_json, result_tool.parameters_json_schema
+                    output_json, output_tool.parameters_json_schema
                 ):
                     continue
 
                 # The following part_id will be thrown away
-                return ToolCallPart(tool_name=result_tool.name, args=output_json)
+                return ToolCallPart(tool_name=output_tool.name, args=output_json)
 
     @staticmethod
     def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
@@ -644,21 +653,21 @@ def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk)
 
 def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None:
     """Maps the delta content from a Mistral Completion Chunk to a string or None."""
-    result: str | None = None
+    output: str | None = None
 
     if isinstance(content, MistralUnset) or not content:
-        result = None
+        output = None
     elif isinstance(content, list):
         for chunk in content:
             if isinstance(chunk, MistralTextChunk):
-                result = result or '' + chunk.text
+                output = output or '' + chunk.text
             else:
                 assert False, f'Other data types like (Image, Reference) are not yet supported,  got {type(chunk)}'
     elif isinstance(content, str):
-        result = content
+        output = 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 result and len(result) == 0:
-        result = None
+    if output and len(output) == 0:  # pragma: no cover
+        output = None
 
-    return result
+    return output