diff --git a/backend/packages/harness/deerflow/models/patched_mimo.py b/backend/packages/harness/deerflow/models/patched_mimo.py index 3851748a46..8cb1d4ddf2 100644 --- a/backend/packages/harness/deerflow/models/patched_mimo.py +++ b/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, ) diff --git a/backend/packages/harness/deerflow/models/patched_stepfun.py b/backend/packages/harness/deerflow/models/patched_stepfun.py index 1a30332a28..281afa4c7a 100644 --- a/backend/packages/harness/deerflow/models/patched_stepfun.py +++ b/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, ) diff --git a/backend/packages/harness/deerflow/models/reasoning_capture.py b/backend/packages/harness/deerflow/models/reasoning_capture.py new file mode 100644 index 0000000000..74f0cf416c --- /dev/null +++ b/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 diff --git a/backend/tests/test_reasoning_capture_single_owner.py b/backend/tests/test_reasoning_capture_single_owner.py new file mode 100644 index 0000000000..4936043719 --- /dev/null +++ b/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"}