Skip to content

Commit 3653cd9

Browse files
authored
fix(chat): parse structured output when reasoning returns chunked content (#621)
With reasoning_effort enabled, the API returns message content as a list of chunks [ThinkChunk, TextChunk] instead of a plain JSON string. convert_to_parsed_chat_completion_response only handled str and None, so parse()/parse_async() raised TypeError. Concatenate the text chunks and parse the result. Content with only thinking chunks yields parsed=None.
1 parent 878fdda commit 3653cd9

2 files changed

Lines changed: 83 additions & 1 deletion

File tree

‎src/mistralai/extra/struct_chat.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
11
import json
22
from typing import Generic
33

4-
from mistralai.client.models import AssistantMessage, ChatCompletionChoice, ChatCompletionResponse
4+
from mistralai.client.models import (
5+
AssistantMessage,
6+
ChatCompletionChoice,
7+
ChatCompletionResponse,
8+
TextChunk,
9+
)
510
from .utils.response_format import CustomPydanticModel, pydantic_model_from_json
611

712

@@ -34,6 +39,16 @@ def convert_to_parsed_chat_completion_response(
3439
parsed_message.parsed = pydantic_model_from_json(json.loads(parsed_message.content), response_format)
3540
elif parsed_message.content is None:
3641
parsed_message.parsed = None
42+
elif isinstance(parsed_message.content, list):
43+
final_text = "".join(
44+
chunk.text
45+
for chunk in parsed_message.content
46+
if isinstance(chunk, TextChunk)
47+
)
48+
if not final_text:
49+
parsed_message.parsed = None
50+
else:
51+
parsed_message.parsed = pydantic_model_from_json(json.loads(final_text), response_format)
3752
else:
3853
raise TypeError(f"Unexpected type for message.content: {type(parsed_message.content)}")
3954
choice_dict = choice.model_dump()

‎src/mistralai/extra/tests/test_struct_chat.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
UsageInfo,
1111
ChatCompletionChoice,
1212
AssistantMessage,
13+
ThinkChunk,
14+
TextChunk,
1315
)
1416
from pydantic import BaseModel
1517

@@ -98,6 +100,71 @@ def test_convert_to_parsed_chat_completion_response(self):
98100
)
99101
self.assertEqual(output, expected_response)
100102

103+
def test_convert_to_parsed_chat_completion_response_with_reasoning_chunks(self):
104+
reasoning_response = ChatCompletionResponse(
105+
id="chunked-response",
106+
object="chat.completion",
107+
model="mistral-medium-3-5",
108+
usage=UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
109+
created=1737727558,
110+
choices=[
111+
ChatCompletionChoice(
112+
index=0,
113+
message=AssistantMessage(
114+
content=[
115+
ThinkChunk(
116+
thinking=[
117+
TextChunk(text="Compute 8x + 7 = -23 step by step.")
118+
]
119+
),
120+
TextChunk(
121+
text='{"steps": [], "final_answer": "x = -4"}'
122+
),
123+
],
124+
role="assistant",
125+
),
126+
finish_reason="stop",
127+
)
128+
],
129+
)
130+
output = convert_to_parsed_chat_completion_response(
131+
reasoning_response, MathDemonstration
132+
)
133+
assert output.choices is not None
134+
assert output.choices[0].message is not None
135+
self.assertEqual(output.choices[0].message.parsed, MathDemonstration(steps=[], final_answer="x = -4"))
136+
137+
def test_convert_to_parsed_chat_completion_response_with_only_reasoning_chunks(self):
138+
reasoning_only_response = ChatCompletionResponse(
139+
id="reasoning-only-response",
140+
object="chat.completion",
141+
model="mistral-medium-3-5",
142+
usage=UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
143+
created=1737727558,
144+
choices=[
145+
ChatCompletionChoice(
146+
index=0,
147+
message=AssistantMessage(
148+
content=[
149+
ThinkChunk(
150+
thinking=[
151+
TextChunk(text="Still reasoning about the answer.")
152+
]
153+
),
154+
],
155+
role="assistant",
156+
),
157+
finish_reason="stop",
158+
)
159+
],
160+
)
161+
output = convert_to_parsed_chat_completion_response(
162+
reasoning_only_response, MathDemonstration
163+
)
164+
assert output.choices is not None
165+
assert output.choices[0].message is not None
166+
self.assertIsNone(output.choices[0].message.parsed)
167+
101168

102169
if __name__ == "__main__":
103170
unittest.main()

0 commit comments

Comments
 (0)