diff --git a/src/mistralai/extra/struct_chat.py b/src/mistralai/extra/struct_chat.py index d3fd3f5a..4023b758 100644 --- a/src/mistralai/extra/struct_chat.py +++ b/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() diff --git a/src/mistralai/extra/tests/test_struct_chat.py b/src/mistralai/extra/tests/test_struct_chat.py index 7b79bf77..789ce5fd 100644 --- a/src/mistralai/extra/tests/test_struct_chat.py +++ b/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()