Add prediction support for MistralModel and associated tests.

Author: TKaluza

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -77,6 +77,10 @@
     )
     from mistralai.models.assistantmessage import AssistantMessage as MistralAssistantMessage
     from mistralai.models.function import Function as MistralFunction
+    from mistralai.models.prediction import (
+        Prediction as MistralPrediction,
+        PredictionTypedDict as MistralPredictionTypedDict,
+    )
     from mistralai.models.systemmessage import SystemMessage as MistralSystemMessage
     from mistralai.models.toolmessage import ToolMessage as MistralToolMessage
     from mistralai.models.usermessage import UserMessage as MistralUserMessage
@@ -114,8 +118,13 @@ class MistralModelSettings(ModelSettings, total=False):
     """Settings used for a Mistral model request."""
 
     # ALL FIELDS MUST BE `mistral_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    mistral_prediction: str | MistralPrediction | MistralPredictionTypedDict | None
+    """Prediction content for the model to use as a prefix. See Predictive outputs.
 
-    # This class is a placeholder for any future mistral-specific settings
+    This feature is currently only supported for certain Mistral models. See the model cards at Models.
+    For example, it is supported for the latest Mistral Serie Large (> 2), Medium (> 3), Small (> 3) and Pixtral models,
+    but not for reasoning or coding models yet.
+    """
 
 
 @dataclass(init=False)
@@ -241,6 +250,7 @@ async def _completions_create(
                 timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
                 random_seed=model_settings.get('seed', UNSET),
                 stop=model_settings.get('stop_sequences', None),
+                prediction=self._map_setting_prediction(model_settings.get('mistral_prediction', None)),
                 http_headers={'User-Agent': get_user_agent()},
             )
         except SDKError as e:
@@ -281,6 +291,7 @@ async def _stream_completions_create(
                 presence_penalty=model_settings.get('presence_penalty'),
                 frequency_penalty=model_settings.get('frequency_penalty'),
                 stop=model_settings.get('stop_sequences', None),
+                prediction=self._map_setting_prediction(model_settings.get('mistral_prediction', None)),
                 http_headers={'User-Agent': get_user_agent()},
             )
 
@@ -298,6 +309,7 @@ async def _stream_completions_create(
                     'type': 'json_object'
                 },  # TODO: Should be able to use json_schema now: https://docs.mistral.ai/capabilities/structured-output/custom_structured_output/, https://github.com/mistralai/client-python/blob/bc4adf335968c8a272e1ab7da8461c9943d8e701/src/mistralai/extra/utils/response_format.py#L9
                 stream=True,
+                prediction=self._map_setting_prediction(model_settings.get('mistral_prediction', None)),
                 http_headers={'User-Agent': get_user_agent()},
             )
 
@@ -307,6 +319,7 @@ async def _stream_completions_create(
                 model=str(self._model_name),
                 messages=mistral_messages,
                 stream=True,
+                prediction=self._map_setting_prediction(model_settings.get('mistral_prediction', None)),
                 http_headers={'User-Agent': get_user_agent()},
             )
         assert response, 'A unexpected empty response from Mistral.'
@@ -427,6 +440,24 @@ def _map_tool_call(t: ToolCallPart) -> MistralToolCall:
             function=MistralFunctionCall(name=t.tool_name, arguments=t.args or {}),
         )
 
+    @staticmethod
+    def _map_setting_prediction(
+        prediction: str | MistralPredictionTypedDict | MistralPrediction | None,
+    ) -> MistralPrediction | None:
+        """Maps various prediction input types to a MistralPrediction object."""
+        if not prediction:
+            return None
+        if isinstance(prediction, MistralPrediction):
+            return prediction
+        elif isinstance(prediction, str):
+            return MistralPrediction(content=prediction)
+        elif isinstance(prediction, dict):
+            return MistralPrediction.model_validate(prediction)
+        else:
+            raise RuntimeError(
+                f'Unsupported prediction type: {type(prediction)} for MistralModelSettings. Expected str, dict, or MistralPrediction.'
+            )
+
     def _generate_user_output_format(self, schemas: list[dict[str, Any]]) -> MistralUserMessage:
         """Get a message with an example of the expected output format."""
         examples: list[dict[str, Any]] = []

tests/models/cassettes/test_mistral/test_mistral_chat_with_prediction.yaml

@@ -0,0 +1,67 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '246'
+      content-type:
+      - application/json
+      host:
+      - api.mistral.ai
+    method: POST
+    parsed_body:
+      messages:
+      - content:
+        - text: Correct only the math, respond with no explanation, no formatting.
+          type: text
+        - text: The result of 21+21=99
+          type: text
+        role: user
+      model: mistral-small-latest
+      n: 1
+      stream: false
+      top_p: 1.0
+    uri: https://api.mistral.ai/v1/chat/completions
+  response:
+    headers:
+      access-control-allow-origin:
+      - '*'
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '321'
+      content-type:
+      - application/json
+      mistral-correlation-id:
+      - 019a639b-7bf4-7481-96af-078cd1a7d277
+      strict-transport-security:
+      - max-age=15552000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        message:
+          content: The result of 21+21=42
+          role: assistant
+          tool_calls: null
+      created: 1762607725
+      id: a7952046ef794d1697627b54231df17a
+      model: mistral-small-latest
+      object: chat.completion
+      usage:
+        completion_tokens: 13
+        prompt_tokens: 28
+        total_tokens: 41
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_mistral.py

@@ -54,6 +54,10 @@
         ToolCall as MistralToolCall,
     )
     from mistralai.types.basemodel import Unset as MistralUnset
+    from mistralai.models.prediction import (
+        Prediction as MistralPrediction,
+        PredictionTypedDict as MistralPredictionTypedDict,
+    )
 
     from pydantic_ai.models.mistral import MistralModel, MistralStreamedResponse
     from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIResponsesModelSettings
@@ -1677,8 +1681,44 @@ async def get_location(loc_name: str) -> str:
 #####################
 ## Test methods
 #####################
+# --- _map_setting_prediction --------------------------------------------------
+@pytest.fixture
+def example_dict() -> MistralPredictionTypedDict:
+    """Fixture providing a typed dict for prediction."""
+    return {"type": "content", "content": "foo"}
+
 
+@pytest.fixture
+def example_prediction() -> MistralPrediction:
+    """Fixture providing a MistralPrediction object."""
+    return MistralPrediction(content="bar")
 
+
+@pytest.mark.parametrize(
+    "input_value,expected_content",
+    [
+        ("plain text", "plain text"),
+        ("example_prediction", "bar"),
+        ("example_dict", "foo"),
+        (None, None),
+    ],
+)
+def test_map_setting_prediction_valid(request, input_value, expected_content):
+    """
+    Accepted input types (str, dict, MistralPrediction, None) must be correctlyconverted to a MistralPrediction or None.
+    """
+    # If the parameter is a fixture name, resolve it using request
+    if isinstance(input_value, str) and input_value in {"example_dict", "example_prediction"}:
+        input_value = request.getfixturevalue(input_value)
+
+    result = MistralModel._map_setting_prediction(input_value)  # pyright: ignore[reportPrivateUsage]
+
+    if input_value is None:
+        assert result is None
+    else:
+        assert isinstance(result, MistralPrediction)
+        assert result.content == expected_content
+# -----------------------------------------------------
 def test_generate_user_output_format_complex(mistral_api_key: str):
     """
     Single test that includes properties exercising every branch
@@ -2263,3 +2303,22 @@ async def test_mistral_model_thinking_part_iter(allow_model_requests: None, mist
             ),
         ]
     )
+
+@pytest.mark.vcr()
+async def test_mistral_chat_with_prediction(allow_model_requests: None, mistral_api_key: str):
+    """Test chat completion with prediction parameter using a math query."""
+    from pydantic_ai.models.mistral import MistralModelSettings
+
+    model = MistralModel(
+        'mistral-small-latest',
+        provider=MistralProvider(api_key=mistral_api_key)
+    )
+    prediction = "The result of 21+21=99"
+    settings = MistralModelSettings(prediction=prediction)
+    agent = Agent(model, model_settings=settings)
+
+    result = await agent.run(['Correct only the math, respond with no explanation, no formatting.',"The result of 21+21=99"])
+
+    # Verify that the response uses the expected prediction
+    assert 'The result of 21+21=' in result.output
+    assert '42' in result.output