mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-26 18:06:00 +00:00
fix(todo): reuse thread state schema (#3206)
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
@@ -17,6 +17,7 @@ from deerflow.agents.middlewares.todo_middleware import (
|
||||
_reminder_in_messages,
|
||||
_todos_in_messages,
|
||||
)
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
|
||||
|
||||
def _ai_with_write_todos():
|
||||
@@ -510,6 +511,42 @@ class TestWrapModelCall:
|
||||
|
||||
|
||||
class TestTodoMiddlewareAgentGraphIntegration:
|
||||
def test_reuses_thread_state_todos_schema_in_real_agent_graph(self):
|
||||
mw = TodoMiddleware()
|
||||
model = _CapturingFakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "write_todos",
|
||||
"id": "todos-1",
|
||||
"args": {
|
||||
"todos": [
|
||||
{"content": "Step 1", "status": "pending"},
|
||||
]
|
||||
},
|
||||
}
|
||||
],
|
||||
),
|
||||
AIMessage(content="final"),
|
||||
],
|
||||
)
|
||||
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[],
|
||||
middleware=[mw],
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [("user", "create a todo")]},
|
||||
context={"thread_id": "schema-thread", "run_id": "schema-run"},
|
||||
)
|
||||
|
||||
assert result["todos"] == [{"content": "Step 1", "status": "pending"}]
|
||||
|
||||
def test_completion_reminder_is_transient_in_real_agent_graph(self):
|
||||
mw = TodoMiddleware()
|
||||
model = _CapturingFakeMessagesListChatModel(
|
||||
|
||||
Reference in New Issue
Block a user