133 lines
7.3 KiB
Python
133 lines
7.3 KiB
Python
"""Opt-in real-model multi-turn probe with checkpointing and accelerated summary triggers.
|
|
|
|
Run from backend: .venv/Scripts/python scripts/probe_multiturn_compression.py --live
|
|
Only synthetic project facts are sent. Config changes stay inside this process.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from deerflow.agents.lead_agent.agent import _create_summarization_middleware
|
|
from deerflow.agents.middlewares.context_window_guard_middleware import ContextWindowGuardMiddleware
|
|
from deerflow.agents.middlewares.llm_error_handling_middleware import is_llm_error_message
|
|
from deerflow.agents.middlewares.summarization_middleware import _clear_global_summary_cache
|
|
from deerflow.agents.thread_state import ThreadState
|
|
from deerflow.config.app_config import get_app_config
|
|
from deerflow.models.factory import create_chat_model
|
|
|
|
|
|
def visible_text(message):
|
|
text = re.sub(r"<think>.*?</think>", "", message.text, flags=re.S | re.I)
|
|
return re.split(r"<think>", text, flags=re.I)[0].strip()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--live", action="store_true", required=True)
|
|
parser.add_argument("--model", default="MiniMax-M2.7")
|
|
parser.add_argument("--output", default=None)
|
|
args = parser.parse_args()
|
|
logging.disable(logging.CRITICAL)
|
|
config = get_app_config().model_copy(deep=True)
|
|
selected = config.get_model_config(args.model)
|
|
selected.max_retries = 0
|
|
selected.request_timeout = 90
|
|
selected.max_tokens = 2048
|
|
# Lower only the summary trigger for this probe; never invent an input ceiling.
|
|
from deerflow.config.summarization_config import ContextSize
|
|
|
|
config.summarization.trigger = [ContextSize(type="messages", value=8)]
|
|
config.summarization.keep = ContextSize(type="messages", value=2)
|
|
config.summarization.model_name = args.model
|
|
main_model = create_chat_model(name=args.model, app_config=config, thinking_enabled=False, force_disable_thinking=True)
|
|
prompts = [
|
|
"请记住本次项目：代号 ORBIT-42；预算35万元；截止日期2026-12-20。只回复已记住。",
|
|
"补充范围：系统需支持资料检索、问答和报告导出。只回复已记住。",
|
|
"补充约束：不能上传真实客户数据，只能用合成测试材料。只回复已记住。",
|
|
"补充验收项：检查引用来源、输出格式和接口稳定性。只回复已记住。",
|
|
"请仅用一句话复述项目代号、预算和截止日期。",
|
|
"更新项目预算为48万元，代号和其他信息保持不变。只回复已更新。",
|
|
"补充约束：报告必须用中文，先写结论再写依据。只回复已记住。",
|
|
"更新截止日期为2026-12-28，其他信息不变。只回复已更新。",
|
|
"验收补充：必须覆盖多轮问答和长材料。只回复已记住。",
|
|
"继续补充：需要留存测试报告，不能删除已有聊天记录。只回复已记住。",
|
|
"补充：测试使用合成材料，不更改刚才确认的代号、预算和日期。只回复已记住。",
|
|
"请仅输出JSON，包含code、budget、deadline，分别给出当前项目代号、最新预算（万元数字）、最新截止日期。",
|
|
]
|
|
saver = InMemorySaver()
|
|
rows = []
|
|
started = time.monotonic()
|
|
report = {"model": args.model, "trigger_messages": 8, "keep_messages": 2, "rows": rows}
|
|
_clear_global_summary_cache()
|
|
final = None
|
|
try:
|
|
for turn, prompt in enumerate(prompts, 1):
|
|
if turn == 9:
|
|
# Emulate a worker restart: only the persisted checkpoint survives.
|
|
_clear_global_summary_cache()
|
|
middleware = _create_summarization_middleware(app_config=config, force_enabled=True)
|
|
action = {"value": "passthrough"}
|
|
original_plan = middleware._plan_compression
|
|
|
|
def inspect_plan(request):
|
|
plan = original_plan(request)
|
|
action["value"] = plan[0] if plan else "passthrough"
|
|
return plan
|
|
|
|
middleware._plan_compression = inspect_plan
|
|
agent = create_agent(model=main_model, middleware=[middleware, ContextWindowGuardMiddleware()], state_schema=ThreadState, checkpointer=saver)
|
|
turn_started = time.monotonic()
|
|
previous_ids = set((final.get("context_summary_cache") or {}).get("summarized_ids", [])) if final else set()
|
|
final = agent.invoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "synthetic-multiturn-probe"}})
|
|
message = final["messages"][-1]
|
|
visible = visible_text(message)
|
|
row = {
|
|
"turn": turn,
|
|
"compression": action["value"],
|
|
"history_messages": len(final["messages"]),
|
|
"has_cache": bool(final.get("context_summary_cache")),
|
|
"visible": visible[:500],
|
|
"elapsed_seconds": round(time.monotonic() - turn_started, 2),
|
|
}
|
|
row["summary_updated"] = set((final.get("context_summary_cache") or {}).get("summarized_ids", [])) != previous_ids
|
|
if row["summary_updated"]:
|
|
row["summary_text"] = final["context_summary_cache"]["summary_text"]
|
|
rows.append(row)
|
|
print(json.dumps({key: value for key, value in row.items() if key != "summary_text"}, ensure_ascii=True), flush=True)
|
|
if not visible or is_llm_error_message(message):
|
|
raise RuntimeError("empty_or_error_answer")
|
|
if len(final["messages"]) != turn * 2:
|
|
raise RuntimeError("history_was_lost")
|
|
if action["value"] == "summarize" and not row["summary_updated"]:
|
|
raise RuntimeError("summary_did_not_complete")
|
|
if turn == 5 and not all(value in visible for value in ("ORBIT-42", "35", "2026")):
|
|
raise RuntimeError("first_summary_lost_project_facts")
|
|
text = visible_text(final["messages"][-1])
|
|
parsed = json.loads(re.search(r"\{.*\}", text, flags=re.S).group())
|
|
report["fact_check"] = parsed.get("code") == "ORBIT-42" and str(parsed.get("budget")) == "48" and parsed.get("deadline") == "2026-12-28"
|
|
report["summary_rounds"] = [row["turn"] for row in rows if row["compression"] == "summarize"]
|
|
report["reuse_rounds"] = [row["turn"] for row in rows if row["compression"] == "reuse"]
|
|
report["outcome"] = "passed" if report["fact_check"] and len(report["summary_rounds"]) >= 3 and report["reuse_rounds"] else "failed"
|
|
except Exception as exc:
|
|
report["outcome"] = type(exc).__name__
|
|
if str(exc) in {"empty_or_error_answer", "history_was_lost", "summary_did_not_complete", "first_summary_lost_project_facts"}:
|
|
report["failure_reason"] = str(exc)
|
|
report["elapsed_seconds"] = round(time.monotonic() - started, 2)
|
|
if args.output:
|
|
from pathlib import Path
|
|
|
|
Path(args.output).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps({key: value for key, value in report.items() if key != "rows"}, ensure_ascii=True), flush=True)
|
|
return 0 if report["outcome"] == "passed" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|