mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-26 16:07:53 +00:00
8c19a2eb36
* perf(checkpoint): linearize message write merging * test(checkpoint): address message reducer review
372 lines
13 KiB
Python
372 lines
13 KiB
Python
import copy
|
|
from typing import get_type_hints
|
|
|
|
import pytest
|
|
from hypothesis import given
|
|
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, ToolMessageChunk
|
|
from langgraph.channels import DeltaChannel
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.graph import StateGraph
|
|
from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
|
|
|
|
from deerflow.agents.thread_state import (
|
|
DeltaThreadState,
|
|
ThreadState,
|
|
adapt_state_schema_for_mode,
|
|
get_thread_state_schema,
|
|
merge_message_writes,
|
|
normalize_middleware_state_schemas,
|
|
)
|
|
|
|
|
|
def _fold(state: list, writes: list) -> list:
|
|
result = list(state)
|
|
for write in writes:
|
|
result = list(add_messages(result, write))
|
|
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",
|
|
[
|
|
[[HumanMessage(id="h1", content="one")], [AIMessage(id="a1", content="two")]],
|
|
[[AIMessage(id="same", content="old")], [AIMessage(id="same", content="new")]],
|
|
[[HumanMessage(id="h1", content="one")], [RemoveMessage(id="h1")]],
|
|
[
|
|
[HumanMessage(id="h1", content="one"), AIMessage(id="a1", content="two")],
|
|
[RemoveMessage(id=REMOVE_ALL_MESSAGES), HumanMessage(id="h2", content="kept")],
|
|
],
|
|
],
|
|
)
|
|
def test_merge_message_writes_matches_sequential_add_messages(writes: list) -> None:
|
|
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")]
|
|
writes = [
|
|
[AIMessage(id="a1", content="first")],
|
|
[AIMessage(id="a1", content="replacement")],
|
|
[HumanMessage(id="h2", content="last")],
|
|
]
|
|
xs = writes[:split]
|
|
ys = writes[split:]
|
|
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")]]
|
|
|
|
with pytest.raises(ValueError) as expected:
|
|
_fold([], writes)
|
|
with pytest.raises(type(expected.value)) as actual:
|
|
merge_message_writes([], writes)
|
|
|
|
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",
|
|
[
|
|
[{"role": "user", "content": "from a dict", "id": "dict-1"}],
|
|
AIMessageChunk(id="chunk-1", content="from a chunk"),
|
|
],
|
|
)
|
|
def test_merge_message_writes_matches_message_coercion(write: object) -> None:
|
|
assert merge_message_writes([], [write]) == _fold([], [write])
|
|
|
|
|
|
def test_raw_tuple_coercion_matches_add_messages_reducer_parity(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: "tuple-id")
|
|
writes = [[("human", "from a tuple")]]
|
|
|
|
assert merge_message_writes([], writes) == _fold([], writes)
|
|
|
|
|
|
def test_mode_selects_expected_state_schema() -> None:
|
|
assert get_thread_state_schema("full") is ThreadState
|
|
assert get_thread_state_schema("delta") is DeltaThreadState
|
|
message_hint = get_type_hints(DeltaThreadState, include_extras=True)["messages"]
|
|
assert any(isinstance(item, DeltaChannel) for item in message_hint.__metadata__)
|
|
|
|
|
|
def test_delta_adaptation_replaces_agent_state_message_reducer() -> None:
|
|
adapted = adapt_state_schema_for_mode(AgentState, "delta")
|
|
hint = get_type_hints(adapted, include_extras=True)["messages"]
|
|
assert any(isinstance(item, DeltaChannel) for item in hint.__metadata__)
|
|
|
|
|
|
def test_agents_package_exports_delta_thread_state() -> None:
|
|
from deerflow.agents import DeltaThreadState as ExportedDeltaThreadState
|
|
|
|
assert ExportedDeltaThreadState is DeltaThreadState
|
|
|
|
|
|
class _FirstState(AgentState):
|
|
first: str
|
|
|
|
|
|
class _SecondState(AgentState):
|
|
second: int
|
|
|
|
|
|
class _FirstMiddleware(AgentMiddleware):
|
|
state_schema = _FirstState
|
|
|
|
|
|
class _SecondMiddleware(AgentMiddleware):
|
|
state_schema = _SecondState
|
|
|
|
|
|
class _FakeModel(FakeMessagesListChatModel):
|
|
def bind_tools(self, tools, **kwargs): # type: ignore[override]
|
|
return self
|
|
|
|
|
|
def _compile_with_middleware(middleware: list[AgentMiddleware], mode: str):
|
|
return create_agent(
|
|
model=_FakeModel(responses=[AIMessage(id="response", content="done")]),
|
|
tools=None,
|
|
middleware=normalize_middleware_state_schemas(middleware, mode),
|
|
state_schema=get_thread_state_schema(mode),
|
|
)
|
|
|
|
|
|
def test_delta_normalization_compiles_stable_channel_without_mutating_middleware() -> None:
|
|
first = _FirstMiddleware()
|
|
second = _SecondMiddleware()
|
|
middleware = [first, second]
|
|
|
|
for _ in range(10):
|
|
graph = _compile_with_middleware(middleware, "delta")
|
|
assert isinstance(graph.channels["messages"], DeltaChannel)
|
|
|
|
assert first.state_schema is _FirstState
|
|
assert second.state_schema is _SecondState
|
|
|
|
full_graph = _compile_with_middleware(middleware, "full")
|
|
assert type(full_graph.channels["messages"]).__name__ == "BinaryOperatorAggregate"
|
|
assert first.state_schema is _FirstState
|
|
assert second.state_schema is _SecondState
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"write",
|
|
[
|
|
HumanMessage(content="root BaseMessage"),
|
|
{"role": "user", "content": "root message dict"},
|
|
[HumanMessage(content="BaseMessage in list")],
|
|
[{"role": "user", "content": "message dict in list"}],
|
|
],
|
|
ids=["base-message", "dict", "base-message-list", "dict-list"],
|
|
)
|
|
def test_production_message_forms_keep_assigned_ids_across_delta_replay(write: object) -> None:
|
|
builder = StateGraph(DeltaThreadState)
|
|
|
|
def write_messages(_state):
|
|
return {"messages": write}
|
|
|
|
builder.add_node("writer", write_messages)
|
|
builder.set_entry_point("writer")
|
|
builder.set_finish_point("writer")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
config = {"configurable": {"thread_id": "stable-message-replay"}}
|
|
|
|
graph.invoke({}, config)
|
|
first = graph.get_state(config).values["messages"]
|
|
second = graph.get_state(config).values["messages"]
|
|
|
|
assert first[0].id is not None
|
|
assert second[0].id == first[0].id
|