137 lines
6.2 KiB
Python
137 lines
6.2 KiB
Python
"""Step-budget recovery must preserve history and allow a real next turn."""
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.graph import END, START, MessagesState, StateGraph
|
|
from pydantic import ValidationError
|
|
|
|
from deerflow.config.system_settings import ConversationContextSettings
|
|
from deerflow.runtime.runs.manager import RunManager
|
|
from deerflow.runtime.runs.schemas import RunStatus
|
|
from deerflow.runtime.runs.worker import RunContext, run_agent
|
|
|
|
|
|
@pytest.mark.parametrize("value", [19, 10001, 25.5, True, "100"])
|
|
def test_invalid_budget(value):
|
|
from app.gateway.routers.system_settings import ConversationContextUpdateRequest
|
|
|
|
with pytest.raises(ValidationError):
|
|
ConversationContextSettings(max_run_steps=value)
|
|
with pytest.raises(ValidationError):
|
|
ConversationContextUpdateRequest(max_run_steps=value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("budget", [3, 4])
|
|
async def test_limit_is_saved_without_error_and_continue_works(budget):
|
|
calls = []
|
|
|
|
def model(state):
|
|
messages = state["messages"]
|
|
if any(isinstance(m, HumanMessage) and m.content == "继续" for m in messages):
|
|
# A provider requires every earlier call to have a result.
|
|
tool_ids = {m.tool_call_id for m in messages if isinstance(m, ToolMessage)}
|
|
assert all(c["id"] in tool_ids for m in messages if isinstance(m, AIMessage) for c in m.tool_calls)
|
|
return {"messages": [AIMessage(content="已继续完成")]}
|
|
return {
|
|
"messages": [
|
|
AIMessage(
|
|
content="正在处理",
|
|
tool_calls=[
|
|
{
|
|
"id": f"call-{len(messages)}",
|
|
"name": "search",
|
|
"args": {},
|
|
}
|
|
],
|
|
)
|
|
]
|
|
}
|
|
|
|
def tool(state):
|
|
call = state["messages"][-1].tool_calls[0]
|
|
calls.append(call["id"])
|
|
return {"messages": [ToolMessage(content="已有检索结果", tool_call_id=call["id"])]}
|
|
|
|
graph = StateGraph(MessagesState)
|
|
graph.add_node("model", model)
|
|
graph.add_node("tools", tool)
|
|
graph.add_edge(START, "model")
|
|
graph.add_conditional_edges("model", lambda s: "tools" if s["messages"][-1].tool_calls else END)
|
|
graph.add_edge("tools", "model")
|
|
saver = InMemorySaver()
|
|
agent = graph.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "step-test"}, "recursion_limit": budget}
|
|
manager = RunManager()
|
|
record = await manager.create("step-test", "lead_agent")
|
|
bridge = AsyncMock()
|
|
await run_agent(bridge, manager, record, ctx=RunContext(checkpointer=saver), agent_factory=lambda config: agent, graph_input={"messages": [HumanMessage(content="研究")]}, config=config, stream_modes=["messages", "values"])
|
|
assert record.status == RunStatus.interrupted
|
|
assert record.error is None
|
|
assert not any(c.args[1] == "error" for c in bridge.publish.call_args_list)
|
|
events = {c.args[1]: c.args[2] for c in bridge.publish.call_args_list}
|
|
assert "步骤上限" in events["messages"][0]["content"]
|
|
assert "步骤上限" in events["values"]["messages"][-1]["content"]
|
|
bridge.publish_end.assert_awaited_once()
|
|
state = await agent.aget_state(config)
|
|
assert not state.next
|
|
assert "已有检索结果" in [m.content for m in state.values["messages"]]
|
|
assert "步骤上限" in state.values["messages"][-1].content
|
|
count = len(calls)
|
|
result = await agent.ainvoke({"messages": [HumanMessage(content="继续")]}, config)
|
|
assert result["messages"][-1].content == "已继续完成"
|
|
assert len(calls) == count
|
|
|
|
|
|
def test_server_budget_overrides_client_and_preserves_other_graphs(monkeypatch):
|
|
from app.gateway.services import build_run_config
|
|
|
|
monkeypatch.setattr("deerflow.config.system_settings.get_conversation_context_settings", lambda: ConversationContextSettings(max_run_steps=800))
|
|
for assistant in (None, "lead_agent", "custom-agent"):
|
|
assert build_run_config("t", {"recursion_limit": 99999}, None, assistant_id=assistant)["recursion_limit"] == 800
|
|
assert build_run_config("t", {"recursion_limit": 250}, None, assistant_id="canvas_agent")["recursion_limit"] == 250
|
|
|
|
|
|
@pytest.mark.parametrize("client_limit", [None, 1, 300])
|
|
@pytest.mark.parametrize("server_limit", [20, 800])
|
|
def test_open_chat_uses_server_budget(monkeypatch, client_limit, server_limit):
|
|
from app.gateway.routers.open_chat import OpenChatRequest, _build_run_context
|
|
|
|
monkeypatch.setattr("deerflow.config.system_settings.get_conversation_context_settings", lambda: ConversationContextSettings(max_run_steps=server_limit))
|
|
body = OpenChatRequest(message="hello", **({"recursion_limit": client_limit} if client_limit is not None else {}))
|
|
config, _, _ = _build_run_context(body)
|
|
assert config["recursion_limit"] == server_limit
|
|
assert config["configurable"]["ephemeral"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setting_requires_admin_and_roundtrips(monkeypatch):
|
|
from fastapi import HTTPException
|
|
|
|
from app.gateway.routers import system_settings as routes
|
|
from deerflow.config.system_settings import SystemSettings
|
|
|
|
current = SystemSettings()
|
|
monkeypatch.setattr(routes, "load_system_settings", lambda: current)
|
|
|
|
def save(value):
|
|
nonlocal current
|
|
current = value
|
|
return value
|
|
|
|
monkeypatch.setattr(routes, "save_system_settings", save)
|
|
guard = AsyncMock(side_effect=HTTPException(status_code=403))
|
|
monkeypatch.setattr(routes, "_require_admin", guard)
|
|
body = routes.ConversationContextUpdateRequest(max_run_steps=600)
|
|
with pytest.raises(HTTPException):
|
|
await routes.update_conversation_context_setting(body, None)
|
|
assert current.conversation_context.max_run_steps == 1000
|
|
guard.side_effect = None
|
|
result = await routes.update_conversation_context_setting(body, None)
|
|
assert result.max_run_steps == 600
|
|
assert (await routes.get_conversation_context_setting()).max_run_steps == 600
|
|
assert current.conversation_context.recent_turn_window_enabled is False
|