fix(chat): parse structured output when reasoning returns chunked content

src/mistralai/extra/struct_chat.py

@@ -1,7 +1,12 @@
 import json
 from typing import Generic
 
-from mistralai.client.models import AssistantMessage, ChatCompletionChoice, ChatCompletionResponse
+from mistralai.client.models import (
+    AssistantMessage,
+    ChatCompletionChoice,
+    ChatCompletionResponse,
+    TextChunk,
+)
 from .utils.response_format import CustomPydanticModel, pydantic_model_from_json
 
 
@@ -34,6 +39,16 @@ def convert_to_parsed_chat_completion_response(
                     parsed_message.parsed = pydantic_model_from_json(json.loads(parsed_message.content), response_format)
                 elif parsed_message.content is None:
                     parsed_message.parsed = None
+                elif isinstance(parsed_message.content, list):
+                    final_text = "".join(
+                        chunk.text
+                        for chunk in parsed_message.content
+                        if isinstance(chunk, TextChunk)
+                    )
+                    if not final_text:
+                        parsed_message.parsed = None
+                    else:
+                        parsed_message.parsed = pydantic_model_from_json(json.loads(final_text), response_format)
                 else:
                     raise TypeError(f"Unexpected type for message.content: {type(parsed_message.content)}")
                 choice_dict = choice.model_dump()

src/mistralai/extra/tests/test_struct_chat.py

@@ -10,6 +10,8 @@
     UsageInfo,
     ChatCompletionChoice,
     AssistantMessage,
+    ThinkChunk,
+    TextChunk,
 )
 from pydantic import BaseModel
 
@@ -98,6 +100,71 @@ def test_convert_to_parsed_chat_completion_response(self):
         )
         self.assertEqual(output, expected_response)
 
+    def test_convert_to_parsed_chat_completion_response_with_reasoning_chunks(self):
+        reasoning_response = ChatCompletionResponse(
+            id="chunked-response",
+            object="chat.completion",
+            model="mistral-medium-3-5",
+            usage=UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
+            created=1737727558,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message=AssistantMessage(
+                        content=[
+                            ThinkChunk(
+                                thinking=[
+                                    TextChunk(text="Compute 8x + 7 = -23 step by step.")
+                                ]
+                            ),
+                            TextChunk(
+                                text='{"steps": [], "final_answer": "x = -4"}'
+                            ),
+                        ],
+                        role="assistant",
+                    ),
+                    finish_reason="stop",
+                )
+            ],
+        )
+        output = convert_to_parsed_chat_completion_response(
+            reasoning_response, MathDemonstration
+        )
+        assert output.choices is not None
+        assert output.choices[0].message is not None
+        self.assertEqual(output.choices[0].message.parsed, MathDemonstration(steps=[], final_answer="x = -4"))
+
+    def test_convert_to_parsed_chat_completion_response_with_only_reasoning_chunks(self):
+        reasoning_only_response = ChatCompletionResponse(
+            id="reasoning-only-response",
+            object="chat.completion",
+            model="mistral-medium-3-5",
+            usage=UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
+            created=1737727558,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message=AssistantMessage(
+                        content=[
+                            ThinkChunk(
+                                thinking=[
+                                    TextChunk(text="Still reasoning about the answer.")
+                                ]
+                            ),
+                        ],
+                        role="assistant",
+                    ),
+                    finish_reason="stop",
+                )
+            ],
+        )
+        output = convert_to_parsed_chat_completion_response(
+            reasoning_only_response, MathDemonstration
+        )
+        assert output.choices is not None
+        assert output.choices[0].message is not None
+        self.assertIsNone(output.choices[0].message.parsed)
+
 
 if __name__ == "__main__":
     unittest.main()