mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-25 23:48:00 +00:00
perf(checkpoint): linearize message write merging (#4421)
* perf(checkpoint): linearize message write merging * test(checkpoint): address message reducer review
This commit is contained in:
Vanzeren
GitHub
parent
3b77a7401b
commit
8c19a2eb36
@ -291,6 +291,15 @@ tool graph or subagent executor during state/schema imports.
|
||||
**ThreadState** (`packages/harness/deerflow/agents/thread_state.py`):
|
||||
- Extends `AgentState` with: `sandbox`, `thread_data`, `title`, `artifacts`, `todos`, `uploaded_files`, `viewed_images`, `goal`, `promoted`, `delegations`, `skill_context`, `summary_text`
|
||||
- Uses custom reducers: `merge_artifacts` (deduplicate), `merge_viewed_images` (merge/clear), `merge_goal` (preserve the active goal across ordinary state updates unless the goal writer replaces it), `merge_promoted` (catalog-hash-scoped deferred tool promotions), `merge_delegations` (append task delegation entries, same id latest wins, terminal status never downgraded, capped to the most recent entries), and `merge_skill_context` (dedupe active-skill references by path, keep the most recently read entries; entries store a name/path/description reference, not the SKILL.md body). `summary_text` is a LastValue channel updated by summarization and projected into model requests as durable context data instead of being stored as a `messages` item.
|
||||
- Delta-mode `merge_message_writes` normalizes the current message state once,
|
||||
then folds normalized writes in order with message-ID position indexes and
|
||||
deferred tombstone compaction. It preserves public `add_messages` behavior,
|
||||
including duplicate IDs, replacement position, removal errors,
|
||||
`REMOVE_ALL_MESSAGES`, null-write errors, and missing-ID allocation order,
|
||||
without rescanning the accumulated state for every write. Keep this
|
||||
full-parity contract covered by differential tests: LangGraph's private
|
||||
`_messages_delta_reducer` is also linear, but intentionally omits some of
|
||||
those public `add_messages` semantics and cannot be substituted directly.
|
||||
|
||||
**Runtime Configuration** (via `config.configurable`):
|
||||
- `thinking_enabled` - Enable model's extended thinking
|
||||
|
||||
@ -1,12 +1,19 @@
|
||||
import copy
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from functools import cache
|
||||
from typing import Annotated, Any, NotRequired, TypedDict, get_type_hints
|
||||
from typing import Annotated, Any, NotRequired, TypedDict, cast, get_type_hints
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain_core.messages import AnyMessage
|
||||
from langchain_core.messages import (
|
||||
AnyMessage,
|
||||
BaseMessageChunk,
|
||||
RemoveMessage,
|
||||
convert_to_messages,
|
||||
message_chunk_to_message,
|
||||
)
|
||||
from langgraph.channels import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
||||
|
||||
import deerflow.checkpoint_patches as _checkpoint_patches # noqa: F401 - import-time saver fixes
|
||||
from deerflow.agents.goal_state import GoalState
|
||||
@ -258,11 +265,91 @@ class ThreadState(AgentState):
|
||||
summary_text: NotRequired[str | None]
|
||||
|
||||
|
||||
def _normalize_messages(value: Any) -> list[AnyMessage]:
|
||||
values = value if isinstance(value, list) else [value]
|
||||
messages = [message_chunk_to_message(cast(BaseMessageChunk, message)) for message in convert_to_messages(values)]
|
||||
for message in messages:
|
||||
if message.id is None:
|
||||
message.id = str(uuid.uuid4())
|
||||
return messages
|
||||
|
||||
|
||||
def _index_messages(
|
||||
messages: list[AnyMessage | None],
|
||||
) -> tuple[dict[str, int], dict[str, list[int]]]:
|
||||
latest_position: dict[str, int] = {}
|
||||
positions_by_id: dict[str, list[int]] = {}
|
||||
for position, message in enumerate(messages):
|
||||
if message is None:
|
||||
continue
|
||||
message_id = cast(str, message.id)
|
||||
latest_position[message_id] = position
|
||||
positions_by_id.setdefault(message_id, []).append(position)
|
||||
return latest_position, positions_by_id
|
||||
|
||||
|
||||
def _raise_null_write(has_messages: bool) -> None:
|
||||
# ``add_messages(left, None)`` reports only ``left`` when the accumulated
|
||||
# message list is non-empty; with an empty list, it reports only ``right``.
|
||||
received = "left" if has_messages else "right"
|
||||
raise ValueError(f"Must specify non-null arguments for both 'left' and 'right'. Only received: '{received}'.")
|
||||
|
||||
|
||||
def merge_message_writes(state: list[AnyMessage], writes: Sequence[Any]) -> list[AnyMessage]:
|
||||
result = list(state)
|
||||
"""Fold DeltaChannel writes with ``add_messages`` semantics in linear time.
|
||||
|
||||
LangGraph's private ``_messages_delta_reducer`` is also linear, but does
|
||||
not preserve the public reducer's full coercion, ID, removal, and
|
||||
``REMOVE_ALL_MESSAGES`` behavior.
|
||||
"""
|
||||
if not writes:
|
||||
return list(state)
|
||||
if writes[0] is None:
|
||||
_raise_null_write(bool(state))
|
||||
|
||||
messages: list[AnyMessage | None] = _normalize_messages(state)
|
||||
latest_position, positions_by_id = _index_messages(messages)
|
||||
|
||||
for write in writes:
|
||||
result = list(add_messages(result, write))
|
||||
return result
|
||||
if write is None:
|
||||
_raise_null_write(bool(latest_position))
|
||||
normalized_write = _normalize_messages(write)
|
||||
remove_all_idx = None
|
||||
for position, message in enumerate(normalized_write):
|
||||
if isinstance(message, RemoveMessage) and message.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = position
|
||||
|
||||
if remove_all_idx is not None:
|
||||
messages = list(normalized_write[remove_all_idx + 1 :])
|
||||
latest_position, positions_by_id = _index_messages(messages)
|
||||
continue
|
||||
|
||||
ids_to_remove: set[str] = set()
|
||||
for message in normalized_write:
|
||||
message_id = cast(str, message.id)
|
||||
existing_position = latest_position.get(message_id)
|
||||
if existing_position is not None:
|
||||
if isinstance(message, RemoveMessage):
|
||||
ids_to_remove.add(message_id)
|
||||
else:
|
||||
ids_to_remove.discard(message_id)
|
||||
messages[existing_position] = message
|
||||
continue
|
||||
|
||||
if isinstance(message, RemoveMessage):
|
||||
raise ValueError(f"Attempting to delete a message with an ID that doesn't exist ('{message_id}')")
|
||||
|
||||
position = len(messages)
|
||||
messages.append(message)
|
||||
latest_position[message_id] = position
|
||||
positions_by_id[message_id] = [position]
|
||||
|
||||
for message_id in ids_to_remove:
|
||||
for position in positions_by_id.pop(message_id):
|
||||
messages[position] = None
|
||||
del latest_position[message_id]
|
||||
|
||||
return [message for message in messages if message is not None]
|
||||
|
||||
|
||||
DELTA_MESSAGES_FIELD = Annotated[
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
import copy
|
||||
from typing import get_type_hints
|
||||
|
||||
import pytest
|
||||
@ -6,7 +7,7 @@ from hypothesis import strategies as st
|
||||
from langchain.agents import AgentState, create_agent
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage, ToolMessageChunk
|
||||
from langgraph.channels import DeltaChannel
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.graph import StateGraph
|
||||
@ -29,6 +30,57 @@ def _fold(state: list, writes: list) -> list:
|
||||
return result
|
||||
|
||||
|
||||
def _outcome(call):
|
||||
try:
|
||||
return ("result", call())
|
||||
except Exception as exc:
|
||||
return ("error", type(exc), str(exc))
|
||||
|
||||
|
||||
@st.composite
|
||||
def _message_merge_cases(draw):
|
||||
message_ids = ["a", "b", "c", "missing"]
|
||||
state_ids = draw(st.lists(st.sampled_from(message_ids[:-1]), max_size=6))
|
||||
state = [
|
||||
{
|
||||
"role": draw(st.sampled_from(["user", "assistant"])),
|
||||
"content": f"state-{index}",
|
||||
"id": message_id,
|
||||
}
|
||||
for index, message_id in enumerate(state_ids)
|
||||
]
|
||||
|
||||
operation = st.one_of(
|
||||
st.tuples(
|
||||
st.just("message"),
|
||||
st.sampled_from(message_ids),
|
||||
st.sampled_from(["user", "assistant", "ai_chunk", "tool_chunk"]),
|
||||
st.text(max_size=12),
|
||||
),
|
||||
st.tuples(
|
||||
st.just("remove"),
|
||||
st.sampled_from([*message_ids, REMOVE_ALL_MESSAGES]),
|
||||
st.none(),
|
||||
st.none(),
|
||||
),
|
||||
)
|
||||
raw_writes = draw(st.lists(st.lists(operation, max_size=6), max_size=8))
|
||||
writes = []
|
||||
for raw_write in raw_writes:
|
||||
write = []
|
||||
for kind, message_id, role, content in raw_write:
|
||||
if kind == "remove":
|
||||
write.append(RemoveMessage(id=message_id))
|
||||
elif role == "ai_chunk":
|
||||
write.append(AIMessageChunk(id=message_id, content=content))
|
||||
elif role == "tool_chunk":
|
||||
write.append(ToolMessageChunk(id=message_id, content=content, tool_call_id=f"call-{message_id}"))
|
||||
else:
|
||||
write.append({"role": role, "content": content, "id": message_id})
|
||||
writes.append(write)
|
||||
return state, writes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"writes",
|
||||
[
|
||||
@ -45,6 +97,16 @@ def test_merge_message_writes_matches_sequential_add_messages(writes: list) -> N
|
||||
assert merge_message_writes([], writes) == _fold([], writes)
|
||||
|
||||
|
||||
@given(case=_message_merge_cases())
|
||||
def test_merge_message_writes_randomized_differential(case: tuple[list, list]) -> None:
|
||||
state, writes = case
|
||||
|
||||
expected = _outcome(lambda: _fold(copy.deepcopy(state), copy.deepcopy(writes)))
|
||||
actual = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
@given(split=st.integers(min_value=0, max_value=3))
|
||||
def test_merge_message_writes_is_batching_invariant(split: int) -> None:
|
||||
state = [HumanMessage(id="h0", content="seed")]
|
||||
@ -58,6 +120,20 @@ def test_merge_message_writes_is_batching_invariant(split: int) -> None:
|
||||
assert merge_message_writes(merge_message_writes(state, xs), ys) == merge_message_writes(state, writes)
|
||||
|
||||
|
||||
@given(case=_message_merge_cases(), data=st.data())
|
||||
def test_merge_message_writes_randomized_batching_invariance(case: tuple[list, list], data: st.DataObject) -> None:
|
||||
state, writes = case
|
||||
split = data.draw(st.integers(min_value=0, max_value=len(writes)))
|
||||
|
||||
expected = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
|
||||
|
||||
def batched():
|
||||
intermediate = merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes[:split]))
|
||||
return merge_message_writes(intermediate, copy.deepcopy(writes[split:]))
|
||||
|
||||
assert _outcome(batched) == expected
|
||||
|
||||
|
||||
def test_merge_message_writes_matches_unknown_remove_error() -> None:
|
||||
writes = [[RemoveMessage(id="missing")]]
|
||||
|
||||
@ -69,6 +145,117 @@ def test_merge_message_writes_matches_unknown_remove_error() -> None:
|
||||
assert str(actual.value) == str(expected.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("state", "writes"),
|
||||
[
|
||||
(
|
||||
[HumanMessage(id="duplicate", content="first"), HumanMessage(id="duplicate", content="second")],
|
||||
[[AIMessage(id="duplicate", content="replacement")]],
|
||||
),
|
||||
(
|
||||
[HumanMessage(id="duplicate", content="first"), HumanMessage(id="duplicate", content="second")],
|
||||
[[RemoveMessage(id="duplicate")]],
|
||||
),
|
||||
(
|
||||
[HumanMessage(id="same", content="old")],
|
||||
[[RemoveMessage(id="same"), AIMessage(id="same", content="same-write replacement")]],
|
||||
),
|
||||
(
|
||||
[HumanMessage(id="same", content="old")],
|
||||
[[RemoveMessage(id="same")], [AIMessage(id="same", content="later-write append")]],
|
||||
),
|
||||
(
|
||||
[HumanMessage(id="seed", content="old")],
|
||||
[
|
||||
[
|
||||
RemoveMessage(id="unknown-but-ignored"),
|
||||
RemoveMessage(id=REMOVE_ALL_MESSAGES),
|
||||
RemoveMessage(id="suffix-is-returned-verbatim"),
|
||||
]
|
||||
],
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"duplicate-id-replacement",
|
||||
"duplicate-id-removal",
|
||||
"same-write-remove-then-replace",
|
||||
"cross-write-remove-then-append",
|
||||
"remove-all-short-circuit",
|
||||
],
|
||||
)
|
||||
def test_merge_message_writes_preserves_add_messages_edge_semantics(state: list, writes: list) -> None:
|
||||
assert merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)) == _fold(copy.deepcopy(state), copy.deepcopy(writes))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("state", "writes"),
|
||||
[
|
||||
([], [None]),
|
||||
([HumanMessage(id="seed", content="seed")], [None]),
|
||||
([], [[HumanMessage(id="added", content="added")], None]),
|
||||
(
|
||||
[HumanMessage(id="removed", content="removed")],
|
||||
[[RemoveMessage(id="removed")], None],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_merge_message_writes_preserves_null_write_errors(state: list, writes: list) -> None:
|
||||
expected = _outcome(lambda: _fold(copy.deepcopy(state), copy.deepcopy(writes)))
|
||||
actual = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
def test_merge_message_writes_preserves_missing_id_allocation_order(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
state = [HumanMessage(content="state")]
|
||||
writes = [
|
||||
[AIMessage(content="first"), HumanMessage(content="second")],
|
||||
[AIMessage(content="third")],
|
||||
]
|
||||
|
||||
expected_ids = iter(["state-id", "first-id", "second-id", "third-id"])
|
||||
monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: next(expected_ids))
|
||||
expected = _fold(copy.deepcopy(state), copy.deepcopy(writes))
|
||||
|
||||
actual_ids = iter(["state-id", "first-id", "second-id", "third-id"])
|
||||
monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: next(actual_ids))
|
||||
actual = merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes))
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
def test_merge_message_writes_normalizes_state_and_each_write_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import deerflow.agents.thread_state as thread_state
|
||||
|
||||
state = [HumanMessage(id="state", content="state")]
|
||||
writes = [
|
||||
[AIMessage(id="first", content="first")],
|
||||
[AIMessage(id="second", content="second")],
|
||||
[AIMessage(id="third", content="third")],
|
||||
]
|
||||
original = thread_state.convert_to_messages
|
||||
normalized_inputs = []
|
||||
|
||||
def record_conversion(messages):
|
||||
normalized_inputs.append(messages)
|
||||
return original(messages)
|
||||
|
||||
monkeypatch.setattr(thread_state, "convert_to_messages", record_conversion)
|
||||
|
||||
merge_message_writes(state, writes)
|
||||
|
||||
assert normalized_inputs == [state, *writes]
|
||||
|
||||
|
||||
def test_merge_message_writes_empty_batch_does_not_assign_state_ids() -> None:
|
||||
state = [HumanMessage(content="unchanged")]
|
||||
|
||||
result = merge_message_writes(state, [])
|
||||
|
||||
assert result == state
|
||||
assert state[0].id is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"write",
|
||||
[
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user