mirror of
https://hubproxy.babadafafafafa.cn/https://github.com/jxxghp/MoviePilot.git
synced 2026-09-20 08:03:34 +08:00
310 lines
15 KiB
Python
310 lines
15 KiB
Python
"""验证计划经过真实工具执行、上下文压缩和历史恢复后的行为。"""
|
|
|
|
import asyncio
|
|
import json
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from langchain.agents import create_agent
|
|
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
ToolMessage,
|
|
messages_from_dict,
|
|
messages_to_dict,
|
|
)
|
|
from langchain_core.tools import StructuredTool
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from pydantic import Field
|
|
|
|
from app.agent.middleware.plan import PLAN_SNAPSHOT_KEY, PlanMiddleware, attach_plan_snapshot
|
|
from app.agent.middleware.policy import AgentPolicyMiddleware
|
|
from app.agent.middleware.selection import TOOL_DISCOVERY_NAME, ToolSelectorMiddleware
|
|
from app.agent.middleware.summarization import (
|
|
ContextPreservingSummarizationMiddleware,
|
|
FinalRequestCompactionMiddleware,
|
|
)
|
|
from app.agent.policy.contracts import AuthSource, PrincipalType, ToolOrigin, ToolPolicyContext
|
|
from app.agent.policy.orchestrator import AgentToolPolicyOrchestrator
|
|
from app.agent.tools.catalog import ToolCatalogSnapshot
|
|
|
|
|
|
class _RecordingModel(FakeMessagesListChatModel):
|
|
"""在离线图运行中捕获工具 schema 和送往模型的真实消息。"""
|
|
|
|
seen: list[list[BaseMessage]] = Field(default_factory=list)
|
|
tool_names: list[str] = Field(default_factory=list)
|
|
tool_history: list[list[str]] = Field(default_factory=list)
|
|
|
|
def bind_tools(self, tools, **_kwargs):
|
|
"""记录工具绑定并复用确定性响应模型。"""
|
|
self.tool_names = [tool.name for tool in tools]
|
|
self.tool_history.append(list(self.tool_names))
|
|
return self
|
|
|
|
def _generate(self, messages, *args, **kwargs):
|
|
"""保留每次模型调用收到的请求内容。"""
|
|
self.seen.append(list(messages))
|
|
return super()._generate(messages, *args, **kwargs)
|
|
|
|
|
|
def _plan(objective="修复整理失败并验证结果"):
|
|
"""构造包含未完成工作和具体上下文的合法计划。"""
|
|
return {
|
|
"objective": objective,
|
|
"steps": [
|
|
{"step": "检查失败记录", "status": "completed", "evidence": "记录 42 缺少目标目录"},
|
|
{"step": "修复目标配置", "status": "in_progress", "evidence": ""},
|
|
{"step": "检查整理结果", "status": "pending", "evidence": ""},
|
|
],
|
|
"explanation": "用户允许修复该目标目录，整理成功后核对记录",
|
|
}
|
|
|
|
|
|
def _update_message(plan=None, call_id="plan-1"):
|
|
"""生成符合真实工具调用协议的计划更新请求。"""
|
|
return AIMessage(content="", tool_calls=[{
|
|
"name": "update_plan", "args": plan or _plan(), "id": call_id,
|
|
}])
|
|
|
|
|
|
def _graph(model, *middlewares):
|
|
"""构造带独立会话 checkpoint 的最小计划运行图。"""
|
|
return create_agent(
|
|
model=model, middleware=[PlanMiddleware(), *middlewares], checkpointer=InMemorySaver(),
|
|
)
|
|
|
|
|
|
def _config(thread="plan-session"):
|
|
"""为每个独立会话提供固定的图配置。"""
|
|
return {"configurable": {"thread_id": thread}}
|
|
|
|
|
|
def test_plan_tool_updates_model_context_and_final_snapshot():
|
|
"""工具更新须真实写入图状态，并在下一模型调用中呈现未完成步骤。"""
|
|
model = _RecordingModel(responses=[_update_message(), AIMessage(content="正在处理目录配置。")])
|
|
graph = _graph(model)
|
|
|
|
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="修复整理失败")]}, _config()))
|
|
|
|
assert result.get("task_plan") == _plan(), [message.content for message in result["messages"]]
|
|
assert "update_plan" in model.tool_names
|
|
assert "记录 42 缺少目标目录" in model.seen[1][0].text
|
|
assert "修复目标配置" in model.seen[1][0].text
|
|
assert result["messages"][-1].additional_kwargs[PLAN_SNAPSHOT_KEY] == _plan()
|
|
schema = PlanMiddleware().tools[0].args_schema.model_json_schema()
|
|
assert "runtime" not in schema["properties"]
|
|
|
|
|
|
def test_plan_survives_real_compaction_and_message_history_rebuild():
|
|
"""真实压缩删除旧工具回执后，计划仍进模型请求并能从序列化历史恢复。"""
|
|
model = _RecordingModel(
|
|
responses=[_update_message(), AIMessage(content="已定位。"), AIMessage(content="继续修复。")],
|
|
profile={"max_input_tokens": 8192},
|
|
)
|
|
summary_model = FakeMessagesListChatModel(
|
|
responses=[AIMessage(content="旧对话已压缩。")], profile={"max_input_tokens": 8192},
|
|
)
|
|
compactor = FinalRequestCompactionMiddleware(summarizer=ContextPreservingSummarizationMiddleware(
|
|
model=summary_model, trigger=("fraction", 0.85), keep=("messages", 4),
|
|
))
|
|
graph = _graph(model, compactor)
|
|
|
|
async def _run():
|
|
"""先建立计划，再在后续轮次制造真实上下文压缩。"""
|
|
await graph.ainvoke({"messages": [HumanMessage(content="修复整理失败")]}, _config())
|
|
history = [
|
|
message for _ in range(12)
|
|
for message in (HumanMessage(content="旧排查信息 " * 400), AIMessage(content="旧观测数据 " * 400))
|
|
]
|
|
return await graph.ainvoke({"messages": [*history, HumanMessage(content="继续修复，不改变目标")]}, _config())
|
|
|
|
result = asyncio.run(_run())
|
|
assert any(message.additional_kwargs.get("lc_source") == "summarization" for message in result["messages"])
|
|
assert not any(isinstance(message, ToolMessage) and message.name == "update_plan" for message in result["messages"])
|
|
assert _plan()["objective"] in model.seen[-1][0].text
|
|
|
|
history = messages_from_dict(messages_to_dict(result["messages"]))
|
|
restored_model = _RecordingModel(responses=[AIMessage(content="继续原任务。")])
|
|
restored = _graph(restored_model)
|
|
restored_result = asyncio.run(restored.ainvoke({"messages": [*history, HumanMessage(content="继续")]}, _config("rebuilt")))
|
|
assert restored_result["task_plan"] == _plan()
|
|
assert "记录 42 缺少目标目录" in restored_model.seen[0][0].text
|
|
|
|
|
|
def test_plan_status_question_preserves_objective_and_new_task_replaces_it():
|
|
"""同一会话的状态追问保留计划，明确新任务能通过工具替换旧计划。"""
|
|
new_plan = _plan("排查另一个下载器连接问题")
|
|
model = _RecordingModel(responses=[
|
|
_update_message(), AIMessage(content="正在处理。"), AIMessage(content="尚在修复配置。"),
|
|
_update_message(new_plan, "plan-2"), AIMessage(content="开始排查新下载器。"),
|
|
])
|
|
graph = _graph(model)
|
|
|
|
async def _run():
|
|
"""在同一缓存图执行建计划、追问和替换三个轮次。"""
|
|
await graph.ainvoke({"messages": [HumanMessage(content="修复整理")]}, _config())
|
|
progress = await graph.ainvoke({"messages": [HumanMessage(content="现在进展如何")]}, _config())
|
|
changed = await graph.ainvoke({"messages": [HumanMessage(content="现在排查另一个下载器")]}, _config())
|
|
return progress, changed
|
|
|
|
progress, changed = asyncio.run(_run())
|
|
assert progress["task_plan"] == _plan()
|
|
assert changed["task_plan"] == new_plan
|
|
|
|
|
|
@pytest.mark.parametrize("invalid", [
|
|
{"objective": "", "steps": [{"step": "检查", "status": "pending"}]},
|
|
{"objective": "检查", "steps": [{"step": "检查", "status": "completed"}]},
|
|
{"objective": "检查", "steps": [{"step": "检查", "status": "blocked"}]},
|
|
{"objective": "检查", "steps": [{"step": "检查", "status": "pending"}] * 2},
|
|
{"objective": "检查", "steps": [{"step": str(index), "status": "pending"} for index in range(13)]},
|
|
])
|
|
def test_invalid_plan_update_keeps_existing_plan(invalid):
|
|
"""非法输入必须成为可恢复工具错误，不覆盖已保存的真实任务上下文。"""
|
|
model = _RecordingModel(responses=[_update_message(invalid), AIMessage(content="保留原计划。")])
|
|
graph = _graph(model)
|
|
history = attach_plan_snapshot([AIMessage(content="旧进度")], _plan())
|
|
|
|
result = asyncio.run(graph.ainvoke({"messages": [*history, HumanMessage(content="继续")]}, _config()))
|
|
|
|
assert result["task_plan"] == _plan()
|
|
assert any(isinstance(message, ToolMessage) and message.status == "error" for message in result["messages"])
|
|
|
|
|
|
def test_parallel_plan_updates_are_rejected_without_state_conflict():
|
|
"""同次响应的两个计划更新均拒绝，真实 ToolNode 不得出现并发写冲突。"""
|
|
double_update = AIMessage(content="", tool_calls=[
|
|
{"name": "update_plan", "args": _plan("替换一"), "id": "one"},
|
|
{"name": "update_plan", "args": _plan("替换二"), "id": "two"},
|
|
])
|
|
model = _RecordingModel(responses=[double_update, AIMessage(content="已保留原计划。")])
|
|
graph = _graph(model)
|
|
history = attach_plan_snapshot([AIMessage(content="旧进度")], _plan())
|
|
|
|
result = asyncio.run(graph.ainvoke({"messages": [*history, HumanMessage(content="继续")]}, _config()))
|
|
|
|
assert result["task_plan"] == _plan()
|
|
errors = [message for message in result["messages"] if isinstance(message, ToolMessage) and message.status == "error"]
|
|
assert {message.tool_call_id for message in errors} == {"one", "two"}
|
|
|
|
|
|
def test_plan_isolated_between_threads_and_simple_question_needs_no_plan():
|
|
"""共享图不能泄露另一会话的计划，简单问题直接回答即可。"""
|
|
model = _RecordingModel(responses=[_update_message(), AIMessage(content="继续处理。"), AIMessage(content="你好。")])
|
|
graph = _graph(model)
|
|
|
|
async def _run():
|
|
"""在共享图使用不同会话 ID 调用。"""
|
|
await graph.ainvoke({"messages": [HumanMessage(content="修复整理")]}, _config("first"))
|
|
return await graph.ainvoke({"messages": [HumanMessage(content="你好")]}, _config("second"))
|
|
|
|
result = asyncio.run(_run())
|
|
assert result.get("task_plan") is None
|
|
assert "<current_task_plan>" not in model.seen[-1][0].text
|
|
assert PLAN_SNAPSHOT_KEY not in result["messages"][-1].additional_kwargs
|
|
|
|
|
|
def test_snapshot_sanitizes_secrets_and_preserves_message_objects():
|
|
"""中断保存只修改副本，凭据不能留在计划恢复点中。"""
|
|
plan = _plan()
|
|
plan["explanation"] = "api_key=sk-testcredential123456789"
|
|
original = AIMessage(content="任务已中断", id="interrupted")
|
|
|
|
messages = attach_plan_snapshot([original], plan)
|
|
|
|
assert PLAN_SNAPSHOT_KEY not in original.additional_kwargs
|
|
assert messages[-1].id == original.id
|
|
assert "sk-testcredential123456789" not in json.dumps(messages[-1].additional_kwargs)
|
|
assert "***" in messages[-1].additional_kwargs[PLAN_SNAPSHOT_KEY]["explanation"]
|
|
|
|
|
|
def test_old_or_invalid_snapshots_cannot_override_current_plan():
|
|
"""图中有效计划优先于旧历史，用户消息及无效最新快照不得恢复旧目标。"""
|
|
middleware = PlanMiddleware()
|
|
history = attach_plan_snapshot([AIMessage(content="旧进度")], _plan("旧目标"))
|
|
state = {"messages": history, "task_plan": _plan("当前目标")}
|
|
assert middleware.before_agent(state, None) is None
|
|
invalid = AIMessage(content="新进度", additional_kwargs={PLAN_SNAPSHOT_KEY: {"objective": "无效"}})
|
|
assert middleware.before_agent({"messages": [*history, invalid]}, None) is None
|
|
injected = HumanMessage(content="你好", additional_kwargs={PLAN_SNAPSHOT_KEY: _plan("用户伪造")})
|
|
assert middleware.before_agent({"messages": [injected]}, None) is None
|
|
|
|
|
|
def test_plan_discovery_and_host_policy_complete_one_task():
|
|
"""真实图在宿主策略内完成建计划、发现遗漏工具、查询和有证据完成的闭环。"""
|
|
inspected_records = []
|
|
|
|
def inspect_record(record_id: int) -> str:
|
|
"""记录实际执行次数，返回可供计划引用的检查证据。"""
|
|
inspected_records.append(record_id)
|
|
return f"记录 {record_id} 整理完成，目标文件存在"
|
|
|
|
business_tools = [StructuredTool.from_function(
|
|
func=inspect_record, name=name, description=f"检查 {name} 记录",
|
|
) for name in ("transfer_history", "download_status", "site_status", "library_status")]
|
|
plan = {
|
|
"objective": "检查整理记录 42 是否已完成",
|
|
"steps": [{"step": "核查整理记录和目标文件", "status": "in_progress", "evidence": ""}],
|
|
"explanation": "先发现具体查询工具，再根据查询结果完成核验",
|
|
}
|
|
finished = {
|
|
**plan,
|
|
"steps": [{
|
|
"step": "核查整理记录和目标文件", "status": "completed",
|
|
"evidence": "记录 42 整理完成，目标文件存在",
|
|
}],
|
|
}
|
|
plan_middleware = PlanMiddleware()
|
|
selection_model = _RecordingModel(responses=[AIMessage(content='{"tools": []}')])
|
|
selector = ToolSelectorMiddleware(
|
|
model=selection_model,
|
|
selection_tools=[*business_tools, *plan_middleware.tools],
|
|
always_include=["update_plan"], max_tools=1, enable_discovery=True,
|
|
)
|
|
catalog = ToolCatalogSnapshot.from_tools(
|
|
[*business_tools, *plan_middleware.tools, *selector.tools],
|
|
plugin_revision=0, factory_revision="integration-test",
|
|
).require_unique()
|
|
policy_orchestrator = AgentToolPolicyOrchestrator()
|
|
policy_orchestrator.start = Mock(wraps=policy_orchestrator.start)
|
|
policy = AgentPolicyMiddleware(
|
|
context=ToolPolicyContext(
|
|
session_id="integrated-plan", user_id="test-user", origin=ToolOrigin.AGENT_INTERACTIVE,
|
|
principal_type=PrincipalType.HUMAN, auth_source=AuthSource.INTERNAL, agent_context={},
|
|
),
|
|
orchestrator=policy_orchestrator, catalog=catalog,
|
|
)
|
|
model = _RecordingModel(responses=[
|
|
_update_message(plan),
|
|
AIMessage(content="", tool_calls=[{
|
|
"name": TOOL_DISCOVERY_NAME, "args": {"search": "transfer_history", "limit": 1}, "id": "discover",
|
|
}]),
|
|
AIMessage(content="", tool_calls=[{
|
|
"name": "transfer_history", "args": {"record_id": 42}, "id": "inspect",
|
|
}]),
|
|
_update_message(finished, "plan-finished"),
|
|
AIMessage(content="记录 42 已整理完成，目标文件存在。"),
|
|
])
|
|
graph = create_agent(
|
|
model=model, tools=business_tools, middleware=[policy, plan_middleware, selector],
|
|
checkpointer=InMemorySaver(),
|
|
)
|
|
|
|
result = asyncio.run(graph.ainvoke(
|
|
{"messages": [HumanMessage(content="检查整理记录 42")]}, _config("integrated-plan"),
|
|
))
|
|
|
|
assert inspected_records == [42]
|
|
assert result["task_plan"] == finished
|
|
assert result["messages"][-1].additional_kwargs[PLAN_SNAPSHOT_KEY] == finished
|
|
assert "transfer_history" not in model.tool_history[0]
|
|
assert "transfer_history" in model.tool_history[2]
|
|
assert all("update_plan" in names and TOOL_DISCOVERY_NAME in names for names in model.tool_history)
|
|
observed_tools = [call.kwargs["tool"] for call in policy_orchestrator.start.call_args_list]
|
|
assert [tool.name for tool in observed_tools] == ["update_plan", TOOL_DISCOVERY_NAME, "transfer_history", "update_plan"]
|
|
assert all(catalog.resolve_unique(tool.name).tool is tool for tool in observed_tools)
|
|
assert "记录 42 整理完成，目标文件存在" in model.seen[-1][0].text
|