255 lines
14 KiB
Python
255 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""链式压缩块重构验证：原生历史重建 / 压缩节点 / 双次压缩块链 / payload 透传。
|
||
|
||
运行（Agent 目录下）: python tests/verify_compression_chain.py
|
||
覆盖:
|
||
T0 单元渲染/合并/token 估算
|
||
T1 _rebuild_turn_messages 原生重建顺序（并行合并/尾部去重/最终回复）
|
||
T2 build_native_history_messages 存储路径 + 游标跳过
|
||
T3 compressed_context 读写往返
|
||
T4 maybe_compress_context 单次压缩（范围/统计/RemoveMessage/块链）
|
||
T5 双次压缩（seq 递增、旧块保留、游标推进）
|
||
T6 llm payload reasoning_content 透传
|
||
"""
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
|
||
|
||
import agent.utils.agent_utils as au
|
||
from agent.history import tool_call_recorder as tcr
|
||
from agent.utils.agent_utils import load_compressed_context, save_compressed_context
|
||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||
|
||
PASSED = []
|
||
|
||
|
||
def check(name: str, cond: bool, detail: str = "") -> None:
|
||
if not cond:
|
||
raise AssertionError(f"FAIL {name}: {detail}")
|
||
PASSED.append(name)
|
||
print(f" ok {name}")
|
||
|
||
|
||
def section(t: str) -> None:
|
||
print(f"\n== {t} ==")
|
||
|
||
|
||
# ---------- 测试目录 ----------
|
||
tmp_root = tempfile.mkdtemp(prefix="verify_compress_")
|
||
SHARE_SID = "20990101_000000"
|
||
try:
|
||
au.SESSIONS_ROOT = os.path.join(tmp_root, "sessions")
|
||
os.makedirs(au.SESSIONS_ROOT, exist_ok=True)
|
||
sess_dir = os.path.join(au.SESSIONS_ROOT, SHARE_SID)
|
||
os.makedirs(sess_dir, exist_ok=True)
|
||
# 测试图片/文件（供 _native_user_message 绝对路径解析）
|
||
Path(os.path.join(au.SESSIONS_ROOT, "20990101_120001")).mkdir(exist_ok=True)
|
||
(Path(au.SESSIONS_ROOT) / "20990101_120001" / "in.png").write_bytes(b"\x89PNG\r\n")
|
||
(Path(au.SESSIONS_ROOT) / "20990101_120001" / "data.pdf").write_bytes(b"1 0 R")
|
||
|
||
# ---------- T0 ----------
|
||
section("T0 单元渲染/合并/token 估算")
|
||
from agent.core import nodes
|
||
check("ascii tokens", nodes._estimate_text_tokens("abcd") == 1, str(nodes._estimate_text_tokens("abcd")))
|
||
check("cjk tokens", nodes._estimate_text_tokens("中文文本") == 4)
|
||
merged = nodes._merge_units([
|
||
HumanMessage(content="m1"),
|
||
HumanMessage(content="m2"),
|
||
AIMessage(content="a1"),
|
||
AIMessage(content="a2", additional_kwargs={"reasoning_content": "r"},
|
||
tool_calls=[{"id": "t1", "name": "t", "args": {}}]),
|
||
ToolMessage(content="tool ok", tool_call_id="t1", name="t"),
|
||
])
|
||
check("合并单元数=3（human 合并 1 + ai 合并 1 + tool 1）", len(merged) == 3, str(len(merged)))
|
||
check("human 合并含两条文本", "m1" in merged[0] and "m2" in merged[0])
|
||
check("ai 合并含思考与调用", "r" in merged[1] and "工具调用" in merged[1])
|
||
|
||
# ---------- T1 ----------
|
||
section("T1 _rebuild_turn_messages 原生重建顺序/id")
|
||
TID1 = "20990101_120001"
|
||
turn_msgs = [
|
||
{"role": "user", "user": {"content": [
|
||
{"type": "text", "text": "帮我处理图片"},
|
||
{"type": "image_url", "relpath": f"{TID1}/in.png", "name": "in.png"},
|
||
{"type": "file", "relpath": f"{TID1}/data.pdf", "name": "data.pdf"},
|
||
]}},
|
||
{"role": "assistant", "assistant": {"text": "处理完成"}},
|
||
]
|
||
calls = [
|
||
{"id": "c1", "name": "read_image", "args": {"path": "in.png"}, "status": "done",
|
||
"thinking": "思考A", "reflection": "反思A", "result_text": "图片内容是文本",
|
||
"group_index": 0, "result_images": [], "result_files": []},
|
||
{"id": "c2", "name": "read_file", "args": {"path": "data.pdf"}, "status": "done",
|
||
"thinking": "", "reflection": "", "result_text": "PDF 正文",
|
||
"group_index": 0, "result_images": [], "result_files": ["out.txt"]},
|
||
{"id": "c3", "name": "call_subagent", "args": {"task": "x"}, "status": "done",
|
||
"thinking": "", "reflection": "", "result_text": "子任务完成",
|
||
"group_index": 1, "result_images": [], "result_files": []},
|
||
]
|
||
tail = [
|
||
{"thought": "思考A", "reflection": "反思A"}, # 与组内重复 → 去重
|
||
{"thought": "尾部思考", "reflection": ""}, # 未消费 → 注入
|
||
]
|
||
msgs = au._rebuild_turn_messages(TID1, turn_msgs, calls, tail)
|
||
kinds = [type(m).__name__ for m in msgs]
|
||
expect_kinds = ["HumanMessage", "AIMessage", "ToolMessage", "ToolMessage",
|
||
"AIMessage", "ToolMessage", "AIMessage", "AIMessage"]
|
||
check("顺序 user->AI(c1c2)->Tool2->AI(c3)->Tool1->tail->final",
|
||
kinds == expect_kinds, str(kinds))
|
||
check("user 多模态还原（text+image_url+file）", len(msgs[0].content) == 3,
|
||
str(msgs[0].content))
|
||
img_part = msgs[0].content[1]
|
||
check("image part 绝对路径", img_part["type"] == "image_url" and
|
||
str(img_part["image_url"]["url"]).endswith(os.path.join(TID1, "in.png")),
|
||
str(img_part))
|
||
file_part = msgs[0].content[2]
|
||
check("file part 绝对路径", file_part["type"] == "file" and
|
||
str(file_part["path"]).endswith(os.path.join(TID1, "data.pdf")), str(file_part))
|
||
a0 = msgs[1]
|
||
check("AIMessage0 合并两条并行调用", len(a0.tool_calls) == 2, str(a0.tool_calls))
|
||
check("AIMessage0 反思→content", a0.content == "反思A", repr(a0.content))
|
||
check("AIMessage0 思考→reasoning_content",
|
||
(a0.additional_kwargs or {}).get("reasoning_content") == "思考A",
|
||
str(a0.additional_kwargs))
|
||
check("AIMessage0 id", a0.id == f"hist_assist_{TID1}_0", str(a0.id))
|
||
check("Tool c1 匹配 id/name", msgs[2].tool_call_id == "c1" and msgs[2].name == "read_image")
|
||
check("Tool c2 含工具产物提示", "out.txt" in msgs[3].content, msgs[3].content)
|
||
check("AIMessage1 单调用", len(msgs[4].tool_calls) == 1 and msgs[4].id == f"hist_assist_{TID1}_1")
|
||
check("Tail 去重+注入", msgs[6].id == f"hist_tail_{TID1}_0" and
|
||
(msgs[6].additional_kwargs or {}).get("reasoning_content") == "尾部思考",
|
||
str(msgs[6].id))
|
||
check("Final id/内容", msgs[7].id == f"hist_final_{TID1}" and msgs[7].content == "处理完成")
|
||
check("全部 synthetic/tool_ctx 标记", all((m.additional_kwargs or {}).get("tool_ctx") for m in msgs))
|
||
|
||
# ---------- T2 ----------
|
||
section("T2 build_native_history_messages 存储路径 + 游标跳过")
|
||
turn_path = Path(sess_dir) / f"turn_{TID1}.json"
|
||
turn_path.write_text(json.dumps({"messages": turn_msgs}, ensure_ascii=False), encoding="utf-8")
|
||
tool_path = Path(sess_dir) / f"tool_{TID1}.json"
|
||
tool_path.write_text(json.dumps({
|
||
"turn_id": TID1,
|
||
"tool_calls": calls,
|
||
"reflections": tail,
|
||
}, ensure_ascii=False), encoding="utf-8")
|
||
hist0 = au.build_native_history_messages(SHARE_SID, 0)
|
||
check("cursor=0 重建 8 条", len(hist0) == 8, str(len(hist0)))
|
||
hist4 = au.build_native_history_messages(SHARE_SID, 4)
|
||
check("cursor=4 跳过前 4 条", len(hist4) == 4 and hist4[0].id == f"hist_assist_{TID1}_1",
|
||
str([m.id for m in hist4]))
|
||
check("cursor=8 无剩余", len(au.build_native_history_messages(SHARE_SID, 8)) == 0)
|
||
|
||
# ---------- T3 ----------
|
||
section("T3 compressed_context 读写往返（真实文件）")
|
||
ctx1 = {"version": 2, "cursor_msg": 0, "blocks": [
|
||
{"seq": 1, "content": "块1摘要", "msg_from": 1, "msg_to": 8, "units": 8,
|
||
"merged_units": 7, "turns": 1, "tokens_before": 900, "tokens_after": 60,
|
||
"compressed_at": 1},
|
||
]}
|
||
save_compressed_context(SHARE_SID, ctx1)
|
||
loaded = load_compressed_context(SHARE_SID)
|
||
check("往返一致", loaded == ctx1, str(loaded)[:200])
|
||
check("新键结构", "blocks" in loaded and loaded["version"] == 2, str(loaded)[:120])
|
||
|
||
# ---------- T4 ----------
|
||
section("T4 maybe_compress_context 单次压缩（范围/统计/RemoveMessage/块链）")
|
||
seen = {"live": [], "result": []}
|
||
tcr.get_current_session = lambda: SHARE_SID
|
||
tcr.is_sub_agent_context = lambda: False
|
||
tcr.record_tool_call_live = lambda sid, name, args: seen["live"].append((sid, name, args))
|
||
tcr.record_tool_result_live = lambda sid, name, text: seen["result"].append((sid, name, text))
|
||
au.summarize_chat_text = lambda llm, text: "SUM:" + text[:80]
|
||
|
||
state_msgs = [SystemMessage(content="sys")]
|
||
state_msgs.append(SystemMessage(content="【历史对话摘要 #1】\n块1摘要", id="compressed_block_1"))
|
||
state_msgs.extend(au.build_native_history_messages(SHARE_SID, 0))
|
||
state_msgs.append(HumanMessage(content="当前轮用户问题")) # 无 id → 不压缩
|
||
out1 = nodes.maybe_compress_context({"messages": state_msgs, "_need_compress": True})
|
||
upd = out1["messages"]
|
||
rm = [u for u in upd if type(u).__name__ == "RemoveMessage"]
|
||
new_sys = [u for u in upd if isinstance(u, SystemMessage)]
|
||
check("RemoveMessage 移除全部 8 条 hist_*", len(rm) == 8, str(len(rm)))
|
||
check("新增 1 个块 SystemMessage", len(new_sys) == 1)
|
||
check("新块 id=compressed_block_2", new_sys[0].id == "compressed_block_2", str(new_sys[0].id))
|
||
check("新块内容带序号前缀", new_sys[0].content.startswith("【历史对话摘要 #2】\nSUM:"),
|
||
str(new_sys[0].content)[:40])
|
||
check("_need_compress 复位", out1.get("_need_compress") is False)
|
||
ctx2 = load_compressed_context(SHARE_SID)
|
||
check("块链追加（2 块，新块 seq=2）", len(ctx2["blocks"]) == 2 and ctx2["blocks"][1]["seq"] == 2,
|
||
str(len(ctx2["blocks"])))
|
||
b2 = ctx2["blocks"][1]
|
||
check("统计 msg 区间 1-8", b2["msg_from"] == 1 and b2["msg_to"] == 8, str(b2))
|
||
check("统计 turns=1 / units=8 / merged=7", b2["turns"] == 1 and b2["units"] == 8 and b2["merged_units"] == 7,
|
||
f"turns={b2['turns']} units={b2['units']} merged={b2['merged_units']}")
|
||
check("统计 tokens_before>tokens_after>0",
|
||
b2["tokens_before"] > b2["tokens_after"] > 0,
|
||
f"{b2.get('tokens_before')}->{b2.get('tokens_after')}")
|
||
check("游标推进到 8", ctx2["cursor_msg"] == 8, str(ctx2["cursor_msg"]))
|
||
check("SSE 压缩开始事件", seen["live"] and seen["live"][-1][0] == SHARE_SID and seen["live"][-1][2]["action"] == "上下文压缩中",
|
||
str(seen["live"])[:120])
|
||
check("SSE 压缩完成事件", seen["result"] and "压缩完成" in seen["result"][-1][2],
|
||
str(seen["result"])[:120])
|
||
|
||
# ---------- T5 ----------
|
||
section("T5 双次压缩（seq 递增、旧块保留、游标推进）")
|
||
TID2 = "21090102_090000"
|
||
turn2_msgs = [
|
||
{"role": "user", "user": {"content": "第二轮问题"}},
|
||
{"role": "assistant", "assistant": {"text": "第二轮答复"}},
|
||
]
|
||
(Path(sess_dir) / f"turn_{TID2}.json").write_text(
|
||
json.dumps({"messages": turn2_msgs}, ensure_ascii=False), encoding="utf-8")
|
||
calls2 = [
|
||
{"id": "c4", "name": "search", "args": {"q": "xy"}, "status": "done",
|
||
"thinking": "思考B", "reflection": "", "result_text": "结果B",
|
||
"group_index": 0, "result_images": [], "result_files": []},
|
||
]
|
||
(Path(sess_dir) / f"tool_{TID2}.json").write_text(
|
||
json.dumps({"turn_id": TID2, "tool_calls": calls2, "reflections": []},
|
||
ensure_ascii=False), encoding="utf-8")
|
||
hist_t2 = au.build_native_history_messages(SHARE_SID, 8)
|
||
check("游标后新内容 4 条（user+AI+tool+final）", len(hist_t2) == 4, str([m.id for m in hist_t2]))
|
||
au.summarize_chat_text = lambda llm, text: "SUM2:" + text[:60]
|
||
state2_msgs = [SystemMessage(content="sys")]
|
||
for b in load_compressed_context(SHARE_SID)["blocks"]:
|
||
state2_msgs.append(SystemMessage(content=f"#{b['seq']}:{b['content']}", id=f"compressed_block_{b['seq']}"))
|
||
state2_msgs.extend(hist_t2)
|
||
state2_msgs.append(HumanMessage(content="第二轮当前问题"))
|
||
out2 = nodes.maybe_compress_context({"messages": state2_msgs, "_need_compress": True})
|
||
upd2 = out2["messages"]
|
||
rm2 = [u for u in upd2 if type(u).__name__ == "RemoveMessage"]
|
||
new_sys2 = [u for u in upd2 if isinstance(u, SystemMessage)]
|
||
check("第二次 RemoveMessage 4 条", len(rm2) == 4, str(len(rm2)))
|
||
check("第二次新块 seq=3", new_sys2[0].id == "compressed_block_3", str(new_sys2[0].id))
|
||
ctx3 = load_compressed_context(SHARE_SID)
|
||
check("块链 3 块且旧块内容未变",
|
||
len(ctx3["blocks"]) == 3 and ctx3["blocks"][0]["content"] == "块1摘要" and
|
||
ctx3["blocks"][1]["content"].startswith("SUM:"), str(len(ctx3["blocks"])))
|
||
b3 = ctx3["blocks"][2]
|
||
check("第 2 次统计区间 9-12", b3["msg_from"] == 9 and b3["msg_to"] == 12, str(b3))
|
||
check("游标推进到 12", ctx3["cursor_msg"] == 12, str(ctx3["cursor_msg"]))
|
||
|
||
# ---------- T6 ----------
|
||
section("T6 llm payload reasoning_content 透传")
|
||
from agent.core.llm import Multimodel_LLM
|
||
llm = Multimodel_LLM(model="test-model", api_key="sk-test", base_url="http://localhost:9/v1")
|
||
payload = llm._get_request_payload([
|
||
HumanMessage(content="用户问题"),
|
||
AIMessage(content="反思A", additional_kwargs={"reasoning_content": "思考A"},
|
||
tool_calls=[{"id": "c1", "name": "t", "args": {}}]),
|
||
ToolMessage(content="结果", tool_call_id="c1", name="t"),
|
||
])
|
||
pdicts = payload["messages"]
|
||
check("payload 3 条消息", len(pdicts) == 3, str(len(pdicts)))
|
||
asst = [p for p in pdicts if p.get("role") == "assistant"][0]
|
||
check("reasoning_content 透传", asst.get("reasoning_content") == "思考A", str(asst))
|
||
check("tool_calls 保留", len(asst.get("tool_calls") or []) == 1)
|
||
|
||
print(f"\n全部通过: {len(PASSED)} 项")
|
||
finally:
|
||
shutil.rmtree(tmp_root, ignore_errors=True)
|