Update tests and model documentation for revised Mistral prediction handling.

Author: TKaluza

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -122,8 +122,7 @@ class MistralModelSettings(ModelSettings, total=False):
     """Prediction content for the model to use as a prefix. See Predictive outputs.
 
     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.
+    As of now, codestral-latest and mistral-large-2411 support [predicted outputs](https://docs.mistral.ai/capabilities/predicted_outputs). 
     """
 
 

tests/models/cassettes/test_mistral/test_mistral_chat_with_prediction.yaml

@@ -8,7 +8,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '246'
+      - '315'
       content-type:
       - application/json
       host:
@@ -17,13 +17,16 @@ interactions:
     parsed_body:
       messages:
       - content:
-        - text: Correct only the math, respond with no explanation, no formatting.
+        - text: Correct the math, keep everything else. No explanation, no formatting.
           type: text
         - text: The result of 21+21=99
           type: text
         role: user
-      model: mistral-small-latest
+      model: mistral-large-2411
       n: 1
+      prediction:
+        content: The result of 21+21=99
+        type: content
       stream: false
       top_p: 1.0
     uri: https://api.mistral.ai/v1/chat/completions
@@ -36,11 +39,11 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '321'
+      - '319'
       content-type:
       - application/json
       mistral-correlation-id:
-      - 019a639b-7bf4-7481-96af-078cd1a7d277
+      - 019a63b7-40ba-70cb-94d0-84f036d7c76f
       strict-transport-security:
       - max-age=15552000; includeSubDomains; preload
       transfer-encoding:
@@ -53,14 +56,14 @@ interactions:
           content: The result of 21+21=42
           role: assistant
           tool_calls: null
-      created: 1762607725
-      id: a7952046ef794d1697627b54231df17a
-      model: mistral-small-latest
+      created: 1762609545
+      id: 6c36e8b6c3c145bd8ada32f9bd0f6be9
+      model: mistral-large-2411
       object: chat.completion
       usage:
         completion_tokens: 13
-        prompt_tokens: 28
-        total_tokens: 41
+        prompt_tokens: 33
+        total_tokens: 46
     status:
       code: 200
       message: OK

tests/models/test_mistral.py

@@ -53,11 +53,11 @@
         SDKError,
         ToolCall as MistralToolCall,
     )
-    from mistralai.types.basemodel import Unset as MistralUnset
     from mistralai.models.prediction import (
         Prediction as MistralPrediction,
         PredictionTypedDict as MistralPredictionTypedDict,
     )
+    from mistralai.types.basemodel import Unset as MistralUnset
 
     from pydantic_ai.models.mistral import MistralModel, MistralStreamedResponse
     from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIResponsesModelSettings
@@ -1681,44 +1681,45 @@ 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"}
+    return {'type': 'content', 'content': 'foo'}
 
 
 @pytest.fixture
 def example_prediction() -> MistralPrediction:
     """Fixture providing a MistralPrediction object."""
-    return MistralPrediction(content="bar")
+    return MistralPrediction(content='bar')
 
 
 @pytest.mark.parametrize(
-    "input_value,expected_content",
+    'input_value,expected_content',
     [
-        ("plain text", "plain text"),
-        ("example_prediction", "bar"),
-        ("example_dict", "foo"),
+        ('plain text', 'plain text'),
+        ('example_prediction', 'bar'),
+        ('example_dict', 'foo'),
         (None, None),
     ],
 )
-def test_map_setting_prediction_valid(request, input_value, expected_content):
+def test_map_setting_prediction_valid(request: pytest.FixtureRequest, input_value: str, expected_content: str | None):
     """
     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)
+    resolved_value: str | MistralPredictionTypedDict | MistralPrediction | None = input_value
+    if isinstance(input_value, str) and input_value in {'example_dict', 'example_prediction'}:
+        resolved_value = request.getfixturevalue(input_value)
 
-    result = MistralModel._map_setting_prediction(input_value)  # pyright: ignore[reportPrivateUsage]
+    result = MistralModel._map_setting_prediction(resolved_value)  # pyright: ignore[reportPrivateUsage]
 
-    if input_value is None:
+    if resolved_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
@@ -2304,21 +2305,19 @@ 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)
+    model = MistralModel('mistral-large-2411', provider=MistralProvider(api_key=mistral_api_key))
+    prediction = 'The result of 21+21=99'
+    settings = MistralModelSettings(mistral_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"])
+    result = await agent.run(['Correct the math, keep everything else. No explanation, no formatting.', prediction])
 
     # Verify that the response uses the expected prediction
-    assert 'The result of 21+21=' in result.output
     assert '42' in result.output
+    assert 'The result of 21+21=' in result.output