mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-28 17:06:05 +00:00
e01173d8b2
* feat(checkpoint): production-shaped full/delta benchmark with configurable snapshot frequency - Group benchmark scripts into per-family folders (checkpoint/, sandbox/) - Extract shared benchmark infrastructure into checkpoint_bench_common.py - Add checkpoint_delta_snapshot_frequency config (default 1000, process-frozen); freeze it in make_lead_agent and DeerFlowClient; key the state-schema adaptation cache by resolved frequency - New bench_production.py: per-case child processes run N ainvoke turns through the real lead-agent graph (scripted deterministic model, real AsyncSqliteSaver), then measure GET /state + POST /history through the real Gateway route stack in one event loop (httpx ASGITransport), cold/warm accessor-cache split, cross-mode digest gates - New summarize_production.py: delta/full ratios plus decision metrics (snapshot_write_spike, cache_effect_ms, checkpoint_write_share, auto-discovered history per-limit ratios) * fix(checkpoint): address production benchmark review
416 lines
15 KiB
Python
416 lines
15 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 import thread_state
|
|
from deerflow.agents.thread_state import (
|
|
DeltaThreadState,
|
|
ThreadState,
|
|
adapt_state_schema_for_mode,
|
|
get_thread_state_schema,
|
|
merge_message_writes,
|
|
normalize_middleware_state_schemas,
|
|
)
|
|
from deerflow.runtime.checkpoint_mode import CheckpointModeReconfigurationError
|
|
|
|
|
|
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_delta_adaptation_cache_varies_by_resolved_snapshot_frequency(monkeypatch) -> None:
|
|
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
|
|
thread_state._adapt_state_schema_for_mode.cache_clear()
|
|
default_adapted = adapt_state_schema_for_mode(AgentState, "delta")
|
|
|
|
thread_state.freeze_delta_snapshot_frequency(7)
|
|
configured_adapted = adapt_state_schema_for_mode(AgentState, "delta")
|
|
|
|
assert configured_adapted is not default_adapted
|
|
hint = get_type_hints(configured_adapted, include_extras=True)["messages"]
|
|
channel = next(item for item in hint.__metadata__ if isinstance(item, DeltaChannel))
|
|
assert channel.snapshot_frequency == 7
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_delta_snapshot_frequency_defaults_to_1000(monkeypatch) -> None:
|
|
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
|
|
assert thread_state.resolved_delta_snapshot_frequency() == 1000
|
|
|
|
|
|
def test_freeze_delta_snapshot_frequency_is_process_frozen(monkeypatch) -> None:
|
|
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
|
|
assert thread_state.freeze_delta_snapshot_frequency(50) == 50
|
|
assert thread_state.freeze_delta_snapshot_frequency(50) == 50
|
|
with pytest.raises(CheckpointModeReconfigurationError):
|
|
thread_state.freeze_delta_snapshot_frequency(10)
|
|
|
|
|
|
def test_get_thread_state_schema_honors_frozen_frequency(monkeypatch) -> None:
|
|
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
|
|
thread_state.freeze_delta_snapshot_frequency(7)
|
|
hint = get_type_hints(thread_state.get_thread_state_schema("delta"), include_extras=True)["messages"]
|
|
channel = next(item for item in hint.__metadata__ if isinstance(item, DeltaChannel))
|
|
assert channel.snapshot_frequency == 7
|
|
|
|
|
|
def test_database_config_default_snapshot_frequency() -> None:
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
|
|
assert DatabaseConfig().checkpoint_delta_snapshot_frequency == 1000
|
|
assert DatabaseConfig(checkpoint_delta_snapshot_frequency=25).checkpoint_delta_snapshot_frequency == 25
|