474 lines
17 KiB
Python
474 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""卡索引驱动的冻结原文检索测试。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import copy
|
||
import pathlib
|
||
import sys
|
||
import unittest
|
||
|
||
PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[3]
|
||
ASSEMBLE_CONTEXT_SCRIPTS = PROJECT_ROOT / "muse" / "lifecycle" / "context" / "skills" / "assemble-context" / "scripts"
|
||
FREEZE_CONTEXT_SCRIPTS = PROJECT_ROOT / "muse" / "lifecycle" / "context" / "skills" / "freeze-context" / "scripts"
|
||
SEARCH_KNOWLEDGE_SCRIPTS = PROJECT_ROOT / "muse" / "lifecycle" / "context" / "skills" / "search-knowledge" / "scripts"
|
||
for script_dir in (
|
||
ASSEMBLE_CONTEXT_SCRIPTS,
|
||
FREEZE_CONTEXT_SCRIPTS,
|
||
SEARCH_KNOWLEDGE_SCRIPTS,
|
||
):
|
||
sys.path.insert(0, str(script_dir))
|
||
|
||
from load_reference_work import begin_read_snapshot # noqa: E402
|
||
from search import search_cards # noqa: E402
|
||
from retrieve_writer_sources import ( # noqa: E402
|
||
ProductionCardIndexRepository,
|
||
ReplayCardIndexRepository,
|
||
RetrievalError,
|
||
build_retrieval_plan,
|
||
freeze_card,
|
||
retrieve_writer_sources,
|
||
stable_sort_cards,
|
||
)
|
||
|
||
|
||
SOURCE_VERSION = "raw-file-v1:sha256:" + "a" * 64
|
||
|
||
|
||
def source_ref(chapter: int, block_id: int, offset: int = 0) -> dict:
|
||
"""构造可回读的冻结原文引用。"""
|
||
|
||
return {
|
||
"sourceId": f"chapter:{chapter}:block:{block_id}",
|
||
"sourceVersion": SOURCE_VERSION,
|
||
"chapter": chapter,
|
||
"blockId": block_id,
|
||
"startCodePoint": offset,
|
||
"endCodePoint": offset + 4,
|
||
}
|
||
|
||
|
||
def card(card_id: str, score: float, chapter: int = 3, offset: int = 0) -> dict:
|
||
"""构造含历史里程碑和原文指针的抽取卡索引。"""
|
||
|
||
return {
|
||
"cardId": card_id,
|
||
"type": "character",
|
||
"name": f"角色{card_id}",
|
||
"score": score,
|
||
"sourceVersion": SOURCE_VERSION,
|
||
"sourceId": f"canonical-entity:{card_id}",
|
||
"sourceOffset": offset,
|
||
"sourceRefs": [source_ref(chapter, int(card_id), offset)],
|
||
"milestones": [
|
||
{"chapter": 1, "fact": "初始状态"},
|
||
{"chapter": chapter, "fact": "冻结线内状态"},
|
||
{"chapter": 9, "fact": "未来终态"},
|
||
],
|
||
"sourceKind": "canonical_entity",
|
||
"sourceStatus": "active",
|
||
"bindingStatus": "active",
|
||
"productionRetrievalEligible": True,
|
||
}
|
||
|
||
|
||
class FakeCardRepository:
|
||
"""只返回预置卡片并记录固定计划查询。"""
|
||
|
||
def __init__(self, cards):
|
||
self.cards = cards
|
||
self.plans = []
|
||
|
||
def search(self, plan):
|
||
self.plans.append(copy.deepcopy(plan))
|
||
return copy.deepcopy(self.cards)
|
||
|
||
|
||
class FakeProseRepository:
|
||
"""按来源引用返回冻结片段，不访问数据库。"""
|
||
|
||
def __init__(self):
|
||
self.calls = []
|
||
|
||
def read_source_refs(self, *, work_id, as_of, source_refs):
|
||
self.calls.append((work_id, as_of, copy.deepcopy(source_refs)))
|
||
rows = []
|
||
for ref in source_refs:
|
||
rows.append(
|
||
{
|
||
"sourceRef": copy.deepcopy(ref),
|
||
"chapter": ref["chapter"],
|
||
"text": "历史原文",
|
||
"purpose": "card_source",
|
||
}
|
||
)
|
||
return rows
|
||
|
||
|
||
class FakeConnection:
|
||
"""记录事务声明，证明读取器先锁定只读可重复读快照。"""
|
||
|
||
def __init__(self):
|
||
self.statements = []
|
||
|
||
def execute(self, statement, params=None):
|
||
self.statements.append((" ".join(statement.split()), params))
|
||
return self
|
||
|
||
|
||
class FakeQueryResult:
|
||
"""模拟 psycopg 查询结果。"""
|
||
|
||
def __init__(self, rows):
|
||
self.rows = rows
|
||
|
||
def fetchall(self):
|
||
return self.rows
|
||
|
||
|
||
class FakeSearchConnection:
|
||
"""为 search_cards 提供字段策略与单条 Canonical 卡结果。"""
|
||
|
||
def __init__(self):
|
||
self.statements = []
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, traceback):
|
||
return False
|
||
|
||
def execute(self, statement, params=None):
|
||
normalized = " ".join(statement.split())
|
||
self.statements.append((normalized, params))
|
||
if "muse_meta_schema" in normalized:
|
||
return FakeQueryResult([("character", {"秘密": False})])
|
||
return FakeQueryResult(
|
||
[
|
||
(
|
||
"entity",
|
||
11,
|
||
{
|
||
"型": "character",
|
||
"名称": "甲",
|
||
"一句话摘要": "历史角色",
|
||
"字段": {
|
||
"秘密": "不可见",
|
||
"sourceRefs": [source_ref(3, 11)],
|
||
"milestones": [{"chapter": 3, "fact": "历史事实"}],
|
||
},
|
||
},
|
||
"active",
|
||
0.9,
|
||
2,
|
||
{"sourceRefs": [source_ref(3, 11)]},
|
||
"active",
|
||
"active",
|
||
)
|
||
]
|
||
)
|
||
|
||
|
||
class RetrieveWriterSourcesTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.plan = build_retrieval_plan(
|
||
run_id="run-a",
|
||
work_id=8,
|
||
target_chapter=6,
|
||
as_of=5,
|
||
fine_outline={
|
||
"entities": [
|
||
{"id": "character:甲", "type": "character", "name": "甲"},
|
||
{"id": "item:剑", "type": "item", "name": "剑"},
|
||
],
|
||
"relations": [],
|
||
"locations": [],
|
||
"powerSystems": [],
|
||
"hardConstraints": ["甲使用剑"],
|
||
},
|
||
card_index_version="cards-v1",
|
||
prose_index_version="prose-v1",
|
||
token_budget={"maxContextChars": 20000},
|
||
)
|
||
|
||
def test_cards_use_documented_stable_order_and_ties_repeat(self):
|
||
cards = [card("3", 0.8, offset=2), card("2", 0.8, offset=1), card("1", 0.9)]
|
||
expected = ["1", "2", "3"]
|
||
for _ in range(3):
|
||
self.assertEqual([item["cardId"] for item in stable_sort_cards(cards)], expected)
|
||
|
||
def test_plan_allows_asof_zero_for_first_chapter(self):
|
||
"""开篇基线章：asOf=0 是合法冻结线，计划与过滤器都接受。"""
|
||
|
||
plan = build_retrieval_plan(
|
||
run_id="run-ch1-baseline",
|
||
work_id=8,
|
||
target_chapter=1,
|
||
as_of=0,
|
||
fine_outline={
|
||
"entities": [],
|
||
"relations": [],
|
||
"locations": [],
|
||
"powerSystems": [],
|
||
"hardConstraints": ["开篇建立基调"],
|
||
},
|
||
card_index_version="cards-v1",
|
||
prose_index_version="prose-v1",
|
||
token_budget={"maxContextChars": 20000},
|
||
)
|
||
self.assertEqual(plan["asOf"], 0)
|
||
self.assertEqual(plan["filters"]["asOfChapter"], 0)
|
||
|
||
def test_plan_identity_excludes_run_id(self):
|
||
other = copy.deepcopy(self.plan)
|
||
other["runId"] = "run-b"
|
||
rebuilt = build_retrieval_plan(
|
||
run_id="run-b",
|
||
work_id=8,
|
||
target_chapter=6,
|
||
as_of=5,
|
||
fine_outline={
|
||
"entities": [
|
||
{"id": "character:甲", "type": "character", "name": "甲"},
|
||
{"id": "item:剑", "type": "item", "name": "剑"},
|
||
],
|
||
"relations": [],
|
||
"locations": [],
|
||
"powerSystems": [],
|
||
"hardConstraints": ["甲使用剑"],
|
||
},
|
||
card_index_version="cards-v1",
|
||
prose_index_version="prose-v1",
|
||
token_budget={"maxContextChars": 20000},
|
||
)
|
||
self.assertEqual(self.plan["planId"], rebuilt["planId"])
|
||
|
||
def test_state_as_of_uses_only_milestones_at_or_before_freeze(self):
|
||
frozen = freeze_card(card("1", 1.0), as_of=5)
|
||
self.assertEqual([item["chapter"] for item in frozen["stateAsOf"]], [1, 3])
|
||
self.assertNotIn("未来终态", str(frozen))
|
||
|
||
def test_target_future_and_unprovable_sources_fail_closed(self):
|
||
for bad_ref in (
|
||
source_ref(6, 1),
|
||
source_ref(7, 1),
|
||
{"sourceId": "unknown", "sourceVersion": SOURCE_VERSION},
|
||
):
|
||
unsafe = card("1", 1.0)
|
||
unsafe["sourceRefs"] = [bad_ref]
|
||
with self.subTest(bad_ref=bad_ref), self.assertRaises(RetrievalError):
|
||
retrieve_writer_sources(
|
||
plan=self.plan,
|
||
card_repository=FakeCardRepository([unsafe]),
|
||
prose_repository=FakeProseRepository(),
|
||
)
|
||
|
||
def test_cards_expand_to_prose_and_authoritative_facts_keep_independent_refs(self):
|
||
prose = FakeProseRepository()
|
||
no_ref_card = card("2", 0.7)
|
||
no_ref_card["sourceRefs"] = []
|
||
result = retrieve_writer_sources(
|
||
plan=self.plan,
|
||
card_repository=FakeCardRepository([card("1", 0.9), no_ref_card]),
|
||
prose_repository=prose,
|
||
authoritative_facts=[
|
||
{
|
||
"factId": "setting:1",
|
||
"fact": "正式设定事实",
|
||
"sourceType": "formal_setting",
|
||
"sourceRef": {
|
||
"sourceId": "setting:8",
|
||
"sourceVersion": "setting-v2",
|
||
},
|
||
"riskLevel": "high",
|
||
},
|
||
{
|
||
"factId": "outline:new:1",
|
||
"fact": "细纲声明的新事实",
|
||
"sourceType": "fine_outline_declared_new",
|
||
"sourceRef": {
|
||
"sourceId": "fine-outline:6",
|
||
"sourceVersion": "outline-v3",
|
||
"chapter": 6,
|
||
},
|
||
"riskLevel": "medium",
|
||
},
|
||
],
|
||
)
|
||
self.assertEqual(len(prose.calls[0][2]), 1)
|
||
self.assertEqual(result["proseEvidence"][0]["text"], "历史原文")
|
||
self.assertEqual(
|
||
{item["sourceType"] for item in result["factEvidence"]},
|
||
{"historical_prose", "formal_setting", "fine_outline_declared_new"},
|
||
)
|
||
self.assertEqual([item["cardId"] for item in result["indexHints"]], ["1", "2"])
|
||
self.assertEqual(
|
||
result["indexHints"][0],
|
||
{
|
||
"cardId": "1",
|
||
"name": "角色1",
|
||
"type": "character",
|
||
"content": "冻结线内状态",
|
||
"sourceId": "canonical-entity:1",
|
||
"sourceVersion": SOURCE_VERSION,
|
||
"asOf": 5,
|
||
},
|
||
)
|
||
self.assertEqual(result["unverifiedIndexHints"][0]["cardId"], "2")
|
||
self.assertEqual(
|
||
result["manifest"]["omittedSources"],
|
||
[{"sourceId": "canonical-entity:2", "reason": "missing_source_refs"}],
|
||
)
|
||
|
||
def test_eval_draft_hint_only_contains_index_metadata_and_proxy_is_explicit(self):
|
||
"""评测卡不透传里程碑台阶，整章代理必须保留降级来源和用途。"""
|
||
|
||
replay_card = card("1", 0.9)
|
||
replay_card["sourceKind"] = "eval_draft"
|
||
replay_card["sourceRefs"][0]["sourceType"] = "card_chapter_proxy"
|
||
result = retrieve_writer_sources(
|
||
plan=self.plan,
|
||
card_repository=FakeCardRepository([replay_card]),
|
||
prose_repository=FakeProseRepository(),
|
||
)
|
||
hint_content = result["indexHints"][0]["content"]
|
||
self.assertEqual(
|
||
hint_content,
|
||
'{"name":"角色1","sourceChapters":[3],"type":"character"}',
|
||
)
|
||
self.assertNotIn("冻结线内状态", hint_content)
|
||
self.assertEqual(
|
||
result["proseEvidence"][0]["sourceRef"]["sourceType"],
|
||
"card_chapter_proxy",
|
||
)
|
||
self.assertEqual(result["proseEvidence"][0]["purpose"], "card_chapter_proxy")
|
||
|
||
def test_snapshot_transaction_is_repeatable_read_and_read_only(self):
|
||
connection = FakeConnection()
|
||
begin_read_snapshot(connection)
|
||
self.assertEqual(
|
||
connection.statements[0][0],
|
||
"SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY",
|
||
)
|
||
|
||
def test_production_repository_only_accepts_active_canonical_binding(self):
|
||
calls = []
|
||
|
||
def search_function(**kwargs):
|
||
calls.append(kwargs)
|
||
return [card("1", 0.9)]
|
||
|
||
repository = ProductionCardIndexRepository(search_function=search_function)
|
||
result = repository.search(self.plan)
|
||
self.assertEqual(result[0]["cardId"], "1")
|
||
self.assertEqual(calls[0]["scope"], "work")
|
||
|
||
for field, value in (
|
||
("sourceKind", "eval_draft"),
|
||
("sourceStatus", "draft"),
|
||
("bindingStatus", "inactive"),
|
||
("productionRetrievalEligible", False),
|
||
):
|
||
unsafe = card("1", 0.9)
|
||
unsafe[field] = value
|
||
repository = ProductionCardIndexRepository(search_function=lambda **_: [unsafe])
|
||
with self.subTest(field=field), self.assertRaises(RetrievalError):
|
||
repository.search(self.plan)
|
||
|
||
def test_search_cards_reuses_active_entity_and_binding_sql(self):
|
||
connection = FakeSearchConnection()
|
||
result = search_cards(
|
||
"甲的历史状态",
|
||
scope="work",
|
||
work_id=8,
|
||
ttype="character",
|
||
purpose="generation",
|
||
top=5,
|
||
connection_factory=lambda _: connection,
|
||
embedder=lambda _: [0.1, 0.2],
|
||
)
|
||
sql = connection.statements[1][0]
|
||
self.assertIn("en.status='active'", sql)
|
||
self.assertIn("b.binding_status='active'", sql)
|
||
self.assertIn("en.source_action_policy='allowed'", sql)
|
||
self.assertEqual(result[0]["sourceKind"], "canonical_entity")
|
||
self.assertTrue(result[0]["productionRetrievalEligible"])
|
||
self.assertNotIn("秘密", result[0]["visibleFields"])
|
||
|
||
def test_replay_repository_is_preregistered_upgrade_book_and_never_production_eligible(self):
|
||
replay_card = card("11", 0.9)
|
||
replay_card.update(
|
||
{
|
||
"sourceKind": "eval_draft",
|
||
"evaluationStatus": "eval_draft",
|
||
"sourceType": "upgrade_book",
|
||
"productionRetrievalEligible": False,
|
||
}
|
||
)
|
||
replay_config = {
|
||
"targetChapter": 6,
|
||
"snapshot": {"asOfChapter": 5, "data": {"chapters": [{"chapter": 5, "text": "安全历史"}]}},
|
||
"sources": [
|
||
{"sourceId": "chapter:5", "sourceVersion": "chapter-v1", "chapter": 5}
|
||
],
|
||
"authorization": {
|
||
"sourceStatus": "active",
|
||
"copyrightStatus": "research_only",
|
||
"sourceHash": "sha256:" + "a" * 64,
|
||
"sourceVersion": SOURCE_VERSION,
|
||
"allowedPurpose": ["offline_evaluation"],
|
||
"forbiddenPurpose": ["external_distribution"],
|
||
"authorizationSnapshot": {
|
||
"id": "auth-1",
|
||
"version": "v1",
|
||
"immutable": True,
|
||
"sourceHash": "sha256:" + "a" * 64,
|
||
"sourceVersion": SOURCE_VERSION,
|
||
"sourceStatus": "active",
|
||
"copyrightStatus": "research_only",
|
||
"authorizationBasis": "user_authorization",
|
||
"allowedPurpose": ["offline_evaluation"],
|
||
"forbiddenPurpose": ["external_distribution"],
|
||
"checkedAt": "2026-07-20T00:00:00Z",
|
||
"revalidationAt": "2099-07-21T00:00:00Z",
|
||
},
|
||
},
|
||
"leakageAudit": {
|
||
"targetFacts": {"targetChapter": 6, "forbiddenFacts": []}
|
||
},
|
||
}
|
||
repository = ReplayCardIndexRepository.from_replay_config(
|
||
replay_config,
|
||
cards=[replay_card],
|
||
preregistered_card_ids=["11"],
|
||
)
|
||
self.assertFalse(repository.search(self.plan)[0]["productionRetrievalEligible"])
|
||
|
||
denied = copy.deepcopy(replay_config)
|
||
denied["authorization"] = {}
|
||
with self.assertRaises(RetrievalError):
|
||
ReplayCardIndexRepository.from_replay_config(
|
||
denied,
|
||
cards=[replay_card],
|
||
preregistered_card_ids=["11"],
|
||
)
|
||
|
||
wrong = copy.deepcopy(replay_card)
|
||
wrong["sourceType"] = "extract_chapter"
|
||
with self.assertRaises(RetrievalError):
|
||
ReplayCardIndexRepository.from_replay_config(
|
||
replay_config,
|
||
cards=[wrong],
|
||
preregistered_card_ids=["11"],
|
||
)
|
||
with self.assertRaises(RetrievalError):
|
||
ReplayCardIndexRepository.from_replay_config(
|
||
replay_config,
|
||
cards=[replay_card],
|
||
preregistered_card_ids=["12"],
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|