401 lines
17 KiB
Python
401 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
端到端写入路径测试。
|
||
|
||
把 judge / preproc / extractor / reconciler / store 串起来，
|
||
验证 CodingWriter 协调正确。
|
||
"""
|
||
|
||
import json
|
||
|
||
import pytest
|
||
|
||
from hy_memory.coding.extractor import CodingMemoryExtractor
|
||
from hy_memory.coding.reconciler import CodingMemoryReconciler
|
||
from hy_memory.coding.writer import CodingWriter, is_coding_segment
|
||
from hy_memory.pipelines.base import ChatMessage, ToolCall
|
||
|
||
from tests.coding._helpers import MockLLMProvider, FailingLLMProvider
|
||
|
||
|
||
def _classify_resp(is_coding):
|
||
return json.dumps({"is_coding": is_coding})
|
||
|
||
|
||
def _extract_resp(experiences, fact=None):
|
||
if isinstance(fact, list):
|
||
if len(fact) <= 1:
|
||
fact = fact[0] if fact else None
|
||
else:
|
||
# 仅供兼容性回归：旧模型可能仍返回多元素 facts 数组。
|
||
return json.dumps({"experiences": experiences, "facts": fact})
|
||
return json.dumps({"experiences": experiences, "fact": fact})
|
||
|
||
|
||
def _reconcile_resp(ops):
|
||
return json.dumps({"ops": ops})
|
||
|
||
|
||
def _coding_messages():
|
||
return [
|
||
ChatMessage(role="user", content="帮我打包发布到私有包索引"),
|
||
ChatMessage(role="assistant", content="", tool_calls=[
|
||
ToolCall(id="1", name="Read", arguments={"path": "scripts/publish_package.sh"}),
|
||
]),
|
||
ChatMessage(role="tool", content="200 lines", tool_call_id="1", tool_name="Read"),
|
||
ChatMessage(role="assistant", content="", tool_calls=[
|
||
ToolCall(id="2", name="Bash", arguments={"command": "twine upload"}),
|
||
]),
|
||
ChatMessage(role="tool", content="OK uploaded", tool_call_id="2", tool_name="Bash"),
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
class TestIsCodingSegment:
|
||
async def test_no_tool_short_circuit(self):
|
||
msgs = [ChatMessage(role="user", content="hi"),
|
||
ChatMessage(role="assistant", content="hello")]
|
||
# 即使 LLM 返回 true，has_any_tool_message False 应短路 False
|
||
llm = MockLLMProvider('{"is_coding": true}')
|
||
result = await is_coding_segment(msgs, llm)
|
||
assert result is False
|
||
# LLM 不应被调用
|
||
assert len(llm.calls) == 0
|
||
|
||
async def test_with_tool_classify_coding(self):
|
||
llm = MockLLMProvider(_classify_resp(True))
|
||
result = await is_coding_segment(_coding_messages(), llm)
|
||
assert result is True
|
||
|
||
async def test_with_tool_classify_chat(self):
|
||
llm = MockLLMProvider(_classify_resp(False))
|
||
result = await is_coding_segment(_coding_messages(), llm)
|
||
assert result is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
class TestCodingWriterE2E:
|
||
async def test_trace_history_shared_across_working_and_coding_variants(
|
||
self, coding_config, mock_embed,
|
||
):
|
||
from hy_memory.coding.store import CodingMemoryStore
|
||
|
||
working = CodingMemoryStore(
|
||
coding_config, embed_service=mock_embed, variant="working",
|
||
)
|
||
coding = CodingMemoryStore(
|
||
coding_config, embed_service=mock_embed, variant="coding",
|
||
)
|
||
await working.initialize()
|
||
await coding.initialize()
|
||
try:
|
||
await working.append_trace(
|
||
[ChatMessage(role="user", content="先整理财务报表")],
|
||
request_id="shared-trace", user_id="u1", agent_id="a1",
|
||
)
|
||
recent = await coding.get_recent_trace_messages(
|
||
user_id="u1", agent_id="a1", max_user_turns=5,
|
||
)
|
||
assert [m.content for m in recent if m.role == "user"] == [
|
||
"先整理财务报表"
|
||
]
|
||
assert working._db_path != coding._db_path
|
||
assert working._trace_db_path == coding._trace_db_path
|
||
finally:
|
||
await coding.close()
|
||
await working.close()
|
||
|
||
async def test_trace_history_keeps_last_five_user_turns(self, coding_store):
|
||
msgs = []
|
||
for i in range(7):
|
||
msgs.extend([
|
||
ChatMessage(role="user", content=f"history-{i}"),
|
||
ChatMessage(role="assistant", content="A" * 8000),
|
||
ChatMessage(role="tool", content="T" * 5000, tool_name="Read"),
|
||
])
|
||
await coding_store.append_trace(
|
||
msgs, request_id="trace-1", user_id="u1", agent_id="a1",
|
||
)
|
||
recent = await coding_store.get_recent_trace_messages(
|
||
user_id="u1", agent_id="a1", max_user_turns=5,
|
||
)
|
||
assert [m.content for m in recent if m.role == "user"] == [
|
||
"history-2", "history-3", "history-4", "history-5", "history-6"
|
||
]
|
||
assert all("chars omitted" in m.content for m in recent if m.role == "assistant")
|
||
assert all("bytes omitted" in m.content for m in recent if m.role == "tool")
|
||
|
||
async def test_trace_history_retention_is_bounded(self, coding_store):
|
||
for i in range(105):
|
||
await coding_store.append_trace(
|
||
[ChatMessage(role="user", content=f"turn-{i}")],
|
||
request_id=f"trace-{i}", user_id="u1", agent_id="a1",
|
||
)
|
||
with coding_store._trace_lock:
|
||
count = coding_store._trace_conn.execute(
|
||
"SELECT COUNT(*) AS n FROM coding_trace WHERE user_id=? AND agent_id=?",
|
||
("u1", "a1"),
|
||
).fetchone()["n"]
|
||
assert count == 100
|
||
|
||
async def test_full_path_add(self, coding_store, mock_embed):
|
||
"""完整 e2e：含 tool 消息 → judge=True → extract → reconcile=ADD → 持久化"""
|
||
# 用同一个 LLM mock 序列响应 extract + reconcile（注意：我们这里直接调
|
||
# writer.write，不经 judge；所以 LLM 序列只需覆盖 extract + reconcile）
|
||
ext_resp = _extract_resp([{
|
||
"task": "打包发布到私有包索引",
|
||
"search_keys": ["凭证位置", "上传命令"],
|
||
"solution": "运行 scripts/publish_package.sh X.Y.Z；凭证通过环境变量提供",
|
||
"boundary_envs": "twine>=6\\nurl: packages.example.com",
|
||
"confidence": 0.9,
|
||
}])
|
||
rec_resp = _reconcile_resp([{"draft_idx": 0, "action": "ADD"}])
|
||
llm = MockLLMProvider([ext_resp, rec_resp])
|
||
|
||
ext = CodingMemoryExtractor(llm)
|
||
rec = CodingMemoryReconciler(coding_store, mock_embed, llm)
|
||
writer = CodingWriter(
|
||
store=coding_store, extractor=ext, reconciler=rec,
|
||
llm_provider=llm, embed_service=mock_embed,
|
||
)
|
||
resp = await writer.write(
|
||
_coding_messages(),
|
||
user_id="u1", agent_id="a1",
|
||
workspace_id="github.com/x/y", branch="main",
|
||
session_id="s1",
|
||
)
|
||
assert resp["success"] is True
|
||
assert resp["scene"] == "productivity"
|
||
assert len(resp["ops"]) == 1
|
||
assert resp["ops"][0]["action"] == "ADD"
|
||
assert len(resp["experience_ids"]) == 1
|
||
|
||
# 验证持久化
|
||
experience_id = resp["experience_ids"][0]
|
||
got = await coding_store.get_by_id(experience_id, user_id="u1")
|
||
assert got is not None
|
||
assert got.task == "打包发布到私有包索引"
|
||
assert got.search_keys == ["凭证位置", "上传命令"]
|
||
assert got.boundary_envs.startswith("twine>=6")
|
||
assert got.workspace_id == "github.com/x/y"
|
||
assert got.files == ["scripts/publish_package.sh"]
|
||
|
||
async def test_value_bar_returns_empty(self, coding_store, mock_embed):
|
||
"""extractor 返回 0 条 memory（价值门槛未达）"""
|
||
llm = MockLLMProvider(_extract_resp([]))
|
||
ext = CodingMemoryExtractor(llm)
|
||
rec = CodingMemoryReconciler(coding_store, mock_embed, llm)
|
||
writer = CodingWriter(
|
||
store=coding_store, extractor=ext, reconciler=rec,
|
||
llm_provider=llm, embed_service=mock_embed,
|
||
)
|
||
resp = await writer.write(
|
||
_coding_messages(), user_id="u1", workspace_id="ws/a",
|
||
)
|
||
assert resp["success"] is True
|
||
assert resp["ops"] == []
|
||
assert resp["experience_ids"] == []
|
||
|
||
async def test_no_workspace_still_writes(self, coding_store, mock_embed):
|
||
"""缺 workspace_id 不再拦稿，两条 draft 都应落库"""
|
||
ext_resp = _extract_resp([
|
||
{"task": "T1", "solution": "S", "confidence": 0.9},
|
||
{"task": "T2", "solution": "S", "confidence": 0.9},
|
||
])
|
||
rec_resp = _reconcile_resp([
|
||
{"draft_idx": 0, "action": "ADD"},
|
||
{"draft_idx": 1, "action": "ADD"},
|
||
])
|
||
llm = MockLLMProvider([ext_resp, rec_resp])
|
||
ext = CodingMemoryExtractor(llm)
|
||
rec = CodingMemoryReconciler(coding_store, mock_embed, llm)
|
||
writer = CodingWriter(
|
||
store=coding_store, extractor=ext, reconciler=rec,
|
||
llm_provider=llm, embed_service=mock_embed,
|
||
)
|
||
resp = await writer.write(
|
||
_coding_messages(), user_id="u1", workspace_id=None,
|
||
)
|
||
assert len(resp["experience_ids"]) == 2
|
||
assert sorted(await coding_store.list_user_tasks("u1")) == ["T1", "T2"]
|
||
|
||
async def test_extractor_failure(self, coding_store, mock_embed):
|
||
"""extractor LLM 失败 → 空 drafts → 空 ops"""
|
||
ext = CodingMemoryExtractor(FailingLLMProvider())
|
||
rec = CodingMemoryReconciler(coding_store, mock_embed, MockLLMProvider("{}"))
|
||
writer = CodingWriter(
|
||
store=coding_store, extractor=ext, reconciler=rec,
|
||
llm_provider=FailingLLMProvider(), embed_service=mock_embed,
|
||
)
|
||
resp = await writer.write(_coding_messages(), user_id="u1", workspace_id="ws")
|
||
assert resp["success"] is True
|
||
assert resp["ops"] == []
|
||
|
||
async def test_existing_tasks_not_shown_to_extractor(self, coding_store, mock_embed):
|
||
"""库里已有的 task 不得进抽取 prompt——抽取只读这一段，去重是 reconcile 的事"""
|
||
from hy_memory.coding.types import CodingMemory
|
||
await coding_store.insert(CodingMemory(
|
||
experience_id="e_old", user_id="u1", task="老 task", solution="S",
|
||
))
|
||
|
||
llm = MockLLMProvider([_extract_resp([]), _reconcile_resp([])])
|
||
ext = CodingMemoryExtractor(llm)
|
||
rec = CodingMemoryReconciler(coding_store, mock_embed, llm)
|
||
writer = CodingWriter(
|
||
store=coding_store, extractor=ext, reconciler=rec,
|
||
llm_provider=llm, embed_service=mock_embed,
|
||
)
|
||
await writer.write(_coding_messages(), user_id="u1", workspace_id="ws/a")
|
||
assert "老 task" not in llm.calls[0]["prompt"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
class TestDualLayerWrite:
|
||
"""双层写入 + 同批全互链（§4）"""
|
||
|
||
def _writer(self, store, embed, llm, *, verbatim_fallback_enabled=False):
|
||
return CodingWriter(
|
||
store=store,
|
||
extractor=CodingMemoryExtractor(
|
||
llm,
|
||
verbatim_fallback_enabled=verbatim_fallback_enabled,
|
||
),
|
||
reconciler=CodingMemoryReconciler(store, embed, llm),
|
||
llm_provider=llm, embed_service=embed,
|
||
)
|
||
|
||
async def test_both_layers_written_and_crosslinked(self, coding_store, mock_embed):
|
||
llm = MockLLMProvider([
|
||
_extract_resp(
|
||
[{"task": "打包发布到私有包索引", "solution": "S", "confidence": 0.9}],
|
||
[{"statement": "发布脚本在 scripts/publish_package.sh", "files": ["scripts/publish_package.sh"]}],
|
||
),
|
||
_reconcile_resp([{"draft_idx": 0, "action": "ADD"}]),
|
||
])
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(
|
||
_coding_messages(), user_id="u1", workspace_id="ws/a", session_id="s1",
|
||
)
|
||
|
||
assert len(resp["experience_ids"]) == 1
|
||
assert len(resp["fact_ids"]) == 1
|
||
assert resp["fact_ops_summary"] == {"insert": 1, "supersede": 0, "total": 1}
|
||
|
||
eid, fid = resp["experience_ids"][0], resp["fact_ids"][0]
|
||
exp = await coding_store.get_by_id(eid, user_id="u1")
|
||
fact = await coding_store.get_fact(fid, user_id="u1")
|
||
assert exp.source_fact_ids == [fid]
|
||
assert fact.derived_experience_ids == [eid]
|
||
assert fact.statement == "发布脚本在 scripts/publish_package.sh"
|
||
|
||
async def test_all_experiences_linked_to_unique_fact(self, coding_store, mock_embed):
|
||
"""一个 trace 的多条经验都与唯一会话 fact 双向关联。"""
|
||
llm = MockLLMProvider([
|
||
_extract_resp(
|
||
[{"task": "T1", "solution": "S"}, {"task": "T2", "solution": "S"}],
|
||
[{"statement": "事实 A"}, {"statement": "事实 B"}],
|
||
),
|
||
_reconcile_resp([
|
||
{"draft_idx": 0, "action": "ADD"},
|
||
{"draft_idx": 1, "action": "ADD"},
|
||
]),
|
||
])
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
eids, fids = resp["experience_ids"], resp["fact_ids"]
|
||
assert len(eids) == 2 and len(fids) == 1
|
||
|
||
for eid in eids:
|
||
exp = await coding_store.get_by_id(eid, user_id="u1")
|
||
assert exp.source_fact_ids == fids
|
||
fact = await coding_store.get_fact(fids[0], user_id="u1")
|
||
assert sorted(fact.derived_experience_ids) == sorted(eids)
|
||
|
||
async def test_skip_not_linked(self, coding_store, mock_embed):
|
||
"""SKIP 的经验没真入库，不该被连上这批事实。"""
|
||
llm = MockLLMProvider([
|
||
_extract_resp(
|
||
[{"task": "T1", "solution": "S"}],
|
||
[{"statement": "事实 A"}],
|
||
),
|
||
_reconcile_resp([{"draft_idx": 0, "action": "SKIP", "reason": "dup"}]),
|
||
])
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
assert resp["experience_ids"] == []
|
||
assert len(resp["fact_ids"]) == 1
|
||
fact = await coding_store.get_fact(resp["fact_ids"][0], user_id="u1")
|
||
assert fact.derived_experience_ids == []
|
||
|
||
async def test_update_keeps_old_fact_link(self, coding_store, mock_embed):
|
||
"""UPDATE 覆盖正文不留痕，但事实链取并集，旧链要保住。"""
|
||
from hy_memory.coding.types import CodingMemory
|
||
await coding_store.insert(CodingMemory(
|
||
experience_id="e1", user_id="u1", task="老 task", solution="旧 solution",
|
||
source_fact_ids=["f_before"],
|
||
))
|
||
llm = MockLLMProvider([
|
||
_extract_resp(
|
||
[{"task": "老 task", "solution": "新 solution"}],
|
||
[{"statement": "本轮新事实"}],
|
||
),
|
||
_reconcile_resp([
|
||
{"draft_idx": 0, "action": "UPDATE", "target_experience_id": "e1"},
|
||
]),
|
||
])
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
exp = await coding_store.get_by_id("e1", user_id="u1")
|
||
assert exp.solution == "新 solution"
|
||
assert exp.source_fact_ids == ["f_before", resp["fact_ids"][0]]
|
||
|
||
async def test_facts_only_no_experience(self, coding_store, mock_embed):
|
||
"""经验层空、事实层非空：事实照样落库，不触发经验层 reconcile。"""
|
||
llm = MockLLMProvider(_extract_resp([], [{"statement": "只有事实"}]))
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
assert resp["experience_ids"] == []
|
||
assert len(resp["fact_ids"]) == 1
|
||
# extract 一次；事实层库里没候选所以不调 LLM；经验层 reconcile 完全没跑
|
||
assert len(llm.calls) == 1
|
||
|
||
async def test_empty_experiences_do_not_persist_verbatim_fact_by_default(
|
||
self, coding_store, mock_embed,
|
||
):
|
||
llm = MockLLMProvider(_extract_resp([], []))
|
||
writer = self._writer(coding_store, mock_embed, llm)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
assert resp["ops"] == []
|
||
assert resp["fact_ids"] == []
|
||
assert resp["fact_ops"] == []
|
||
|
||
async def test_empty_experiences_can_persist_verbatim_fact_when_enabled(
|
||
self, coding_store, mock_embed,
|
||
):
|
||
"""显式开启时，两层都抽空仍可落一条原文事实。
|
||
|
||
端到端确认开关开启后，fallback 能穿过 reconcile 真正落库。
|
||
"""
|
||
llm = MockLLMProvider(_extract_resp([], []))
|
||
writer = self._writer(
|
||
coding_store,
|
||
mock_embed,
|
||
llm,
|
||
verbatim_fallback_enabled=True,
|
||
)
|
||
|
||
resp = await writer.write(_coding_messages(), user_id="u1")
|
||
assert resp["ops"] == [] # 经验层不兜底
|
||
assert len(resp["fact_ids"]) == 1
|
||
assert [op["action"] for op in resp["fact_ops"]] == ["INSERT"]
|
||
|
||
fact = await coding_store.get_fact(resp["fact_ids"][0], user_id="u1")
|
||
assert fact.statement == "帮我打包发布到私有包索引"
|
||
assert "twine upload" in fact.details # 轨迹原文进了 details
|