1037 lines
38 KiB
Python
1037 lines
38 KiB
Python
"""写作模式：仅在开始写 md 时改走 deep_research_report，且不强制 tool_choice。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
|
|
from deerflow.agents.analysis_project import (
|
|
SOLUTION_REQUIREMENTS_CONFIRMATION_TOKEN,
|
|
SOLUTION_REQUIREMENTS_FORCE_TOKEN,
|
|
)
|
|
from deerflow.agents.middlewares.writing_mode_report_middleware import (
|
|
TOOL_NAME,
|
|
WritingModeReportMiddleware,
|
|
extract_report_focus,
|
|
extract_report_topic,
|
|
is_handwritten_report_call,
|
|
is_report_reread_call,
|
|
pipeline_already_succeeded,
|
|
rewrite_handwritten_report_calls,
|
|
)
|
|
from deerflow.tools.builtins.deep_research_report_tool import is_writing_mode_report_path
|
|
|
|
|
|
def test_only_report_md_is_writing_mode_follow_up() -> None:
|
|
assert is_writing_mode_report_path("/mnt/user-data/outputs/report.md") is True
|
|
assert is_writing_mode_report_path("report.md") is True
|
|
assert is_writing_mode_report_path("/mnt/user-data/outputs/notes.md") is False
|
|
assert is_writing_mode_report_path("/mnt/user-data/outputs/analysis.markdown") is False
|
|
assert is_writing_mode_report_path("") is False
|
|
assert is_writing_mode_report_path(None) is False
|
|
|
|
|
|
def test_handwritten_report_detects_outputs_markdown() -> None:
|
|
assert is_handwritten_report_call({"name": "write_file", "args": {"path": "/mnt/user-data/outputs/report.md"}, "id": "c1"})
|
|
assert is_handwritten_report_call({"name": "begin_write_file", "args": {"path": "/mnt/user-data/outputs/report.md"}, "id": "c1"})
|
|
assert is_handwritten_report_call({"name": "write_file", "args": {"path": "/mnt/user-data/outputs/analysis.md"}, "id": "c1"})
|
|
assert is_handwritten_report_call({"name": "str_replace", "args": {"path": "report.md"}, "id": "c1"})
|
|
assert not is_handwritten_report_call({"name": "write_file", "args": {"path": "/mnt/user-data/workspace/notes.md"}, "id": "c1"})
|
|
assert not is_handwritten_report_call({"name": "web_search", "args": {"query": "课题"}, "id": "c1"})
|
|
|
|
|
|
def test_extract_topic_uses_latest_visible_user_text() -> None:
|
|
topic = extract_report_topic(
|
|
[
|
|
HumanMessage(content="旧课题"),
|
|
HumanMessage(content="hidden", additional_kwargs={"hide_from_ui": True}),
|
|
HumanMessage(content="写一份新能源汽车产业报告，重点看产能"),
|
|
]
|
|
)
|
|
assert topic.startswith("写一份新能源汽车产业报告")
|
|
|
|
|
|
def test_extract_topic_skips_solution_generation_control_message() -> None:
|
|
topic = extract_report_topic(
|
|
[
|
|
HumanMessage(content="为普通大众写一份 AI 趋势与个人应对方案"),
|
|
AIMessage(content="需求已经足够。"),
|
|
HumanMessage(
|
|
content=(
|
|
"直接生成方案"
|
|
f"{SOLUTION_REQUIREMENTS_CONFIRMATION_TOKEN}"
|
|
f"{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"
|
|
)
|
|
),
|
|
]
|
|
)
|
|
assert topic == "为普通大众写一份 AI 趋势与个人应对方案"
|
|
|
|
|
|
def test_extract_topic_skips_missing_button_complaint() -> None:
|
|
topic = extract_report_topic(
|
|
[
|
|
HumanMessage(content="为普通大众写一份 AI 趋势与个人应对方案"),
|
|
AIMessage(content="我已理解需求。"),
|
|
HumanMessage(content="结束问答并生成方案按钮 没有按钮，我走不了下一步"),
|
|
HumanMessage(content=f"直接生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
]
|
|
)
|
|
assert topic == "为普通大众写一份 AI 趋势与个人应对方案"
|
|
|
|
|
|
def test_extract_topic_skips_missing_report_display_complaint() -> None:
|
|
topic = extract_report_topic(
|
|
[
|
|
HumanMessage(content="为普通大众写一份 AI 趋势与个人应对方案"),
|
|
AIMessage(content="正在处理。"),
|
|
HumanMessage(content="我根本没看到报告，右侧没显示，报告在哪？"),
|
|
HumanMessage(content=f"重新生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
]
|
|
)
|
|
assert topic == "为普通大众写一份 AI 趋势与个人应对方案"
|
|
|
|
|
|
def test_rewrite_write_file_to_pipeline() -> None:
|
|
last = AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "# 草稿"},
|
|
"type": "tool_call",
|
|
},
|
|
{
|
|
"id": "c2",
|
|
"name": "present_files",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
},
|
|
],
|
|
)
|
|
|
|
patched = rewrite_handwritten_report_calls(last, topic="新能源汽车")
|
|
|
|
assert patched is not None
|
|
assert len(patched.tool_calls) == 1
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
assert patched.tool_calls[0]["args"] == {"topic": "新能源汽车"}
|
|
assert patched.tool_calls[0]["id"] == "c1"
|
|
|
|
|
|
def test_rewrite_keeps_existing_pipeline_and_drops_handwritten() -> None:
|
|
last = AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c1",
|
|
"name": TOOL_NAME,
|
|
"args": {"topic": "课题"},
|
|
"type": "tool_call",
|
|
},
|
|
{
|
|
"id": "c2",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "x"},
|
|
"type": "tool_call",
|
|
},
|
|
],
|
|
)
|
|
|
|
patched = rewrite_handwritten_report_calls(last, topic="课题")
|
|
|
|
assert patched is not None
|
|
assert [call["name"] for call in patched.tool_calls] == [TOOL_NAME]
|
|
|
|
|
|
def test_rewrite_leaves_chat_and_search_alone() -> None:
|
|
chat = AIMessage(content="先确认一下研究范围。")
|
|
assert rewrite_handwritten_report_calls(chat, topic="课题") is None
|
|
|
|
search = AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "s1", "name": "web_search", "args": {"query": "课题"}, "type": "tool_call"},
|
|
],
|
|
)
|
|
assert rewrite_handwritten_report_calls(search, topic="课题") is None
|
|
|
|
pipeline = AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"},
|
|
],
|
|
)
|
|
assert rewrite_handwritten_report_calls(pipeline, topic="课题") is None
|
|
|
|
|
|
def test_middleware_allows_chat_before_writing() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="先聊聊这个课题"),
|
|
AIMessage(content="好的，我们先明确范围。"),
|
|
]
|
|
}
|
|
assert middleware.after_model(state, runtime=SimpleNamespace()) is None
|
|
|
|
|
|
def test_required_pipeline_replaces_plain_fake_completion() -> None:
|
|
middleware = WritingModeReportMiddleware(require_pipeline=True)
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(content="方案已经生成并展示在右侧。"),
|
|
],
|
|
"artifacts": [],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.content == ""
|
|
assert len(patched.tool_calls) == 1
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
assert patched.tool_calls[0]["args"]["topic"] == "写一份新能源汽车报告"
|
|
|
|
|
|
def test_required_pipeline_prefers_validated_intent_over_button_complaint() -> None:
|
|
middleware = WritingModeReportMiddleware(
|
|
require_pipeline=True,
|
|
solution_intent={
|
|
"goal": "帮助普通大众理解 AI 发展趋势及个人影响",
|
|
"audience": "不具备技术背景的普通大众",
|
|
"success_outcome": "读者形成清晰认知并能采取行动",
|
|
"scope": "技术突破、就业、生活与应对策略",
|
|
"constraints": "不展开产业竞争和实现细节",
|
|
"deliverable": "结构化 Markdown 报告",
|
|
"acceptance_criteria": "通俗、可靠、可操作",
|
|
"open_questions": ["篇幅尚未确认"],
|
|
},
|
|
)
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="结束问答并生成方案按钮 没有按钮"),
|
|
HumanMessage(content=f"直接生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
AIMessage(content="方案已经生成。"),
|
|
],
|
|
"artifacts": [],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
args = update["messages"][0].tool_calls[0]["args"]
|
|
assert args["topic"] == "帮助普通大众理解 AI 发展趋势及个人影响"
|
|
assert "结构化 Markdown 报告" in args["focus"]
|
|
assert "暂定假设" in args["focus"]
|
|
|
|
|
|
def test_required_pipeline_replaces_clarification_in_generation_turn() -> None:
|
|
middleware = WritingModeReportMiddleware(require_pipeline=True)
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "q1",
|
|
"name": "ask_clarification",
|
|
"args": {
|
|
"question": "预算是多少？",
|
|
"clarification_type": "missing_info",
|
|
"options": ["100 万", "500 万", "其他/我补充"],
|
|
},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": [],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert [call["name"] for call in patched.tool_calls] == [TOOL_NAME]
|
|
|
|
|
|
def test_required_pipeline_accepts_completion_only_with_report_artifact() -> None:
|
|
middleware = WritingModeReportMiddleware(require_pipeline=True)
|
|
messages = [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content="研究报告已完成：已写入 report.md。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
AIMessage(content="方案已经生成并展示在右侧。"),
|
|
]
|
|
|
|
missing = middleware.after_model(
|
|
{"messages": messages, "artifacts": []},
|
|
runtime=SimpleNamespace(),
|
|
)
|
|
assert missing is not None
|
|
assert missing["jump_to"] == "end"
|
|
assert missing["messages"][0].tool_calls == []
|
|
assert "未成功" in missing["messages"][0].content
|
|
|
|
completed = middleware.after_model(
|
|
{"messages": messages, "artifacts": ["/mnt/user-data/outputs/report.md"]},
|
|
runtime=SimpleNamespace(),
|
|
)
|
|
assert completed is None
|
|
|
|
|
|
def test_required_pipeline_failure_ends_without_automatic_retry() -> None:
|
|
middleware = WritingModeReportMiddleware(require_pipeline=True)
|
|
messages = [
|
|
HumanMessage(content=f"直接生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content="Error: 深度研究报告生成失败：模型暂时不可用。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
AIMessage(
|
|
content="我再试一次。",
|
|
tool_calls=[
|
|
{"id": "c2", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
]
|
|
|
|
update = middleware.after_model(
|
|
{"messages": messages, "artifacts": []},
|
|
runtime=SimpleNamespace(),
|
|
)
|
|
|
|
assert update is not None
|
|
assert update["jump_to"] == "end"
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls == []
|
|
assert "未成功" in patched.content
|
|
assert "重新生成方案" in patched.content
|
|
|
|
|
|
def test_new_user_retry_is_not_blocked_by_prior_pipeline_failure() -> None:
|
|
middleware = WritingModeReportMiddleware(require_pipeline=True)
|
|
messages = [
|
|
HumanMessage(content=f"直接生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content="Error: 深度研究报告生成失败：模型暂时不可用。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
AIMessage(content="本轮生成失败，请重试。"),
|
|
HumanMessage(content=f"重新生成方案{SOLUTION_REQUIREMENTS_FORCE_TOKEN}"),
|
|
AIMessage(content="开始重试。"),
|
|
]
|
|
|
|
update = middleware.after_model(
|
|
{"messages": messages, "artifacts": []},
|
|
runtime=SimpleNamespace(),
|
|
)
|
|
|
|
assert update is not None
|
|
assert update.get("jump_to") is None
|
|
assert update["messages"][0].tool_calls[0]["name"] == TOOL_NAME
|
|
|
|
|
|
def test_middleware_rewrites_write_file_on_after_model() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "# 报告"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
assert patched.tool_calls[0]["args"]["topic"].startswith("写一份新能源汽车报告")
|
|
|
|
|
|
def test_wrap_tool_call_rejects_handwritten_report() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
request = SimpleNamespace(
|
|
tool_call={
|
|
"id": "c1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "x"},
|
|
}
|
|
)
|
|
called = {"handler": False}
|
|
|
|
def handler(_request):
|
|
called["handler"] = True
|
|
return ToolMessage(content="wrote", tool_call_id="c1", name="write_file")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert called["handler"] is False
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.content.startswith("Error:")
|
|
assert TOOL_NAME in result.content
|
|
|
|
|
|
def test_pipeline_success_is_detected() -> None:
|
|
messages = [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
]
|
|
assert pipeline_already_succeeded(messages) is True
|
|
assert pipeline_already_succeeded(messages[:-1]) is False
|
|
failed = [
|
|
*messages[:-1],
|
|
ToolMessage(content="Error: 深度研究报告生成失败。", tool_call_id="c1", name=TOOL_NAME),
|
|
]
|
|
assert pipeline_already_succeeded(failed) is False
|
|
refused = [
|
|
*messages[:-1],
|
|
ToolMessage(
|
|
content="当前仍处于正常问答阶段，不能生成首份方案。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
]
|
|
assert pipeline_already_succeeded(refused) is False
|
|
|
|
|
|
def test_policy_refusal_never_injects_completed_report_reminder() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="开始"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "当前课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content="当前仍处于正常问答阶段，不能生成首份方案。请继续自然理解用户意图。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
HumanMessage(content="结束问答并生成方案按钮 没有按钮"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c2", "name": TOOL_NAME, "args": {"topic": "当前课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": [],
|
|
}
|
|
|
|
assert middleware.after_model(state, runtime=SimpleNamespace()) is None
|
|
|
|
|
|
def test_reread_detects_outputs_markdown() -> None:
|
|
assert is_report_reread_call({"name": "read_file", "args": {"path": "/mnt/user-data/outputs/report.md"}, "id": "c1"})
|
|
assert is_report_reread_call({"name": "read_file", "args": {"path": "/mnt/user-data/outputs/ai-agents-paper.md"}, "id": "c1"})
|
|
assert not is_report_reread_call({"name": "read_file", "args": {"path": "/mnt/user-data/workspace/notes.md"}, "id": "c1"})
|
|
assert not is_report_reread_call({"name": "write_file", "args": {"path": "/mnt/user-data/outputs/report.md"}, "id": "c1"})
|
|
|
|
|
|
def test_rewrite_after_success_drops_deliverable_tools() -> None:
|
|
last = AIMessage(
|
|
content="我再重写一版",
|
|
tool_calls=[
|
|
{
|
|
"id": "c2",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/ai-agents-paper.md", "content": "# 新稿"},
|
|
"type": "tool_call",
|
|
},
|
|
{
|
|
"id": "c3",
|
|
"name": TOOL_NAME,
|
|
"args": {"topic": "课题"},
|
|
"type": "tool_call",
|
|
},
|
|
],
|
|
)
|
|
|
|
patched = rewrite_handwritten_report_calls(last, topic="课题", pipeline_already_done=True)
|
|
|
|
assert patched is not None
|
|
assert patched.tool_calls == []
|
|
|
|
|
|
def test_rewrite_after_success_drops_reread() -> None:
|
|
last = AIMessage(
|
|
content="先看看报告",
|
|
tool_calls=[
|
|
{
|
|
"id": "c2",
|
|
"name": "read_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
)
|
|
|
|
patched = rewrite_handwritten_report_calls(last, topic="课题", pipeline_already_done=True)
|
|
|
|
assert patched is not None
|
|
assert patched.tool_calls == []
|
|
assert patched.content == ""
|
|
|
|
|
|
def test_middleware_nudges_summary_after_successful_pipeline() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
AIMessage(
|
|
content="报告太泛化，我要重写",
|
|
tool_calls=[
|
|
{
|
|
"id": "c2",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/ai-agents-paper.md", "content": "# 重写"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": ["/mnt/user-data/outputs/report.md"],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
assert update["jump_to"] == "model"
|
|
assert update["messages"][0].tool_calls == []
|
|
assert update["messages"][0].content == ""
|
|
nudge = update["messages"][1]
|
|
assert nudge.additional_kwargs["hide_from_ui"] is True
|
|
assert nudge.additional_kwargs["writing_mode_summary_nudge"] is True
|
|
assert "2-3" in nudge.content
|
|
|
|
|
|
def test_middleware_ends_turn_after_ignored_summary_nudge() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
HumanMessage(
|
|
content="请总结",
|
|
additional_kwargs={"hide_from_ui": True, "writing_mode_summary_nudge": True},
|
|
),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c2",
|
|
"name": "read_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": ["/mnt/user-data/outputs/report.md"],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
assert update["jump_to"] == "end"
|
|
assert update["messages"][0].tool_calls == []
|
|
|
|
|
|
def test_wrap_tool_call_rejects_rerun_after_success() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
prior = [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
]
|
|
request = SimpleNamespace(
|
|
tool_call={"id": "c2", "name": TOOL_NAME, "args": {"topic": "课题"}},
|
|
state={"messages": prior, "artifacts": ["/mnt/user-data/outputs/report.md"]},
|
|
)
|
|
called = {"handler": False}
|
|
|
|
def handler(_request):
|
|
called["handler"] = True
|
|
return ToolMessage(content="ran", tool_call_id="c2", name=TOOL_NAME)
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert called["handler"] is False
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.content.startswith("Error:")
|
|
assert "已经完成" in result.content
|
|
assert "请立即调用" not in result.content
|
|
|
|
|
|
def test_wrap_tool_call_does_not_trust_success_text_without_artifact() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
prior = [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content="研究报告已完成：已写入 report.md。",
|
|
tool_call_id="c1",
|
|
name=TOOL_NAME,
|
|
),
|
|
]
|
|
request = SimpleNamespace(
|
|
tool_call={"id": "c2", "name": TOOL_NAME, "args": {"topic": "课题"}},
|
|
state={"messages": prior, "artifacts": []},
|
|
)
|
|
expected = ToolMessage(content="ran", tool_call_id="c2", name=TOOL_NAME)
|
|
|
|
result = middleware.wrap_tool_call(request, lambda _request: expected)
|
|
|
|
assert result is expected
|
|
|
|
|
|
def test_wrap_tool_call_allows_new_workbench_write_after_success() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
prior = [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
HumanMessage(
|
|
content="请立即开始撰写研究报告。按新结构撰写研究报告。",
|
|
additional_kwargs={"hide_from_ui": True, "workbench_write_now": True, "workbench_topic": "课题"},
|
|
),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c2", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
]
|
|
request = SimpleNamespace(
|
|
tool_call={"id": "c2", "name": TOOL_NAME, "args": {"topic": "课题"}},
|
|
state={"messages": prior, "artifacts": ["/mnt/user-data/outputs/report.md"]},
|
|
)
|
|
expected = ToolMessage(content="ran", tool_call_id="c2", name=TOOL_NAME)
|
|
|
|
result = middleware.wrap_tool_call(request, lambda _request: expected)
|
|
|
|
assert result is expected
|
|
|
|
|
|
def test_wrap_tool_call_rejects_reread_after_success() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
prior = [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
]
|
|
request = SimpleNamespace(
|
|
tool_call={
|
|
"id": "c2",
|
|
"name": "read_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
},
|
|
state={"messages": prior, "artifacts": ["/mnt/user-data/outputs/report.md"]},
|
|
)
|
|
called = {"handler": False}
|
|
|
|
def handler(_request):
|
|
called["handler"] = True
|
|
return ToolMessage(content="read", tool_call_id="c2", name="read_file")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert called["handler"] is False
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.content.startswith("Error:")
|
|
assert "已经完成" in result.content
|
|
|
|
|
|
def test_wrap_tool_call_allows_unrelated_writes() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
request = SimpleNamespace(
|
|
tool_call={
|
|
"id": "c1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/workspace/notes.md", "content": "x"},
|
|
}
|
|
)
|
|
expected = ToolMessage(content="wrote", tool_call_id="c1", name="write_file")
|
|
|
|
result = middleware.wrap_tool_call(request, lambda _request: expected)
|
|
|
|
assert result is expected
|
|
|
|
|
|
def test_extract_topic_skips_clarification_answers() -> None:
|
|
"""澄清回答不是课题：向上取真正的写作请求原文。"""
|
|
topic = extract_report_topic(
|
|
[
|
|
HumanMessage(content="写人工智能领域的论文md"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "q1",
|
|
"name": "ask_clarification",
|
|
"args": {"question": "方向？"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
ToolMessage(content="需要澄清：方向？", tool_call_id="q1", name="ask_clarification"),
|
|
HumanMessage(content="方向：AI Agent；用途：学术期刊论文；篇幅：短篇2000-4000字"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
)
|
|
assert topic.startswith("写人工智能领域的论文md")
|
|
|
|
|
|
def test_bash_heredoc_write_is_redirected_to_pipeline() -> None:
|
|
"""模型被剥离不了工具时会绕道 bash heredoc 手写报告——同样改写为管线。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写人工智能领域的论文md"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "b1",
|
|
"name": "bash",
|
|
"args": {"command": "cat > /mnt/user-data/outputs/report.md << 'ENDOFFILE'\n# 报告全文……\nENDOFFILE"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
assert patched.tool_calls[0]["args"]["topic"].startswith("写人工智能领域的论文md")
|
|
|
|
|
|
def test_bash_python_open_write_is_redirected_to_pipeline() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "b2",
|
|
"name": "bash",
|
|
"args": {"command": ("python3 << 'PYEOF'\ncontent = r'''# 报告'''\nwith open('/mnt/user-data/outputs/report.md', 'w', encoding='utf-8') as f:\n f.write(content)\nPYEOF")},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
|
|
|
|
def test_bash_readonly_and_non_outputs_commands_are_left_alone() -> None:
|
|
"""mkdir / wc / cat 读取、非 outputs 的 .md 写入都不是手写报告。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
for command in (
|
|
"mkdir -p /mnt/user-data/outputs",
|
|
"wc -m /mnt/user-data/outputs/report.md",
|
|
"cat /mnt/user-data/outputs/report.md",
|
|
"echo hi > /tmp/scratch.md",
|
|
):
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "b3", "name": "bash", "args": {"command": command}, "type": "tool_call"}],
|
|
),
|
|
]
|
|
}
|
|
assert middleware.after_model(state, runtime=SimpleNamespace()) is None, command
|
|
|
|
|
|
def test_bash_write_after_success_is_dropped() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：基于 3 条资料…", tool_call_id="c1", name=TOOL_NAME),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "b4",
|
|
"name": "bash",
|
|
"args": {"command": "echo '# 新版' > /mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": ["/mnt/user-data/outputs/report.md"],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls == []
|
|
nudged = update["messages"][1]
|
|
assert getattr(nudged, "additional_kwargs", {}).get("writing_mode_summary_nudge") is True
|
|
|
|
|
|
def test_wrap_tool_call_rejects_bash_report_write() -> None:
|
|
"""漏网到执行层的 bash 手写（如裸 report.md 重定向）也要拒绝并引导走管线。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
request = SimpleNamespace(
|
|
tool_call={
|
|
"id": "b5",
|
|
"name": "bash",
|
|
"args": {"command": "cat > report.md << 'EOF'\n# x\nEOF"},
|
|
}
|
|
)
|
|
called = {"handler": False}
|
|
|
|
def handler(_request):
|
|
called["handler"] = True
|
|
return ToolMessage(content="ran", tool_call_id="b5", name="bash")
|
|
|
|
result = middleware.wrap_tool_call(request, handler)
|
|
|
|
assert called["handler"] is False
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.content.startswith("Error:")
|
|
assert TOOL_NAME in result.content
|
|
|
|
|
|
def test_wrap_model_call_strips_handwrite_tools_from_binding() -> None:
|
|
"""首报轮从模型绑定中剥离 write_file/str_replace/bash——bash 是唯一能写
|
|
文件的执行类工具（实测模型会把全文起草进 heredoc 参数流式几十秒后才被
|
|
after_model 拦截），剥离后模型没有任何可发起文件写入的工具。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
|
|
def make_request(tools):
|
|
return SimpleNamespace(
|
|
tools=tools,
|
|
override=lambda **kwargs: SimpleNamespace(tools=kwargs.get("tools", [])),
|
|
)
|
|
|
|
seen: dict[str, list] = {}
|
|
|
|
def handler(request):
|
|
seen["tools"] = list(request.tools)
|
|
return SimpleNamespace()
|
|
|
|
middleware.wrap_model_call(
|
|
make_request(
|
|
[
|
|
SimpleNamespace(name="write_file"),
|
|
SimpleNamespace(name="begin_write_file"),
|
|
SimpleNamespace(name="web_search"),
|
|
SimpleNamespace(name="str_replace"),
|
|
SimpleNamespace(name="bash"),
|
|
SimpleNamespace(name=TOOL_NAME),
|
|
]
|
|
),
|
|
handler,
|
|
)
|
|
|
|
assert [getattr(t, "name", "") for t in seen["tools"]] == ["web_search", TOOL_NAME]
|
|
|
|
|
|
def test_wrap_model_call_passes_through_when_no_handwrite_tools() -> None:
|
|
middleware = WritingModeReportMiddleware()
|
|
request = SimpleNamespace(
|
|
tools=[SimpleNamespace(name="web_search"), SimpleNamespace(name=TOOL_NAME)],
|
|
override=lambda **kwargs: SimpleNamespace(tools=kwargs.get("tools", [])),
|
|
)
|
|
|
|
def handler(passed):
|
|
assert passed is request
|
|
return "unchanged"
|
|
|
|
assert middleware.wrap_model_call(request, handler) == "unchanged"
|
|
|
|
|
|
def test_extract_focus_joins_clarification_answers() -> None:
|
|
focus = extract_report_focus(
|
|
[
|
|
HumanMessage(content="写人工智能领域的论文md"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "q1", "name": "ask_clarification", "args": {"question": "方向？"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="需要澄清", tool_call_id="q1", name="ask_clarification"),
|
|
HumanMessage(content="主题：LLM的能力边界；类型：观点评论型"),
|
|
],
|
|
topic="写人工智能领域的论文md",
|
|
)
|
|
assert "LLM的能力边界" in focus
|
|
assert "观点评论型" in focus
|
|
assert "写人工智能领域的论文md" not in focus
|
|
|
|
|
|
def test_rewrite_carries_focus_from_clarification_answers() -> None:
|
|
"""改写进管线的调用要带上澄清回答（方向/类型/篇幅），否则管线只拿到
|
|
原始请求一句泛写。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写人工智能领域的论文md，有不懂的可以问我"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "q1", "name": "ask_clarification", "args": {"question": "主题？"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="需要澄清", tool_call_id="q1", name="ask_clarification"),
|
|
HumanMessage(content="主题：LLM的能力边界；篇幅：3000-5000字"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c1",
|
|
"name": "bash",
|
|
"args": {"command": "cat > /mnt/user-data/outputs/report.md << 'EOF'\n# x\nEOF"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls[0]["name"] == TOOL_NAME
|
|
assert patched.tool_calls[0]["args"]["topic"].startswith("写人工智能领域的论文md")
|
|
assert "LLM的能力边界" in patched.tool_calls[0]["args"]["focus"]
|
|
|
|
|
|
def test_second_offense_after_nudge_gets_fallback_summary() -> None:
|
|
"""模型无视总结提示仍要继续：丢弃调用后必须留一段兜底总结，不能静默结束。"""
|
|
middleware = WritingModeReportMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="写一份新能源汽车报告"),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "c1", "name": TOOL_NAME, "args": {"topic": "课题"}, "type": "tool_call"}],
|
|
),
|
|
ToolMessage(content="研究报告已完成：已写入 report.md。", tool_call_id="c1", name=TOOL_NAME),
|
|
HumanMessage(
|
|
content="请总结",
|
|
additional_kwargs={"hide_from_ui": True, "writing_mode_summary_nudge": True},
|
|
),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "c2",
|
|
"name": "read_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
],
|
|
"artifacts": ["/mnt/user-data/outputs/report.md"],
|
|
}
|
|
|
|
update = middleware.after_model(state, runtime=SimpleNamespace())
|
|
|
|
assert update is not None
|
|
assert update["jump_to"] == "end"
|
|
patched = update["messages"][0]
|
|
assert patched.tool_calls == []
|
|
assert "report.md" in patched.content
|
|
assert "已完成" in patched.content
|