141 lines
5.1 KiB
Python
141 lines
5.1 KiB
Python
"""Normalize OpenAI-compatible thinking fields across provider versions.
|
|
|
|
Older vLLM / DeepSeek put thinking in ``reasoning_content``. Newer vLLM
|
|
reasoning parsers put it in ``reasoning``. DeerFlow's UI and downstream code
|
|
canonicalize on ``AIMessage.additional_kwargs.reasoning_content``.
|
|
|
|
Only fields that actually appeared on a response are echoed back on later
|
|
turns, so models that still use ``reasoning_content`` (or send neither) are
|
|
unchanged.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from langchain_core.messages import AIMessage, AIMessageChunk
|
|
|
|
|
|
def reasoning_to_text(reasoning: Any) -> str:
|
|
"""Best-effort extraction of readable reasoning text from provider payloads."""
|
|
if isinstance(reasoning, str):
|
|
return reasoning
|
|
|
|
if isinstance(reasoning, list):
|
|
parts = [reasoning_to_text(item) for item in reasoning]
|
|
return "".join(part for part in parts if part)
|
|
|
|
if isinstance(reasoning, dict):
|
|
for key in ("text", "content", "reasoning"):
|
|
value = reasoning.get(key)
|
|
if isinstance(value, str):
|
|
return value
|
|
if value is not None:
|
|
text = reasoning_to_text(value)
|
|
if text:
|
|
return text
|
|
try:
|
|
return json.dumps(reasoning, ensure_ascii=False)
|
|
except TypeError:
|
|
return str(reasoning)
|
|
|
|
if reasoning is None:
|
|
return ""
|
|
try:
|
|
return json.dumps(reasoning, ensure_ascii=False)
|
|
except TypeError:
|
|
return str(reasoning)
|
|
|
|
|
|
def extract_reasoning_text(payload: dict | None) -> str | None:
|
|
"""Read thinking text from ``reasoning_content`` or ``reasoning``.
|
|
|
|
Prefer ``reasoning_content`` when both are present so older DeepSeek /
|
|
vLLM payloads keep their original string. Newer vLLM that only sends
|
|
``reasoning`` (string or structured) is still captured.
|
|
"""
|
|
if not isinstance(payload, dict):
|
|
return None
|
|
for key in ("reasoning_content", "reasoning"):
|
|
value = payload.get(key)
|
|
if value is None:
|
|
continue
|
|
text = reasoning_to_text(value)
|
|
if text:
|
|
return text
|
|
return None
|
|
|
|
|
|
def _merge_streamed_value(existing: Any, incoming: Any) -> Any:
|
|
if isinstance(existing, str) and isinstance(incoming, str):
|
|
if incoming and incoming not in existing:
|
|
return f"{existing}{incoming}"
|
|
return existing or incoming
|
|
return incoming if incoming is not None else existing
|
|
|
|
|
|
def apply_reasoning_from_payload(
|
|
message: AIMessage | AIMessageChunk,
|
|
payload: dict | None,
|
|
*,
|
|
append: bool = True,
|
|
) -> AIMessage | AIMessageChunk:
|
|
"""Copy provider thinking fields onto ``additional_kwargs`` for UI + echo-back."""
|
|
if not isinstance(payload, dict):
|
|
return message
|
|
|
|
additional_kwargs = dict(message.additional_kwargs)
|
|
changed = False
|
|
|
|
if "reasoning" in payload and payload.get("reasoning") is not None:
|
|
incoming = payload.get("reasoning")
|
|
existing = additional_kwargs.get("reasoning")
|
|
merged = _merge_streamed_value(existing, incoming) if append else incoming
|
|
if merged is not None and merged != existing:
|
|
additional_kwargs["reasoning"] = merged
|
|
changed = True
|
|
|
|
text = extract_reasoning_text(payload)
|
|
if text:
|
|
existing = additional_kwargs.get("reasoning_content")
|
|
if append and isinstance(existing, str) and existing:
|
|
if text not in existing:
|
|
additional_kwargs["reasoning_content"] = f"{existing}{text}"
|
|
changed = True
|
|
elif text != existing:
|
|
additional_kwargs["reasoning_content"] = text
|
|
changed = True
|
|
|
|
if not changed:
|
|
return message
|
|
return message.model_copy(update={"additional_kwargs": additional_kwargs})
|
|
|
|
|
|
def fill_ui_reasoning_content(message: AIMessage | AIMessageChunk) -> AIMessage | AIMessageChunk:
|
|
"""If LangChain stored ``reasoning`` but not the UI field, copy the text across."""
|
|
kwargs = message.additional_kwargs
|
|
reasoning = kwargs.get("reasoning")
|
|
if reasoning is None or kwargs.get("reasoning_content"):
|
|
return message
|
|
return apply_reasoning_from_payload(message, {"reasoning": reasoning}, append=False)
|
|
|
|
|
|
def restore_reasoning_fields(payload_msg: dict, orig_msg: AIMessage) -> None:
|
|
"""Re-inject the thinking field this assistant turn originally carried.
|
|
|
|
Newer vLLM stores thinking on ``reasoning`` and we also copy a UI-facing
|
|
``reasoning_content`` string from it. Echo the wire field that actually
|
|
arrived: ``reasoning`` wins when present so we do not send a synthesized
|
|
``reasoning_content`` to a parser that no longer accepts it. Older
|
|
DeepSeek / vLLM turns only have ``reasoning_content`` and stay on that
|
|
field.
|
|
"""
|
|
reasoning = orig_msg.additional_kwargs.get("reasoning")
|
|
if reasoning is not None:
|
|
payload_msg["reasoning"] = reasoning
|
|
return
|
|
reasoning_content = orig_msg.additional_kwargs.get("reasoning_content")
|
|
if reasoning_content is not None:
|
|
payload_msg["reasoning_content"] = reasoning_content
|