242 lines
11 KiB
Python
242 lines
11 KiB
Python
"""Regression tests for evidence retrieval, refusals, citations and streaming failures."""
|
||
|
||
import json
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
import zipfile
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from xml.etree import ElementTree as ET
|
||
|
||
from agent_core import LocalModelRuntime, ProposalAgent, SpecRetrieverTool, SpecValidatorTool
|
||
from rag_evidence import ROOT, Evidence, WORD_NS, extract_tdoc_ids
|
||
|
||
|
||
class FixtureTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.temp = tempfile.TemporaryDirectory(dir=ROOT / "output")
|
||
self.addCleanup(self.temp.cleanup)
|
||
self.root = Path(self.temp.name)
|
||
self.docs = self.root / "docs"
|
||
self.meeting = self.docs / "meeting"
|
||
self.meeting.mkdir(parents=True)
|
||
self.index = self.root / "index.json"
|
||
self.index.write_text(json.dumps({
|
||
"R2-2400001": {"meeting": "meeting", "title": "XR rate control", "source": "A"},
|
||
"R2-2400002": {"meeting": "meeting", "title": "Delay-aware LCP", "source": "B"},
|
||
"R2-2400003": {"meeting": "meeting", "title": "Mobility", "source": "C"},
|
||
"R2-2400004": {"meeting": "meeting", "title": "Metadata only", "source": "D"},
|
||
}), encoding="utf-8")
|
||
self.write_docx("R2-2400001", ["Header " + "x" * 1800, "Proposal 1: recommended bit rate signaling over a DRB."])
|
||
self.write_docx("R2-2400002", ["Delay-aware LCP uses remaining packet delay information."])
|
||
(self.meeting / "R2-2400003.doc").write_bytes(b"unsupported legacy document")
|
||
self.retriever = SpecRetrieverTool(tdoc_cache_path=self.index, docs_root=self.docs)
|
||
|
||
def write_docx(self, number, paragraphs):
|
||
document = ET.Element(WORD_NS + "document")
|
||
body = ET.SubElement(document, WORD_NS + "body")
|
||
for text in paragraphs:
|
||
paragraph = ET.SubElement(body, WORD_NS + "p")
|
||
run = ET.SubElement(paragraph, WORD_NS + "r")
|
||
ET.SubElement(run, WORD_NS + "t").text = text
|
||
with zipfile.ZipFile(self.meeting / (number + ".docx"), "w") as archive:
|
||
archive.writestr("word/document.xml", ET.tostring(document))
|
||
|
||
def test_multiple_ids_survive_top_k_one(self):
|
||
result = self.retriever.retrieve_evidence("compare R2-2400001 and R2-2400002", top_k=1)
|
||
self.assertTrue(result.answerable)
|
||
self.assertEqual({item.tdoc for item in result.evidence}, {"R2-2400001", "R2-2400002"})
|
||
|
||
def test_proposal_after_long_header_is_retrieved(self):
|
||
result = self.retriever.retrieve_evidence("R2-2400001 recommended bit rate")
|
||
passage = next(item for item in result.evidence if "Proposal 1" in item.text)
|
||
self.assertLessEqual(passage.paragraph_start, 2)
|
||
self.assertGreaterEqual(passage.paragraph_end, 2)
|
||
self.assertLessEqual(len(passage.text), 900)
|
||
|
||
def test_unknown_id_does_not_substitute_training_answer(self):
|
||
result = self.retriever.retrieve_evidence("R2-2999999 XR rate control")
|
||
self.assertFalse(result.answerable)
|
||
self.assertEqual(result.evidence, [])
|
||
self.assertEqual(result.documents[0]["status"], "missing")
|
||
|
||
def test_partial_comparison_is_not_answerable(self):
|
||
result = self.retriever.retrieve_evidence("R2-2400001 R2-2999999")
|
||
self.assertTrue(result.evidence)
|
||
self.assertFalse(result.answerable)
|
||
|
||
def test_legacy_doc_is_not_body_evidence(self):
|
||
result = self.retriever.retrieve_evidence("R2-2400003")
|
||
self.assertFalse(result.answerable)
|
||
self.assertEqual(result.documents[0]["status"], "unsupported_format")
|
||
|
||
def test_metadata_only_is_not_body_evidence(self):
|
||
result = self.retriever.retrieve_evidence("R2-2400004")
|
||
self.assertFalse(result.answerable)
|
||
self.assertEqual(result.documents[0]["status"], "metadata_only")
|
||
|
||
def test_bm25_search_finds_title(self):
|
||
result = self.retriever.retrieve_evidence("速率控制", top_k=1)
|
||
self.assertEqual(result.documents[0]["tdoc"], "R2-2400001")
|
||
|
||
def test_no_lexical_match_returns_no_evidence(self):
|
||
self.assertEqual(self.retriever.retrieve_evidence("zzzxqmissing").evidence, [])
|
||
|
||
def test_corrupt_docx_is_reported(self):
|
||
(self.meeting / "R2-2400002.docx").write_bytes(b"invalid zip")
|
||
result = self.retriever.retrieve_evidence("R2-2400002")
|
||
self.assertFalse(result.answerable)
|
||
self.assertEqual(result.documents[0]["status"], "read_error")
|
||
|
||
def test_changed_docx_refreshes_paragraph_cache(self):
|
||
self.retriever.retrieve_evidence("R2-2400002")
|
||
self.write_docx("R2-2400002", ["A newly revised paragraph about logical channels."])
|
||
result = self.retriever.retrieve_evidence("R2-2400002")
|
||
self.assertIn("newly revised", result.evidence[0].text)
|
||
|
||
def test_missing_index_has_visible_warning(self):
|
||
retriever = SpecRetrieverTool(tdoc_cache_path=self.root / "absent.json", docs_root=self.docs)
|
||
self.assertTrue(retriever.retrieve_evidence("XR").warnings)
|
||
|
||
def test_abstention_never_calls_runtime(self):
|
||
class MustNotRun:
|
||
def stream(self, *args):
|
||
raise AssertionError("No model should run without evidence")
|
||
agent = ProposalAgent(retriever=self.retriever, runtime=MustNotRun())
|
||
last = list(agent.run_stream("R2-2400003"))[-1]
|
||
self.assertEqual(last["mode"], "insufficient_evidence")
|
||
|
||
def test_disabled_retrieval_respected_even_with_id(self):
|
||
agent = ProposalAgent(retriever=self.retriever)
|
||
last = list(agent.run_stream("R2-2400001", enable_rag=False))[-1]
|
||
self.assertEqual(last["sources"]["evidence"], [])
|
||
self.assertIsNone(agent.runtime.model)
|
||
|
||
def test_retrieval_only_never_loads_model(self):
|
||
agent = ProposalAgent(retriever=self.retriever)
|
||
last = list(agent.run_stream("R2-2400001", retrieval_only=True))[-1]
|
||
self.assertEqual(last["mode"], "retrieval_only")
|
||
self.assertIsNone(agent.runtime.model)
|
||
|
||
def test_model_receives_both_sources_and_marks_extra_unverified(self):
|
||
class RecordingRuntime:
|
||
messages = None
|
||
def stream(self, messages, *args):
|
||
self.messages = messages
|
||
yield "速率控制见 [S1]。时延处理见 [S3]。"
|
||
runtime = RecordingRuntime()
|
||
agent = ProposalAgent(retriever=self.retriever, runtime=runtime)
|
||
last = list(agent.run_stream("R2-2400001 R2-2400002", extra_context="custom claim"))[-1]
|
||
self.assertEqual(last["mode"], "generated")
|
||
self.assertIn("R2-2400002", runtime.messages[1]["content"])
|
||
self.assertIn("未经核实", runtime.messages[1]["content"])
|
||
self.assertIn("first_text_seconds", last["metrics"])
|
||
|
||
def test_generation_error_is_not_reported_as_completion(self):
|
||
class BrokenRuntime:
|
||
def stream(self, *args):
|
||
yield "partial answer"
|
||
raise RuntimeError("GPU unavailable")
|
||
agent = ProposalAgent(retriever=self.retriever, runtime=BrokenRuntime())
|
||
events = list(agent.run_stream("R2-2400001"))
|
||
self.assertEqual(events[-1]["stage"], "ERROR")
|
||
self.assertFalse(any(event["stage"] == "4. COMPLETE" for event in events))
|
||
|
||
def test_empty_question_is_an_error(self):
|
||
self.assertEqual(list(ProposalAgent(retriever=self.retriever).run_stream(" "))[-1]["stage"], "ERROR")
|
||
|
||
|
||
class CitationTests(unittest.TestCase):
|
||
def test_normalize_and_deduplicate_ids(self):
|
||
self.assertEqual(extract_tdoc_ids("r2 – 2408133 R2-2408133 R2-2408134"), ["R2-2408133", "R2-2408134"])
|
||
|
||
def test_invented_citation_fails(self):
|
||
evidence = [Evidence("S1", "R2-2400001", "m", "title", "A", "file", 1, 2, "source")]
|
||
result = SpecValidatorTool().validate("A statement [S99]", evidence)
|
||
self.assertEqual(result["unknown_citations"], ["S99"])
|
||
self.assertEqual(result["factual_correctness"], "not_verified")
|
||
self.assertNotIn("compliance_score", result)
|
||
|
||
def test_known_citation_does_not_claim_factually_verified(self):
|
||
evidence = [Evidence("S1", "R2-2400001", "m", "title", "A", "file", 1, 2, "source")]
|
||
result = SpecValidatorTool().validate("A wrong statement [S1]", evidence)
|
||
self.assertEqual(result["citation_status"], "format_checked")
|
||
self.assertEqual(result["factual_correctness"], "not_verified")
|
||
|
||
def test_comparison_citing_only_one_document_fails(self):
|
||
evidence = [
|
||
Evidence("S1", "R2-2400001", "m", "title", "A", "file", 1, 2, "source"),
|
||
Evidence("S2", "R2-2400002", "m", "title", "B", "file", 1, 2, "source"),
|
||
]
|
||
result = SpecValidatorTool().validate("Both proposals do this [S1]", evidence)
|
||
self.assertEqual(result["citation_status"], "failed")
|
||
self.assertEqual(result["uncited_documents"], ["R2-2400002"])
|
||
|
||
def test_no_citations_is_not_a_fact_check(self):
|
||
result = SpecValidatorTool().validate("TS 38.331 RRC")
|
||
self.assertEqual(result["factual_correctness"], "not_verified")
|
||
self.assertNotIn("is_valid", result)
|
||
|
||
|
||
class RuntimeTests(unittest.TestCase):
|
||
@staticmethod
|
||
def runtime(behavior):
|
||
import torch
|
||
from transformers import BatchEncoding
|
||
|
||
class Tokenizer:
|
||
pad_token_id = 0
|
||
eos_token_id = 2
|
||
def apply_chat_template(self, *args, **kwargs):
|
||
return "prompt"
|
||
def __call__(self, *args, **kwargs):
|
||
return BatchEncoding({"input_ids": torch.tensor([[1, 3]]), "attention_mask": torch.tensor([[1, 1]])})
|
||
def decode(self, tokens, **kwargs):
|
||
return "".join("word " for _ in tokens)
|
||
|
||
class Model:
|
||
config = SimpleNamespace(max_position_embeddings=8192)
|
||
def get_input_embeddings(self):
|
||
return SimpleNamespace(weight=torch.zeros(1))
|
||
def generate(self, **kwargs):
|
||
kwargs["streamer"].put(kwargs["input_ids"])
|
||
if behavior == "error":
|
||
raise RuntimeError("intentional worker failure")
|
||
if behavior == "timeout":
|
||
time.sleep(0.15)
|
||
else:
|
||
kwargs["streamer"].put(torch.tensor([4]))
|
||
kwargs["streamer"].end()
|
||
|
||
runtime = LocalModelRuntime()
|
||
runtime.model, runtime.tokenizer = Model(), Tokenizer()
|
||
return runtime
|
||
|
||
def test_worker_error_reaches_caller_and_releases_lock(self):
|
||
runtime = self.runtime("error")
|
||
with self.assertRaisesRegex(RuntimeError, "intentional worker failure"):
|
||
list(runtime.stream([{"role": "user", "content": "q"}]))
|
||
self.assertFalse(runtime._lock.locked())
|
||
|
||
def test_stream_deadline_and_lock_cleanup(self):
|
||
runtime = self.runtime("timeout")
|
||
with self.assertRaises(TimeoutError):
|
||
list(runtime.stream([{"role": "user", "content": "q"}], timeout_seconds=0.03))
|
||
self.assertFalse(runtime._lock.locked())
|
||
|
||
def test_successful_stream_returns_text(self):
|
||
runtime = self.runtime("success")
|
||
self.assertIn("word", "".join(runtime.stream([{"role": "user", "content": "q"}])))
|
||
|
||
def test_busy_request_is_rejected(self):
|
||
runtime = self.runtime("success")
|
||
runtime._lock.acquire()
|
||
self.addCleanup(runtime._lock.release)
|
||
with self.assertRaisesRegex(RuntimeError, "另一条请求"):
|
||
list(runtime.stream([{"role": "user", "content": "q"}]))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|