refactor(models): give reasoning capture plumbing a single owner (#6639)

backend/packages/harness/deerflow/models/patched_mimo.py

@@ -18,6 +18,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResu
 from langchain_openai import ChatOpenAI
 
 from deerflow.models.assistant_payload_replay import restore_assistant_payloads, restore_reasoning_content
+from deerflow.models.reasoning_capture import typed_choice_message, with_reasoning_content
 
 _MISSING = object()
 
@@ -40,23 +41,6 @@ def _extract_reasoning_content(value: Any) -> str | object:
     return _MISSING
 
 
-def _with_reasoning_content(message: AIMessage | AIMessageChunk, reasoning: str) -> AIMessage | AIMessageChunk:
-    additional_kwargs = dict(message.additional_kwargs)
-    if additional_kwargs.get("reasoning_content") != reasoning:
-        additional_kwargs["reasoning_content"] = reasoning
-    return message.model_copy(update={"additional_kwargs": additional_kwargs})
-
-
-def _get_typed_choice_message(response: Any, index: int) -> Any:
-    choices = getattr(response, "choices", None)
-    if choices is None:
-        return None
-    try:
-        return choices[index].message
-    except (AttributeError, IndexError, TypeError):
-        return None
-
-
 class PatchedChatMiMo(ChatOpenAI):
     """ChatOpenAI with ``reasoning_content`` preservation for MiMo thinking mode."""
 
@@ -105,7 +89,7 @@ class PatchedChatMiMo(ChatOpenAI):
             reasoning = _extract_reasoning_content(delta)
             if reasoning is not _MISSING and isinstance(generation_chunk.message, AIMessageChunk):
                 generation_chunk = ChatGenerationChunk(
-                    message=_with_reasoning_content(generation_chunk.message, reasoning),
+                    message=with_reasoning_content(generation_chunk.message, reasoning),
                     generation_info=generation_chunk.generation_info,
                 )
 
@@ -126,14 +110,14 @@ class PatchedChatMiMo(ChatOpenAI):
             choice_message = choice.get("message", {}) if isinstance(choice, Mapping) else {}
             reasoning = _extract_reasoning_content(choice_message)
             if reasoning is _MISSING and not isinstance(response, dict):
-                reasoning = _extract_reasoning_content(_get_typed_choice_message(response, index))
+                reasoning = _extract_reasoning_content(typed_choice_message(response, index))
 
             message = generation.message
             if reasoning is not _MISSING and isinstance(message, AIMessage):
                 if patched_generations is None:
                     patched_generations = list(result.generations)
                 patched_generations[index] = ChatGeneration(
-                    message=_with_reasoning_content(message, reasoning),
+                    message=with_reasoning_content(message, reasoning),
                     generation_info=generation.generation_info,
                 )
 

backend/packages/harness/deerflow/models/patched_stepfun.py

@@ -21,6 +21,7 @@ from deerflow.models.assistant_payload_replay import (
     restore_assistant_payloads,
     restore_reasoning_content,
 )
+from deerflow.models.reasoning_capture import typed_choice_message, with_reasoning_content
 
 _MISSING = object()
 
@@ -54,25 +55,6 @@ def _extract_reasoning(value: Any) -> str | object:
     return _MISSING
 
 
-def _with_reasoning_content(message: AIMessage | AIMessageChunk, reasoning: str) -> AIMessage | AIMessageChunk:
-    """Return a copy of *message* with reasoning_content stored in additional_kwargs."""
-    additional_kwargs = dict(message.additional_kwargs)
-    if additional_kwargs.get("reasoning_content") != reasoning:
-        additional_kwargs["reasoning_content"] = reasoning
-    return message.model_copy(update={"additional_kwargs": additional_kwargs})
-
-
-def _get_typed_choice_message(response: Any, index: int) -> Any:
-    """Extract the SDK-typed choice message at *index*, if available."""
-    choices = getattr(response, "choices", None)
-    if choices is None:
-        return None
-    try:
-        return choices[index].message
-    except (AttributeError, IndexError, TypeError):
-        return None
-
-
 class PatchedChatStepFun(ChatOpenAI):
     """ChatOpenAI with full reasoning support for StepFun models.
 
@@ -133,7 +115,7 @@ class PatchedChatStepFun(ChatOpenAI):
             reasoning = _extract_reasoning(delta)
             if reasoning is not _MISSING and isinstance(generation_chunk.message, AIMessageChunk):
                 generation_chunk = ChatGenerationChunk(
-                    message=_with_reasoning_content(generation_chunk.message, reasoning),
+                    message=with_reasoning_content(generation_chunk.message, reasoning),
                     generation_info=generation_chunk.generation_info,
                 )
 
@@ -158,14 +140,14 @@ class PatchedChatStepFun(ChatOpenAI):
             reasoning = _extract_reasoning(choice_message)
 
             if reasoning is _MISSING and not isinstance(response, dict):
-                reasoning = _extract_reasoning(_get_typed_choice_message(response, index))
+                reasoning = _extract_reasoning(typed_choice_message(response, index))
 
             message = generation.message
             if reasoning is not _MISSING and isinstance(message, AIMessage):
                 if patched_generations is None:
                     patched_generations = list(result.generations)
                 patched_generations[index] = ChatGeneration(
-                    message=_with_reasoning_content(message, reasoning),
+                    message=with_reasoning_content(message, reasoning),
                     generation_info=generation.generation_info,
                 )
 

backend/packages/harness/deerflow/models/reasoning_capture.py

@@ -0,0 +1,56 @@
+"""Shared message plumbing for providers that preserve a reasoning string.
+
+``patched_mimo`` and ``patched_stepfun`` both wrap ``ChatOpenAI`` so a provider's
+``reasoning_content`` survives multi-turn tool calls, and each carried a
+byte-identical copy of the two mechanical steps that job needs: attach a
+reasoning string to a LangChain message without mutating the message the caller
+still holds, and read ``choices[index].message`` out of a response object that
+may not be shaped like that at all (dict payloads, SDK objects, or a chunk with
+no choices).
+
+How a provider *finds* its reasoning field stays local on purpose: MiMo reads
+``reasoning_content`` through ``_extract_reasoning_content``, StepFun reads both
+``reasoning`` and ``reasoning_content`` through ``_extract_reasoning``, and
+``patched_minimax`` merges streamed deltas. Those are provider dialects, not this
+rule. This module owns only what the two copies already agreed on, so a fix to
+the attach-without-mutating rule lands once instead of per provider.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from langchain_core.messages import AIMessage, AIMessageChunk
+
+
+def with_reasoning_content(message: AIMessage | AIMessageChunk, reasoning: str) -> AIMessage | AIMessageChunk:
+    """Return a copy of ``message`` carrying ``reasoning`` in ``additional_kwargs``.
+
+    The incoming message is never modified: LangChain hands the same object to
+    callbacks, run journals and the caller's history, so an in-place write would
+    leak reasoning into a message that was already emitted. A copy is returned
+    whether or not the value changed, keeping the result's identity predictable
+    for callers that replay messages through ``model_copy``.
+    """
+    additional_kwargs = dict(message.additional_kwargs)
+    if additional_kwargs.get("reasoning_content") != reasoning:
+        additional_kwargs["reasoning_content"] = reasoning
+    return message.model_copy(update={"additional_kwargs": additional_kwargs})
+
+
+def typed_choice_message(response: Any, index: int) -> Any:
+    """Return ``response.choices[index].message`` when that shape exists, else ``None``.
+
+    Non-streaming results arrive either as a dict payload (no attribute access)
+    or as an SDK object whose choice list may be shorter than the index the
+    caller is replaying. Neither case is an error for the callers here: they
+    fall back to the payload they already have, so this reports "nothing typed"
+    rather than raising.
+    """
+    choices = getattr(response, "choices", None)
+    if choices is None:
+        return None
+    try:
+        return choices[index].message
+    except (AttributeError, IndexError, TypeError):
+        return None

backend/tests/test_reasoning_capture_single_owner.py

@@ -0,0 +1,156 @@
+"""Ownership gate for the reasoning-capture plumbing MiMo and StepFun shared.
+
+``patched_mimo`` and ``patched_stepfun`` each carried a byte-identical copy of
+two helpers: attach a reasoning string to a LangChain message without mutating
+the message the caller still holds, and read ``choices[index].message`` out of a
+response that may be a dict payload, an SDK object, or a chunk with no choices.
+Two owners for one rule means a leak fix (or an identity fix) landing in one
+provider silently leaves the other behind.
+
+These tests pin three things the extraction must not silently give back:
+
+1. the owner's behaviour, case by case — including the two properties that are
+   easy to lose in a refactor: the input message is never mutated, and a copy is
+   returned even when the value already matched;
+2. that both providers actually call the owner (binding identity, plus the
+   streaming path exercised end to end), and the private helpers are gone;
+3. a source census over ``models/``: no module may restate either function body,
+   and exactly the two migrated providers import the owner.
+
+Providers whose reasoning handling genuinely differs are out of scope and stay
+local: ``patched_minimax`` merges streamed deltas, ``patched_deepseek`` and
+``vllm_provider`` read different field spellings, and each keeps its own
+extraction. ``test_the_census_compares_shapes_not_names`` pins that this gate
+does not sweep them in.
+"""
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from langchain_core.messages import AIMessage, AIMessageChunk
+
+import deerflow.models.patched_mimo as patched_mimo
+import deerflow.models.patched_minimax as patched_minimax
+import deerflow.models.patched_stepfun as patched_stepfun
+import deerflow.models.reasoning_capture as owner
+
+MODELS_ROOT = Path(owner.__file__).resolve().parent
+
+PROVIDERS = {
+    "patched_mimo": (patched_mimo, patched_mimo.PatchedChatMiMo, "mimo-v2.5-pro"),
+    "patched_stepfun": (patched_stepfun, patched_stepfun.PatchedChatStepFun, "step-3.7-flash"),
+}
+REASONING = "checked the tool output first"
+
+
+def _functions(tree: ast.Module) -> list[ast.FunctionDef | ast.AsyncFunctionDef]:
+    return [n for n in ast.walk(tree) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))]
+
+
+def _body_hash(fn: ast.AST) -> str:
+    """Hash the executable body, ignoring the docstring a provider may or may not carry."""
+    body = list(fn.body)
+    if body and isinstance(body[0], ast.Expr) and isinstance(getattr(body[0], "value", None), ast.Constant) and isinstance(body[0].value.value, str):
+        body = body[1:]
+    return ast.unparse(ast.Module(body=body, type_ignores=[]))
+
+
+def _owner_hashes() -> dict[str, str]:
+    return {_body_hash(fn): fn.name for fn in _functions(ast.parse(Path(owner.__file__).read_text(encoding="utf-8")))}
+
+
+# --- 1. the bar -----------------------------------------------------------------
+
+
+def test_with_reasoning_content_attaches_without_mutating() -> None:
+    message = AIMessage(content="answer", additional_kwargs={"other": "kept"})
+    patched = owner.with_reasoning_content(message, REASONING)
+
+    assert patched.additional_kwargs["reasoning_content"] == REASONING
+    assert patched.additional_kwargs["other"] == "kept"
+    assert "reasoning_content" not in message.additional_kwargs, "the caller's message was mutated"
+    assert patched is not message
+
+
+def test_with_reasoning_content_still_returns_a_copy_when_value_matches() -> None:
+    message = AIMessage(content="answer", additional_kwargs={"reasoning_content": REASONING})
+    patched = owner.with_reasoning_content(message, REASONING)
+
+    assert patched.additional_kwargs["reasoning_content"] == REASONING
+    assert patched is not message, "callers replay through model_copy; identity must stay predictable"
+
+
+def test_with_reasoning_content_overwrites_a_non_string_and_handles_chunks() -> None:
+    message = AIMessage(content="answer", additional_kwargs={"reasoning_content": None})
+    assert owner.with_reasoning_content(message, REASONING).additional_kwargs["reasoning_content"] == REASONING
+
+    chunk = AIMessageChunk(content="par")
+    patched = owner.with_reasoning_content(chunk, REASONING)
+    assert isinstance(patched, AIMessageChunk) and patched.additional_kwargs["reasoning_content"] == REASONING
+
+
+@pytest.mark.parametrize(
+    ("response", "index", "expected"),
+    [
+        (SimpleNamespace(choices=[SimpleNamespace(message="typed")]), 0, "typed"),
+        (SimpleNamespace(choices=[SimpleNamespace(message="typed")]), 1, None),
+        (SimpleNamespace(choices=None), 0, None),
+        ({"choices": [{"message": "payload shaped"}]}, 0, None),
+        (SimpleNamespace(choices=[SimpleNamespace()]), 0, None),
+        (SimpleNamespace(choices=7), 0, None),
+    ],
+)
+def test_typed_choice_message_reports_absence_rather_than_raising(response, index: int, expected) -> None:
+    assert owner.typed_choice_message(response, index) is expected
+
+
+# --- 2. both providers really go through the owner ----------------------------
+
+
+@pytest.mark.parametrize(("module", "cls", "model_name"), sorted(PROVIDERS.values(), key=lambda x: x[0].__name__))
+def test_provider_names_are_the_owner_not_a_restatement(module, cls, model_name: str) -> None:
+    assert module.with_reasoning_content is owner.with_reasoning_content
+    assert module.typed_choice_message is owner.typed_choice_message
+    for stale in ("_with_reasoning_content", "_get_typed_choice_message"):
+        assert not hasattr(module, stale), f"{module.__name__} still carries its private {stale}"
+
+
+@pytest.mark.parametrize(("module", "cls", "model_name"), sorted(PROVIDERS.values(), key=lambda x: x[0].__name__))
+def test_streaming_path_still_captures_reasoning_through_the_owner(module, cls, model_name: str) -> None:
+    """Behaviour at a real call site, not just an identity check."""
+    model = cls(model=model_name, api_key="test-key", base_url="https://example.invalid/v1")
+    chunk = model._convert_chunk_to_generation_chunk(
+        {"choices": [{"delta": {"role": "assistant", "reasoning_content": REASONING}}]},
+        AIMessageChunk,
+        {},
+    )
+    assert chunk is not None
+    assert chunk.message.additional_kwargs["reasoning_content"] == REASONING
+
+
+# --- 3. census over models/ ---------------------------------------------------
+
+
+def test_no_model_module_restates_an_owned_body() -> None:
+    hashes = _owner_hashes()
+    offenders = []
+    for path in sorted(MODELS_ROOT.glob("*.py")):
+        if path.name == "reasoning_capture.py":
+            continue
+        for fn in _functions(ast.parse(path.read_text(encoding="utf-8"), filename=str(path))):
+            if _body_hash(fn) in hashes:
+                offenders.append(f"{path.name}::{fn.name}")
+    assert offenders == [], f"shared reasoning bodies belong to the owner: found in {offenders}"
+
+
+def test_the_census_compares_shapes_not_names() -> None:
+    """MiMo/StepFun's near-neighbour keeps its own merge logic and must stay untouched."""
+    minimax = next(fn for fn in _functions(ast.parse(Path(patched_minimax.__file__).read_text(encoding="utf-8"))) if fn.name == "_with_reasoning_content")
+    assert _body_hash(minimax) not in _owner_hashes(), "the census would be claiming a fold that never happened"
+
+
+def test_exactly_the_migrated_providers_import_the_owner() -> None:
+    importers = {path.name for path in sorted(MODELS_ROOT.glob("*.py")) if path.name != "reasoning_capture.py" and "from deerflow.models.reasoning_capture import" in path.read_text(encoding="utf-8")}
+    assert importers == {"patched_mimo.py", "patched_stepfun.py"}