205 lines
9.1 KiB
Python
205 lines
9.1 KiB
Python
"""§8.8 Phase 2 折叠逻辑单测:切点选择 / 内存-prepared 映射 / Session.apply_fold /
|
|
maybe_fold 编排(DB 与 LLM 全打桩,不碰真实 PG)。"""
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
from uuid import uuid4
|
|
|
|
import core.context_fold as cf
|
|
from core.context_fold import (
|
|
FOLD_TAIL_RATIO,
|
|
MIN_FOLD_MESSAGES,
|
|
SUMMARY_NOTE_HEADER,
|
|
_kth_user_index,
|
|
_user_ordinal,
|
|
build_summary_note,
|
|
find_cutoff,
|
|
)
|
|
from core.session import Session
|
|
|
|
|
|
def _mk_window(n_turns: int, content_chars: int = 200) -> list:
|
|
"""构造 n_turns 轮 user/assistant 对话(等长消息,方便按条数推算体量占比)。"""
|
|
msgs = []
|
|
for i in range(n_turns):
|
|
msgs.append({"role": "user", "content": f"u{i:03d}" + "x" * content_chars})
|
|
msgs.append({"role": "assistant", "content": f"a{i:03d}" + "y" * content_chars})
|
|
return msgs
|
|
|
|
|
|
class FindCutoffTests(unittest.TestCase):
|
|
def test_cutoff_is_user_message_and_tail_ratio_respected(self) -> None:
|
|
messages = [{"role": "system", "content": "rules"}] + _mk_window(30)
|
|
cutoff = find_cutoff(messages, head_len=1)
|
|
self.assertIsNotNone(cutoff)
|
|
self.assertEqual(messages[cutoff].get("role"), "user")
|
|
# 尾部(切点起)体量 >= FOLD_TAIL_RATIO(等长消息按条数近似)
|
|
tail = len(messages) - cutoff
|
|
window = len(messages) - 1
|
|
self.assertGreaterEqual(tail / window, FOLD_TAIL_RATIO * 0.9)
|
|
# 可折叠段不为空且够长
|
|
self.assertGreaterEqual(cutoff - 1, MIN_FOLD_MESSAGES)
|
|
|
|
def test_too_few_messages_returns_none(self) -> None:
|
|
messages = [{"role": "system", "content": "rules"}] + _mk_window(3)
|
|
self.assertIsNone(find_cutoff(messages, head_len=1))
|
|
|
|
def test_empty_window_returns_none(self) -> None:
|
|
self.assertIsNone(find_cutoff([{"role": "system", "content": "s"}], head_len=1))
|
|
|
|
def test_cutoff_never_splits_tool_pairing(self) -> None:
|
|
# assistant.tool_calls 后面跟 tool 结果 —— 切点是 user 消息,天然不落在中间
|
|
messages = [{"role": "system", "content": "rules"}]
|
|
for i in range(20):
|
|
messages.append({"role": "user", "content": f"u{i}" + "x" * 200})
|
|
messages.append({
|
|
"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": f"tc{i}", "function": {"name": "shell"}}],
|
|
})
|
|
messages.append({"role": "tool", "tool_call_id": f"tc{i}", "content": "r" * 200})
|
|
cutoff = find_cutoff(messages, head_len=1)
|
|
self.assertIsNotNone(cutoff)
|
|
self.assertEqual(messages[cutoff]["role"], "user")
|
|
|
|
|
|
class UserOrdinalMappingTests(unittest.TestCase):
|
|
def test_ordinal_roundtrip_across_prepared_list(self) -> None:
|
|
"""prepare 补占位 tool 消息后 index 平移,但「第 k 条 user」两边一致。"""
|
|
messages = [
|
|
{"role": "system", "content": "s"},
|
|
{"role": "user", "content": "q1"},
|
|
# 悬空 tool_call:prepare 的 repair 会在其后补一条占位 tool 消息
|
|
{"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": "tc1", "function": {"name": "shell"}}]},
|
|
{"role": "user", "content": "q2"},
|
|
{"role": "assistant", "content": "a2"},
|
|
]
|
|
from core.context import prepare_messages_with_stats
|
|
prepared, stats = prepare_messages_with_stats(messages)
|
|
self.assertEqual(stats["repaired_tool_calls"], 1)
|
|
k = _user_ordinal(messages, 3) # q2 是第 2 条 user
|
|
self.assertEqual(k, 2)
|
|
p_idx = _kth_user_index(prepared, k)
|
|
self.assertEqual(prepared[p_idx]["content"], "q2")
|
|
self.assertEqual(p_idx, 4) # 占位 tool 插入使 index +1
|
|
|
|
def test_kth_user_missing_returns_none(self) -> None:
|
|
self.assertIsNone(_kth_user_index([{"role": "assistant", "content": "a"}], 1))
|
|
|
|
|
|
class ApplyFoldTests(unittest.TestCase):
|
|
def test_apply_fold_rebuilds_head_and_advances_base(self) -> None:
|
|
sess = Session(task_id=uuid4(), system_prompt="SYS")
|
|
window = _mk_window(10)
|
|
sess.messages.extend(window) # 模拟 load:base=0,head=1(system)
|
|
cutoff = 11 # 内存 index,某条 user 消息
|
|
self.assertEqual(sess.messages[cutoff]["role"], "user")
|
|
|
|
sess.apply_fold(cutoff, "SUMMARY-TEXT")
|
|
|
|
# 窗口 = [system] + [摘要注入] + 尾段
|
|
self.assertEqual(sess.messages[0]["role"], "system")
|
|
self.assertTrue(sess.messages[1]["content"].startswith(SUMMARY_NOTE_HEADER))
|
|
self.assertIn("SUMMARY-TEXT", sess.messages[1]["content"])
|
|
self.assertEqual(sess.messages[2]["content"], window[cutoff - 1]["content"])
|
|
# base 推进 = 折叠条数(cutoff - 旧 head 1);head 变 2(system + 摘要)
|
|
self.assertEqual(sess.context_base, cutoff - 1)
|
|
self.assertEqual(sess.context_head_len, 2)
|
|
|
|
def test_second_fold_replaces_summary_note(self) -> None:
|
|
sess = Session(task_id=uuid4(), system_prompt="SYS")
|
|
sess.messages.extend(_mk_window(10))
|
|
sess.apply_fold(11, "FIRST")
|
|
base_after_first = sess.context_base
|
|
sess.messages.extend(_mk_window(10))
|
|
# head=2,再折 10 条(切到 index 12 的 user 消息)
|
|
self.assertEqual(sess.messages[12]["role"], "user")
|
|
sess.apply_fold(12, "SECOND")
|
|
notes = [m for m in sess.messages if str(m.get("content", "")).startswith(SUMMARY_NOTE_HEADER)]
|
|
self.assertEqual(len(notes), 1) # 旧摘要被折进新摘要,不叠加
|
|
self.assertIn("SECOND", notes[0]["content"])
|
|
self.assertEqual(sess.context_base, base_after_first + 10)
|
|
self.assertEqual(sess.context_head_len, 2)
|
|
|
|
def test_summary_note_not_counted_as_user_msg(self) -> None:
|
|
sess = Session(task_id=uuid4(), system_prompt="SYS")
|
|
sess.messages.extend(_mk_window(10))
|
|
before = sess.n_user_msgs()
|
|
sess.apply_fold(11, "S")
|
|
# 折掉 5 轮 user(index 1..10 里 5 条 user),注入的摘要不计入
|
|
self.assertEqual(sess.n_user_msgs(), before - 5)
|
|
|
|
|
|
class _FakeLLM:
|
|
def __init__(self, reply: str) -> None:
|
|
self.reply = reply
|
|
self.last_messages = None
|
|
|
|
def chat(self, messages, **kwargs):
|
|
self.last_messages = messages
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content=self.reply))],
|
|
usage=SimpleNamespace(prompt_tokens=100, completion_tokens=20),
|
|
)
|
|
|
|
|
|
_FAKE_CAPS = SimpleNamespace(
|
|
reliable_context=100, # 极小上限 → 体量必然超阈值,不用造几十万字符
|
|
family="deepseek_v4", variant="flash",
|
|
input_cny_per_mtoken=1.0, output_cny_per_mtoken=2.0,
|
|
)
|
|
|
|
|
|
class MaybeFoldTests(unittest.TestCase):
|
|
def _session(self) -> Session:
|
|
sess = Session(task_id=uuid4(), system_prompt="SYS")
|
|
sess.messages.extend(_mk_window(30))
|
|
return sess
|
|
|
|
def test_fold_happy_path(self) -> None:
|
|
sess = self._session()
|
|
llm = _FakeLLM("## 目标\n...")
|
|
events = []
|
|
with patch.object(cf, "persist_fold") as persist, \
|
|
patch.object(cf, "record_chat_usage") as usage:
|
|
stats = cf.maybe_fold(sess, llm, _FAKE_CAPS, user_id=uuid4(), emit=events.append)
|
|
self.assertIsNotNone(stats)
|
|
self.assertEqual(stats["phase"], "done")
|
|
# 摘要指令是折叠调用的最后一条消息(前缀复用构造)
|
|
self.assertEqual(llm.last_messages[-1]["content"], cf.FOLD_INSTRUCTION)
|
|
# DB 持久化先行,参数与内存变更一致
|
|
persist.assert_called_once_with(sess.task_id, stats["new_base"], "## 目标\n...")
|
|
usage.assert_called_once()
|
|
# 独立 kind:折叠次数 = count(usage_events kind='context_fold'),task 详情接口依赖
|
|
self.assertEqual(usage.call_args.kwargs.get("kind"), "context_fold")
|
|
self.assertEqual(sess.context_base, stats["new_base"])
|
|
phases = [e.get("phase") for e in events if e.get("type") == "context_fold"]
|
|
self.assertEqual(phases, ["start", "done"])
|
|
self.assertLess(stats["after_chars"], stats["before_chars"])
|
|
|
|
def test_below_threshold_no_fold(self) -> None:
|
|
sess = self._session()
|
|
caps = SimpleNamespace(**{**_FAKE_CAPS.__dict__, "reliable_context": 10_000_000})
|
|
llm = _FakeLLM("unused")
|
|
with patch.object(cf, "persist_fold") as persist:
|
|
self.assertIsNone(cf.maybe_fold(sess, llm, caps, user_id=uuid4(), emit=lambda e: None))
|
|
persist.assert_not_called()
|
|
self.assertIsNone(llm.last_messages)
|
|
|
|
def test_empty_summary_skips_fold(self) -> None:
|
|
sess = self._session()
|
|
before = list(sess.messages)
|
|
events = []
|
|
with patch.object(cf, "persist_fold") as persist:
|
|
result = cf.maybe_fold(
|
|
sess, _FakeLLM(" "), _FAKE_CAPS, user_id=uuid4(), emit=events.append
|
|
)
|
|
self.assertIsNone(result)
|
|
persist.assert_not_called()
|
|
self.assertEqual(sess.messages, before) # 内存零污染
|
|
self.assertTrue(any(e.get("type") == "warn" for e in events))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|