154 lines
6.5 KiB
Python
154 lines
6.5 KiB
Python
#!/usr/bin/env python3
|
||
"""抽取智能体派发桥离线测试（阶段 F 第二部分）。
|
||
|
||
固定合同：任务包装配（正文注入/白名单/修复提示词）、单轮派发产出合法抽取、
|
||
机械校验失败触发一轮修复重派、两轮失败走保守收口、派发失败失败关闭。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import json
|
||
import pathlib
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from unittest import mock
|
||
|
||
PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[3]
|
||
SCRIPT_DIR = PROJECT_ROOT / ".agent" / "skills" / "extract-chapter-knowledge" / "scripts"
|
||
for path in (SCRIPT_DIR,):
|
||
if str(path) not in sys.path:
|
||
sys.path.insert(0, str(path))
|
||
|
||
import dispatch_extraction_bridge as bridge # noqa: E402
|
||
from dispatch_extraction_bridge import ( # noqa: E402
|
||
DispatchExtractionError,
|
||
build_extraction_task_spec,
|
||
extractor_session_paths,
|
||
run_extraction_via_dispatch,
|
||
)
|
||
|
||
BODY = "第三章 波纹。茧从实验舰外舱门的裂口挤出去，金属边缘像被剥开的肋骨。林深盯着深渊方向。"
|
||
|
||
|
||
def _valid_payload(evidence="茧从实验舰外舱门的裂口挤出去"):
|
||
return {
|
||
"entities": [
|
||
{"type": "character", "name": "林深", "brief": "试机师", "fields": {}, "evidence": "林深盯着深渊方向"},
|
||
{"type": "item", "name": "茧", "brief": "生物机甲", "fields": {}, "evidence": evidence},
|
||
],
|
||
"relations": [
|
||
{"source": "林深", "target": "茧", "type": "神经链接", "description": "链接",
|
||
"evidence": "茧从实验舰外舱门的裂口挤出去"},
|
||
],
|
||
"state": {"currentSituation": "出击", "characterStates": {}, "foreshadowing": {"埋": [], "推": [], "收": []}, "handoff": ""},
|
||
}
|
||
|
||
|
||
def _load_chapter_fake(work_id, chapter_order):
|
||
return ("深渊机神", 301, "波纹", 901, BODY), ["林深", "茧"]
|
||
|
||
|
||
class ExtractionSpecTest(unittest.TestCase):
|
||
|
||
def test_spec_shape(self):
|
||
spec = build_extraction_task_spec(
|
||
work_id=12, chapter_order=3, title="深渊机神", chapter_title="波纹",
|
||
existing_names=["林深"], body=BODY,
|
||
)
|
||
self.assertEqual(spec["role"], "extractor")
|
||
self.assertEqual(spec["input"]["body"], BODY)
|
||
self.assertEqual(spec["input"]["existingEntityNames"], ["林深"])
|
||
self.assertEqual(spec["toolAllowlist"], ["read_chapter_text", "search_entities"])
|
||
self.assertIn("逐字连续片段", spec["taskPrompt"])
|
||
self.assertNotIn("机械校验失败", spec["taskPrompt"])
|
||
|
||
def test_spec_repair_appends_reason(self):
|
||
spec = build_extraction_task_spec(
|
||
work_id=12, chapter_order=3, title="深渊机神", chapter_title="波纹",
|
||
existing_names=[], body=BODY, repair_reason="entities[0] 证据不在正文中",
|
||
)
|
||
self.assertIn("机械校验失败", spec["taskPrompt"])
|
||
self.assertIn("不得新增条目", spec["taskPrompt"])
|
||
|
||
def test_session_paths(self):
|
||
sid, path = extractor_session_paths(12, 3)
|
||
self.assertEqual(sid, "extractor-work12-ch3")
|
||
self.assertIn("extractor-sessions", str(path))
|
||
|
||
|
||
class OrchestrationTest(unittest.TestCase):
|
||
|
||
def setUp(self):
|
||
self._tmp = tempfile.TemporaryDirectory()
|
||
self.tmp = pathlib.Path(self._tmp.name)
|
||
self.dispatches = []
|
||
|
||
def tearDown(self):
|
||
self._tmp.cleanup()
|
||
|
||
def _fake_dispatch(self, outputs):
|
||
"""outputs: 每次派发依次返回的 JSON（Exception 表示派发失败）。"""
|
||
outputs = list(outputs)
|
||
|
||
def fake(spec_file, **kwargs):
|
||
self.dispatches.append(kwargs["run_id"])
|
||
run_dir = self.tmp / kwargs["run_id"]
|
||
run_dir.mkdir(parents=True)
|
||
out = outputs[len(self.dispatches) - 1]
|
||
if isinstance(out, Exception):
|
||
return {"status": "failed", "errorCode": "TIMEOUT"}, 1
|
||
(run_dir / "output.json").write_text(json.dumps(out, ensure_ascii=False), encoding="utf-8")
|
||
return {"status": "completed", "runDir": str(run_dir),
|
||
"actualModelIds": ["catproxy-anthropic/claude-opus-5"],
|
||
"usage": {"input_tokens": 10}}, 0
|
||
|
||
return fake
|
||
|
||
@contextlib.contextmanager
|
||
def _patched(self, outputs):
|
||
fake = self._fake_dispatch(outputs)
|
||
with mock.patch.object(bridge, "run_dispatch", fake), \
|
||
mock.patch.object(bridge, "_load_chapter", _load_chapter_fake), \
|
||
mock.patch.object(bridge, "start_run", lambda **kw: {"run_id": kw["run_id"]}), \
|
||
mock.patch.object(bridge, "finish_run", lambda *a, **kw: None), \
|
||
mock.patch.object(bridge, "record_failure", lambda *a, **kw: None), \
|
||
mock.patch.object(bridge, "_propose_chapter_extract_lesson", lambda **kw: None), \
|
||
mock.patch.object(bridge, "persist_extraction",
|
||
lambda *a, **kw: {"draft_ids": [1, 2, 3], "state_draft_id": 9}):
|
||
yield
|
||
|
||
def _run(self, outputs):
|
||
with self._patched(outputs):
|
||
return run_extraction_via_dispatch(
|
||
12, 3, repo_root=self.tmp, provider="catproxy-anthropic",
|
||
model="claude-opus-5", run_id="run-extract-test", spec_dir=self.tmp,
|
||
)
|
||
|
||
def test_single_dispatch_success(self):
|
||
summary = self._run([_valid_payload()])
|
||
self.assertEqual(summary["run_id"], "run-extract-test")
|
||
self.assertEqual(len(summary["draft_ids"]), 3)
|
||
self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1"])
|
||
|
||
def test_repair_round_when_evidence_invalid(self):
|
||
bad = _valid_payload(evidence="不在正文里的证据")
|
||
summary = self._run([bad, _valid_payload()])
|
||
self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1", "run-extract-test-extractor-v2"])
|
||
self.assertEqual(len(summary["draft_ids"]), 3)
|
||
|
||
def test_salvage_after_two_failures(self):
|
||
bad = _valid_payload(evidence="不在正文里的证据")
|
||
summary = self._run([bad, bad])
|
||
# salvage：证据无效但实体名逐字在正文 → 保留并以实体名为最小证据；关系缺证据丢弃
|
||
self.assertEqual(len(summary["draft_ids"]), 3)
|
||
self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1", "run-extract-test-extractor-v2"])
|
||
|
||
def test_dispatch_failure_fails_closed(self):
|
||
with self.assertRaises(DispatchExtractionError):
|
||
self._run([TimeoutError()])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|