Mistral optimised (#396)

Author: YanSte

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -8,6 +8,7 @@
 from itertools import chain
 from typing import Any, Callable, Literal, Union
 
+import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
 from typing_extensions import assert_never
 
@@ -39,7 +40,6 @@
 )
 
 try:
-    from json_repair import repair_json
     from mistralai import (
         UNSET,
         CompletionChunk as MistralCompletionChunk,
@@ -198,11 +198,10 @@ async def _stream_completions_create(
         """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)))
-
         model_settings = model_settings or {}
 
         if self.result_tools and self.function_tools or self.function_tools:
-            # Function Calling Mode
+            # Function Calling
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -218,9 +217,9 @@ async def _stream_completions_create(
         elif self.result_tools:
             # Json Mode
             parameters_json_schemas = [tool.parameters_json_schema for tool in self.result_tools]
-
             user_output_format_message = self._generate_user_output_format(parameters_json_schemas)
             mistral_messages.append(user_output_format_message)
+
             response = await self.client.chat.stream_async(
                 model=str(self.model_name),
                 messages=mistral_messages,
@@ -270,12 +269,13 @@ def _map_function_and_result_tools_definition(self) -> list[MistralTool] | None:
     @staticmethod
     def _process_response(response: MistralChatCompletionResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
+        assert response.choices, 'Unexpected empty response choice.'
+
         if response.created:
             timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         else:
             timestamp = _now_utc()
 
-        assert response.choices, 'Unexpected empty response choice.'
         choice = response.choices[0]
         content = choice.message.content
         tool_calls = choice.message.tool_calls
@@ -546,20 +546,15 @@ def get(self, *, final: bool = False) -> ModelResponse:
                 calls.append(tool)
 
         elif self._delta_content and self._result_tools:
-            # NOTE: Params set for the most efficient and fastest way.
-            output_json = repair_json(self._delta_content, return_objects=True, skip_json_loads=True)
-            assert isinstance(
-                output_json, dict
-            ), f'Expected repair_json as type dict, invalid type: {type(output_json)}'
+            output_json: dict[str, Any] | None = pydantic_core.from_json(
+                self._delta_content, allow_partial='trailing-strings'
+            )
 
             if output_json:
                 for result_tool in self._result_tools.values():
-                    # NOTE: Additional verification to prevent JSON validation to crash in `result.py`
+                    # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
                     # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
-                    # For example, `return_type=list[str]` expects a 'response' key with value type array of str.
-                    # when `{"response":` then `repair_json` sets `{"response": ""}` (type not found default str)
-                    # when `{"response": {` then `repair_json` sets `{"response": {}}` (type found)
-                    # This ensures it's corrected to `{"response": {}}` and other required parameters and type.
+                    # Example with BaseModel and required fields.
                     if not self._validate_required_json_schema(output_json, result_tool.parameters_json_schema):
                         continue
 

pydantic_ai_slim/pyproject.toml

@@ -46,7 +46,7 @@ openai = ["openai>=1.54.3"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.3"]
 anthropic = ["anthropic>=0.40.0"]
 groq = ["groq>=0.12.0"]
-mistral = ["mistralai>=1.2.5", "json-repair>=0.30.3"]
+mistral = ["mistralai>=1.2.5"]
 logfire = ["logfire>=2.3"]
 
 [dependency-groups]

tests/models/test_mistral.py

@@ -544,7 +544,7 @@ async def test_request_result_type_with_arguments_str_response(allow_model_reque
 #####################
 
 
-async def test_stream_structured_with_all_typd(allow_model_requests: None):
+async def test_stream_structured_with_all_type(allow_model_requests: None):
     class MyTypedDict(TypedDict, total=False):
         first: str
         second: int
@@ -563,19 +563,19 @@ class MyTypedDict(TypedDict, total=False):
             '", "second": 2',
         ),
         text_chunk(
-            '", "bool_value": true',
+            ', "bool_value": true',
         ),
         text_chunk(
-            '", "nullable_value": null',
+            ', "nullable_value": null',
         ),
         text_chunk(
-            '", "array_value": ["A", "B", "C"]',
+            ', "array_value": ["A", "B", "C"]',
         ),
         text_chunk(
-            '", "dict_value": {"A": "A", "B":"B"}',
+            ', "dict_value": {"A": "A", "B":"B"}',
         ),
         text_chunk(
-            '", "dict_int_value": {"A": 1, "B":2}',
+            ', "dict_int_value": {"A": 1, "B":2}',
         ),
         text_chunk('}'),
         chunk([]),
@@ -721,8 +721,8 @@ class MyTypedDict(TypedDict, total=False):
                 {'first': 'One'},
                 {'first': 'One'},
                 {'first': 'One'},
-                {'first': 'One', 'second': ''},
-                {'first': 'One', 'second': ''},
+                {'first': 'One'},
+                {'first': 'One'},
                 {'first': 'One', 'second': ''},
                 {'first': 'One', 'second': 'T'},
                 {'first': 'One', 'second': 'Tw'},
@@ -828,20 +828,21 @@ async def test_stream_result_type_primitif_array(allow_model_requests: None):
         v = [c async for c in result.stream(debounce_by=None)]
         assert v == snapshot(
             [
+                [''],
                 ['f'],
                 ['fi'],
                 ['fir'],
                 ['firs'],
                 ['first'],
                 ['first'],
                 ['first'],
-                ['first'],
+                ['first', ''],
                 ['first', 'O'],
                 ['first', 'On'],
                 ['first', 'One'],
                 ['first', 'One'],
                 ['first', 'One'],
-                ['first', 'One'],
+                ['first', 'One', ''],
                 ['first', 'One', 's'],
                 ['first', 'One', 'se'],
                 ['first', 'One', 'sec'],
@@ -850,7 +851,7 @@ async def test_stream_result_type_primitif_array(allow_model_requests: None):
                 ['first', 'One', 'second'],
                 ['first', 'One', 'second'],
                 ['first', 'One', 'second'],
-                ['first', 'One', 'second'],
+                ['first', 'One', 'second', ''],
                 ['first', 'One', 'second', 'T'],
                 ['first', 'One', 'second', 'Tw'],
                 ['first', 'One', 'second', 'Two'],
@@ -869,10 +870,10 @@ async def test_stream_result_type_primitif_array(allow_model_requests: None):
         assert result.usage().response_tokens == len(stream)
 
 
-async def test_stream_result_type_basemodel(allow_model_requests: None):
+async def test_stream_result_type_basemodel_with_default_params(allow_model_requests: None):
     class MyTypedBaseModel(BaseModel):
-        first: str = ''  # Note: Don't forget to set default values
-        second: str = ''
+        first: str = ''  # Note: Default, set value.
+        second: str = ''  # Note: Default, set value.
 
     # Given
     stream = [
@@ -958,6 +959,79 @@ class MyTypedBaseModel(BaseModel):
         assert result.usage().response_tokens == len(stream)
 
 
+async def test_stream_result_type_basemodel_with_required_params(allow_model_requests: None):
+    class MyTypedBaseModel(BaseModel):
+        first: str  # Note: Required params
+        second: str  # Note: Required params
+
+    # Given
+    stream = [
+        text_chunk('{'),
+        text_chunk('"'),
+        text_chunk('f'),
+        text_chunk('i'),
+        text_chunk('r'),
+        text_chunk('s'),
+        text_chunk('t'),
+        text_chunk('"'),
+        text_chunk(':'),
+        text_chunk(' '),
+        text_chunk('"'),
+        text_chunk('O'),
+        text_chunk('n'),
+        text_chunk('e'),
+        text_chunk('"'),
+        text_chunk(','),
+        text_chunk(' '),
+        text_chunk('"'),
+        text_chunk('s'),
+        text_chunk('e'),
+        text_chunk('c'),
+        text_chunk('o'),
+        text_chunk('n'),
+        text_chunk('d'),
+        text_chunk('"'),
+        text_chunk(':'),
+        text_chunk(' '),
+        text_chunk('"'),
+        text_chunk('T'),
+        text_chunk('w'),
+        text_chunk('o'),
+        text_chunk('"'),
+        text_chunk('}'),
+        chunk([]),
+    ]
+
+    mock_client = MockMistralAI.create_stream_mock(stream)
+    model = MistralModel('mistral-large-latest', client=mock_client)
+    agent = Agent(model=model, result_type=MyTypedBaseModel)
+
+    # When
+    async with agent.run_stream('User prompt value') as result:
+        # Then
+        assert result.is_structured
+        assert not result.is_complete
+        v = [c async for c in result.stream(debounce_by=None)]
+        assert v == snapshot(
+            [
+                MyTypedBaseModel(first='One', second=''),
+                MyTypedBaseModel(first='One', second='T'),
+                MyTypedBaseModel(first='One', second='Tw'),
+                MyTypedBaseModel(first='One', second='Two'),
+                MyTypedBaseModel(first='One', second='Two'),
+                MyTypedBaseModel(first='One', second='Two'),
+                MyTypedBaseModel(first='One', second='Two'),
+            ]
+        )
+        assert result.is_complete
+        assert result.usage().request_tokens == 34
+        assert result.usage().response_tokens == 34
+        assert result.usage().total_tokens == 34
+
+        # double check cost matches stream count
+        assert result.usage().response_tokens == len(stream)
+
+
 #####################
 ## Completion Function call
 #####################
@@ -1693,6 +1767,6 @@ def test_generate_user_output_format_multiple():
         ),
     ],
 )
-def test_validate_required_json_shema(desc: str, schema: dict[str, Any], data: dict[str, Any], expected: bool) -> None:
+def test_validate_required_json_schema(desc: str, schema: dict[str, Any], data: dict[str, Any], expected: bool) -> None:
     result = MistralStreamStructuredResponse._validate_required_json_schema(data, schema)  # pyright: ignore[reportPrivateUsage]
     assert result == expected, f'{desc} — expected {expected}, got {result}'