mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-20 11:36:17 +00:00
e89b128157
* fix(models): preserve DeepSeek thinking tool history * fix(runtime): harden model response recovery * fix(runtime): tighten model response recovery * fix(runtime): protect run-scoped retry state * fix(runtime): complete model recovery review fixes * fix(runtime): preserve empty-response diagnostics * fix(runtime): strip native tool calls on length caps * docs(middleware): fit recovery guidance within inherited size limit --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""Helpers for replaying provider-specific assistant message fields.
|
|
|
|
Several provider adapters need to preserve fields that LangChain stores on the
|
|
original ``AIMessage`` but drops when serializing request payloads. This module
|
|
keeps the assistant-message matching logic shared while letting each provider
|
|
decide which fields to restore.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Callable, Sequence
|
|
from typing import Any
|
|
|
|
from langchain_core.messages import AIMessage, BaseMessage
|
|
|
|
AssistantPayloadRestorer = Callable[[dict[str, Any], AIMessage], None]
|
|
|
|
|
|
def restore_assistant_payloads(
|
|
payload_messages: Sequence[dict[str, Any]],
|
|
original_messages: Sequence[BaseMessage],
|
|
restore: AssistantPayloadRestorer,
|
|
) -> None:
|
|
"""Restore provider-specific fields onto serialized assistant payloads."""
|
|
ai_messages = [m for m in original_messages if isinstance(m, AIMessage)]
|
|
assistant_payloads = [m for m in payload_messages if m.get("role") == "assistant"]
|
|
used_ai_indexes: set[int] = set()
|
|
|
|
for ordinal, payload_msg in enumerate(assistant_payloads):
|
|
ai_msg = _match_ai_message(payload_msg, ai_messages, used_ai_indexes, ordinal)
|
|
if ai_msg is not None:
|
|
restore(payload_msg, ai_msg)
|
|
|
|
|
|
def restore_additional_kwargs_field(payload_msg: dict[str, Any], orig_msg: AIMessage, field_name: str) -> None:
|
|
"""Copy a provider-specific ``additional_kwargs`` field onto a payload message."""
|
|
value = orig_msg.additional_kwargs.get(field_name)
|
|
if value is not None:
|
|
payload_msg[field_name] = value
|
|
|
|
|
|
def restore_reasoning_content(payload_msg: dict[str, Any], orig_msg: AIMessage) -> None:
|
|
"""Copy provider reasoning content onto a serialized assistant payload."""
|
|
restore_additional_kwargs_field(payload_msg, orig_msg, "reasoning_content")
|
|
|
|
|
|
def _match_ai_message(
|
|
payload_msg: dict[str, Any],
|
|
ai_messages: Sequence[AIMessage],
|
|
used_ai_indexes: set[int],
|
|
fallback_ordinal: int,
|
|
) -> AIMessage | None:
|
|
payload_key = _assistant_signature(payload_msg)
|
|
if payload_key is not None:
|
|
matches = [index for index, ai_msg in enumerate(ai_messages) if index not in used_ai_indexes and _ai_signature(ai_msg) == payload_key]
|
|
if len(matches) == 1:
|
|
used_ai_indexes.add(matches[0])
|
|
return ai_messages[matches[0]]
|
|
|
|
fallback_index = _next_unused_index_at_or_after(len(ai_messages), used_ai_indexes, fallback_ordinal)
|
|
if fallback_index is not None:
|
|
used_ai_indexes.add(fallback_index)
|
|
return ai_messages[fallback_index]
|
|
|
|
return None
|
|
|
|
|
|
def _next_unused_index_at_or_after(count: int, used_ai_indexes: set[int], start: int) -> int | None:
|
|
"""Return the next unused AI index at or after ``start``.
|
|
|
|
Scanning forward from the payload's ordinal preserves the positional bias of
|
|
the previous behaviour while still recovering when serialization drops or
|
|
reorders messages so the exact ordinal index is already taken. It does not
|
|
wrap to earlier indexes because those messages may be represented by payload
|
|
entries that were already dropped.
|
|
"""
|
|
if count == 0 or start >= count:
|
|
return None
|
|
for index in range(start, count):
|
|
if index not in used_ai_indexes:
|
|
return index
|
|
return None
|
|
|
|
|
|
def _assistant_signature(payload_msg: dict[str, Any]) -> tuple[str, str] | None:
|
|
return _signature(
|
|
payload_msg.get("content"),
|
|
_tool_call_ids(payload_msg.get("tool_calls") or []),
|
|
)
|
|
|
|
|
|
def _ai_signature(message: AIMessage) -> tuple[str, str] | None:
|
|
tool_calls = message.tool_calls or message.additional_kwargs.get("tool_calls") or []
|
|
return _signature(message.content, _tool_call_ids(tool_calls))
|
|
|
|
|
|
def _signature(content: Any, tool_call_ids: tuple[str, ...]) -> tuple[str, str] | None:
|
|
if content in (None, "") and not tool_call_ids:
|
|
return None
|
|
return (_stable_repr(content), "|".join(tool_call_ids))
|
|
|
|
|
|
def _stable_repr(value: Any) -> str:
|
|
try:
|
|
return json.dumps(value, sort_keys=True, ensure_ascii=False)
|
|
except TypeError:
|
|
return repr(value)
|
|
|
|
|
|
def _tool_call_ids(tool_calls: Sequence[Any]) -> tuple[str, ...]:
|
|
ids: list[str] = []
|
|
for tool_call in tool_calls:
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("id")
|
|
if isinstance(call_id, str) and call_id:
|
|
ids.append(call_id)
|
|
return tuple(ids)
|