test: add 3GPP agent unit tests
Covers prompt assembly and retrieval wiring so refactors fail loudly.pull/893/head
parent
4d960838e4
commit
cd6efc9206
|
|
@ -0,0 +1,241 @@
|
|||
"""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()
|
||||
Loading…
Reference in New Issue