Skip to content

fix(server): keep the visible answer of a thinking model on the MLX chat path #736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
krisztian-gajdar merged 2 commits into from
Oct 14, 2026
Merged

fix(server): keep the visible answer of a thinking model on the MLX chat path #736

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
75c83cf
fix(server): keep the visible answer of a thinking model on the MLX c…
Oct 14, 2026
f34ff02
fix(server): hold content beside split reasoning until the block clos…
Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 33 additions & 8 deletions packages/sie_server/src/sie_server/api/openai_local.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
from sie_sdk.queue_types import denormalize_model_id

from sie_server.adapters._generation_base import (
REASONING_CONTROL_TOKENS,
GenerationAdapter,
GenerationError,
GenerationUnsupportedFieldError,
Expand Down Expand Up @@ -131,6 +132,7 @@
_MAX_CHAT_VIDEOS = 1
_MAX_U32 = (1 << 32) - 1
_ALLOWED_CHAT_ROLES = frozenset({"system", "user", "assistant", "tool", "developer"})
_ANSWER_FINISH_REASONS = frozenset({"stop", "tool_calls"})
# Compatibility requests project onto the same native score bound.
_MAX_RERANK_DOCS = MAX_SCORE_ITEMS

Expand Down Expand Up @@ -1004,6 +1006,7 @@ def _sanitize_chat_payload(
initial_inside_thinking: bool,
reasoning_format: ReasoningFormat,
states: dict[int, ThinkingBlockStripper] | None = None,
held_contents: dict[int, list[str]] | None = None,
terminal: bool,
include_reasoning: bool = False,
) -> dict[str, Any]:
Expand All @@ -1016,6 +1019,7 @@ def _sanitize_chat_payload(

public_reasoning = include_reasoning and hide_thinking_blocks
state_map = states if states is not None else {}
held_map = held_contents if held_contents is not None else {}
for position, choice in enumerate(choices):
if not isinstance(choice, dict):
continue
Expand All @@ -1037,18 +1041,37 @@ def _sanitize_chat_payload(

if not hide_thinking_blocks:
continue
stripper = state_map.setdefault(
index,
ThinkingBlockStripper(
raw_content = container.get("content")
content = raw_content if isinstance(raw_content, str) else ""
stripper = state_map.get(index)
if stripper is None:
if not content and not reasoning_was_present:
continue
stripper = state_map[index] = ThinkingBlockStripper(
inside=initial_inside_thinking,
reasoning_format=reasoning_format,
keep_reasoning=public_reasoning,
),
)
raw_content = container.get("content")
content = raw_content if isinstance(raw_content, str) else ""
visible = stripper.feed(content) if content else ""
)
if initial_inside_thinking and reasoning_was_present:
held_map[index] = []
is_terminal_choice = terminal or choice_obj.get("finish_reason") is not None
text = content
held = held_map.get(index)
if held is not None:
close = REASONING_CONTROL_TOKENS[reasoning_format][1]
closes = close in "".join(held[1 - len(close) :]) + content
if content:
held.append(content)
if closes or is_terminal_choice:
del held_map[index]
text = "".join(held)
if not closes and choice_obj.get("finish_reason") in _ANSWER_FINISH_REASONS:
stripper = state_map[index] = ThinkingBlockStripper(
reasoning_format=reasoning_format, keep_reasoning=public_reasoning
)
else:
text = ""
visible = stripper.feed(text) if text else ""
if is_terminal_choice:
visible += stripper.finish()
if isinstance(raw_content, str) or visible:
Expand Down Expand Up @@ -1093,6 +1116,7 @@ async def _sanitize_sse_stream(
include_reasoning: bool = False,
) -> AsyncIterator[bytes]:
states: dict[int, ThinkingBlockStripper] = {}
held_contents: dict[int, list[str]] = {}
strict_contents: dict[int, list[str]] = {}
strict_tool_choices: set[int] = set()
tool_calls: dict[int, set[object]] = {}
Expand Down Expand Up @@ -1135,6 +1159,7 @@ async def _sanitize_sse_stream(
initial_inside_thinking=initial_inside_thinking,
reasoning_format=reasoning_format,
states=states,
held_contents=held_contents,
terminal=False,
include_reasoning=include_reasoning,
)
Expand Down
240 changes: 240 additions & 0 deletions packages/sie_server/tests/api/test_openai_local.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2418,3 +2418,243 @@ def _handler(request: httpx.Request) -> httpx.Response:
reasoning = "".join(delta.get("reasoning_content", "") for delta in deltas)
assert reasoning == ("private default reasoning" if include_reasoning else "")
assert "<think>" not in response.text


_MLX_REASONING = "The user wants one word.\n"
_MLX_ANSWER = "\n\nready"


@pytest.mark.parametrize("include_reasoning", [False, True])
@pytest.mark.parametrize(
("message", "finish_reason", "content", "reasoning", "hidden"),
[
(
{"role": "assistant", "content": _MLX_ANSWER, "reasoning": _MLX_REASONING},
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
{"role": "assistant", "content": _MLX_ANSWER, "reasoning_content": _MLX_REASONING},
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
{"role": "assistant", "content": _MLX_REASONING + "</think>" + _MLX_ANSWER},
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
{"role": "assistant", "content": "one word.\n</think>" + _MLX_ANSWER, "reasoning": "The user wants "},
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
{"role": "assistant", "content": "ready<think>one word", "reasoning": _MLX_REASONING},
"stop",
"ready",
_MLX_REASONING + "one word",
"one word",
),
(
{"role": "assistant", "content": "still thinking", "reasoning": _MLX_REASONING},
"length",
"",
_MLX_REASONING + "still thinking",
"still thinking",
),
({"role": "assistant", "content": _MLX_ANSWER}, "stop", "", _MLX_ANSWER, "ready"),
],
ids=[
"reasoning",
"reasoning_content",
"raw_tags",
"split_and_raw_close",
"split_then_raw_open",
"split_cut_off",
"no_reasoning_field",
],
)
def test_mlx_thinking_chat_keeps_the_answer_beside_split_reasoning(
monkeypatch: pytest.MonkeyPatch,
message: dict[str, Any],
finish_reason: str,
content: str,
reasoning: str,
hidden: str,
include_reasoning: bool,
) -> None:
payload = {
"id": "chatcmpl-mlx",
"system_fingerprint": "0.32.0",
"object": "chat.completion",
"model": "mlx-community/Qwen3.5-4B-4bit",
"created": 1,
"choices": [{"index": 0, "finish_reason": finish_reason, "message": message}],
"usage": {"prompt_tokens": 14, "completion_tokens": 12, "total_tokens": 26},
}
client, _ = _mlx_chat_client(
monkeypatch, lambda _request: httpx.Response(200, json=payload), chat_template_kwargs={"enable_thinking": True}
)

response = client.post(
"/v1/chat/completions",
json={
"model": "Qwen/Qwen3.5-4B:thinking",
"messages": [{"role": "user", "content": "Reply with one word: ready"}],
"include_reasoning": include_reasoning,
},
)

assert response.status_code == 200
returned = response.json()["choices"][0]["message"]
assert returned["content"] == content
assert "reasoning" not in returned
assert returned.get("reasoning_content") == (reasoning if include_reasoning else None)
assert include_reasoning or hidden not in response.text


@pytest.mark.parametrize("include_reasoning", [False, True])
@pytest.mark.parametrize(
("deltas", "finish_reason", "content", "reasoning", "hidden"),
[
(
[
{"role": "assistant", "reasoning": "The user wants "},
{"role": "assistant", "reasoning": "one word.\n"},
{"role": "assistant", "content": "\n\n"},
{"role": "assistant", "content": "ready"},
],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[
{"role": "assistant", "content": ""},
{"reasoning_content": "The user wants "},
{"reasoning_content": "one word.\n"},
{"content": "\n\nready"},
],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[
{"role": "assistant", "content": "The user wants "},
{"role": "assistant", "content": "one word.\n</th"},
{"role": "assistant", "content": "ink>\n\nready"},
],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[
{"role": "assistant", "content": ""},
{"reasoning_content": "The user wants "},
{"content": "one word.\n</think>\n\nready"},
],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[
{"role": "assistant", "reasoning": "The user wants "},
{"role": "assistant", "content": "one word.\n</th"},
{"role": "assistant", "content": "ink>\n\nready"},
],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[{"role": "assistant", "reasoning": _MLX_REASONING, "content": _MLX_ANSWER}],
"stop",
_MLX_ANSWER,
_MLX_REASONING,
"one word",
),
(
[
{"role": "assistant", "reasoning": _MLX_REASONING},
{"role": "assistant", "content": "Checking."},
{"role": "assistant", "tool_calls": [{"index": 0, **_PARIS_CALL}]},
],
"tool_calls",
"Checking.",
_MLX_REASONING,
"one word",
),
(
[{"role": "assistant", "reasoning": _MLX_REASONING}, {"role": "assistant", "content": "still thinking"}],
"length",
"",
_MLX_REASONING + "still thinking",
"still thinking",
),
([{"role": "assistant", "content": _MLX_ANSWER}], "stop", "", _MLX_ANSWER, "ready"),
],
ids=[
"reasoning",
"reasoning_content_after_role",
"raw_tags",
"split_then_raw_close",
"split_then_raw_close_across_deltas",
"split_and_answer_in_one_delta",
"answer_before_tool_call",
"split_cut_off",
"no_reasoning_field",
],
)
def test_mlx_thinking_chat_stream_keeps_the_answer_beside_split_reasoning(
monkeypatch: pytest.MonkeyPatch,
deltas: list[dict[str, Any]],
finish_reason: str,
content: str,
reasoning: str,
hidden: str,
include_reasoning: bool,
) -> None:
wire = _sse_wire(
[*(_stream_event(delta) for delta in deltas), _stream_event({"role": "assistant"}, finish_reason=finish_reason)]
)
client, _ = _mlx_chat_client(
monkeypatch,
lambda _request: httpx.Response(
200, stream=_ChunkedAsyncStream([wire[:83], wire[83:]]), headers={"content-type": "text/event-stream"}
),
chat_template_kwargs={"enable_thinking": True},
)

response = client.post(
"/v1/chat/completions",
json={
"model": "Qwen/Qwen3.5-4B:thinking",
"messages": [{"role": "user", "content": "Reply with one word: ready"}],
"tools": [_WEATHER_TOOL],
"stream": True,
"include_reasoning": include_reasoning,
},
)

assert response.status_code == 200
returned = _sse_deltas(response.text)
assert "".join(delta.get("content") or "" for delta in returned) == content
assert not any("reasoning" in delta for delta in returned)
assert "".join(delta.get("reasoning_content", "") for delta in returned) == (reasoning if include_reasoning else "")
assert include_reasoning or hidden not in response.text
Loading