"""Tests for Lab 01 — authorized hybrid retrieval. pytest test_lab.py -v LAB_MODULE=solution pytest test_lab.py -v # the reference — must be green """ import importlib import os import pytest lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab")) POLICY_TEXT = """## Hold release A payment on sanctions hold may be released only after a screening review is completed and recorded. Releases above AED 100,000 require dual control by two authorised officers. ## Escalation If a possible match scores above 0.80, escalate to Group Compliance before any release. """ PAYMENT_TEXT = """## Payment PMT-771 Payment reference PMT-771 for AED 250,000 to beneficiary Acme Trading FZE is currently HELD. The hold was applied by the sanctions screening engine on value date 2026-03-12. ## Counterparty Acme Trading FZE is a UAE-registered entity, LEI 5493001KJTIIGC8Y1R12, ultimately controlled by Northgate Holdings Ltd. """ RETAIL_TEXT = """## Collections script When a retail customer disputes a late fee, apologise, check the statement date and offer a one-time goodwill reversal if the account has no prior reversals. """ MNPI_TEXT = """## Project Falcon Project Falcon is a confidential acquisition of Acme Trading FZE by a listed acquirer. Announcement is expected in Q3. This information is material and non-public. """ def doc(doc_id="d1", tenant="wholesale", text=POLICY_TEXT, **kwargs): kwargs.setdefault("title", doc_id) return lab.Document(doc_id=doc_id, tenant=tenant, text=text, **kwargs) # ====================================================================================== # 1. Documents and classification # ====================================================================================== def test_document_requires_identity(): with pytest.raises(ValueError): lab.Document("", "wholesale", "t", "text") with pytest.raises(ValueError): lab.Document("d", "", "t", "text") def test_document_rejects_an_unknown_classification_at_ingestion(): with pytest.raises(ValueError): doc(classification="cosmic") def test_classification_is_ordered(): assert lab.classification_rank("public") < lab.classification_rank("internal") assert lab.classification_rank("confidential") < lab.classification_rank("restricted") with pytest.raises(ValueError): lab.classification_rank("secret-ish") # ====================================================================================== # 2. Chunking # ====================================================================================== def test_chunking_splits_on_sections_first(): chunks = lab.chunk_document(doc(text=POLICY_TEXT), max_tokens=200) assert [c.section for c in chunks] == ["Hold release", "Escalation"] def test_chunking_falls_back_to_body_without_headings(): chunks = lab.chunk_document(doc(text="just a paragraph of text"), max_tokens=200) assert [c.section for c in chunks] == ["body"] def test_chunking_respects_the_token_budget(): chunks = lab.chunk_document(doc(text=PAYMENT_TEXT), max_tokens=20, overlap_tokens=4) assert chunks for chunk in chunks: assert lab.estimate_tokens(chunk.text) <= 20 def test_chunking_never_splits_mid_word(): chunks = lab.chunk_document(doc(text=PAYMENT_TEXT), max_tokens=12, overlap_tokens=2) words = set(PAYMENT_TEXT.split()) for chunk in chunks: for word in chunk.text.split(): assert word in words def test_overlap_must_be_smaller_than_the_chunk(): with pytest.raises(ValueError): lab.chunk_document(doc(), max_tokens=10, overlap_tokens=10) with pytest.raises(ValueError): lab.chunk_document(doc(), max_tokens=0) def test_chunks_carry_provenance(): chunk = lab.chunk_document(doc(doc_id="pol", text=POLICY_TEXT))[0] assert chunk.doc_id == "pol" assert chunk.tenant == "wholesale" assert chunk.doc_version == 1 assert "pol@v1#Hold release" in chunk.citation() def test_chunk_ids_are_unique_and_stable(): first = lab.chunk_document(doc(text=PAYMENT_TEXT), max_tokens=20) second = lab.chunk_document(doc(text=PAYMENT_TEXT), max_tokens=20) ids = [c.chunk_id for c in first] assert len(ids) == len(set(ids)) assert ids == [c.chunk_id for c in second] def test_empty_text_produces_no_chunks(): assert lab.chunk_document(doc(text="")) == [] def test_estimate_tokens(): assert lab.estimate_tokens("") == 0 assert lab.estimate_tokens("abcd") == 1 assert lab.estimate_tokens("abcde") == 2 # ====================================================================================== # 3. Embeddings # ====================================================================================== def test_embedding_is_deterministic_and_normalized(): a = lab.hash_embed("sanctions hold release") assert a == lab.hash_embed("sanctions hold release") assert lab.cosine(a, a) == pytest.approx(1.0, abs=1e-9) def test_empty_text_embeds_to_zero(): assert lab.cosine(lab.hash_embed(""), lab.hash_embed("anything")) == 0.0 def test_related_text_scores_above_unrelated(): query = lab.hash_embed("payment on sanctions hold") near = lab.hash_embed("the payment is on a sanctions hold") far = lab.hash_embed("goodwill reversal of a late fee") assert lab.cosine(query, near) > lab.cosine(query, far) def test_signed_hashing_allows_negative_similarity(): """Without random signs, collisions always add constructively and nothing is ever dissimilar. With them, unrelated text can score at or below zero.""" scores = [lab.cosine(lab.hash_embed(f"alpha{i}"), lab.hash_embed(f"beta{i}")) for i in range(40)] assert min(scores) <= 0.0 # ====================================================================================== # 4. BM25 # ====================================================================================== def build_bm25(*documents): index = lab.BM25Index() for document in documents: for chunk in lab.chunk_document(document): index.add(chunk) return index def test_bm25_finds_an_exact_identifier(): index = build_bm25(doc("pmt", text=PAYMENT_TEXT)) hits = index.search("PMT-771", namespace="wholesale") assert hits assert "PMT-771" in hits[0].chunk.text def test_bm25_returns_nothing_for_an_unknown_namespace(): index = build_bm25(doc("pmt", text=PAYMENT_TEXT)) assert index.search("PMT-771", namespace="retail") == [] def test_bm25_ignores_terms_absent_from_the_corpus(): index = build_bm25(doc("pmt", text=PAYMENT_TEXT)) assert index.search("quantum entanglement", namespace="wholesale") == [] def test_bm25_saturates_term_frequency(): """k1 controls saturation: going 9 -> 10 adds far less than going 1 -> 2. b=0 isolates the effect; with length normalisation on, the longer documents are also penalised and the two effects are not separable. """ index = lab.BM25Index(k1=1.5, b=0.0) for repeats in (1, 2, 9, 10): index.add(lab.Chunk(f"c{repeats:02d}", "d", "t", " ".join(["hold"] * repeats) + " payment", "s", 0, 10, "internal", None, 1, 0)) by_id = {h.chunk.chunk_id: h.score for h in index.search("hold", namespace="t")} assert by_id["c10"] > by_id["c09"] > by_id["c02"] > by_id["c01"] assert (by_id["c10"] - by_id["c09"]) < (by_id["c02"] - by_id["c01"]) def test_bm25_penalises_length_when_b_is_high(): short = lab.Chunk("short", "d", "t", "sanctions hold", "s", 0, 1, "internal", None, 1, 0) long = lab.Chunk("long", "d", "t", "sanctions hold " + "filler " * 40, "s", 0, 1, "internal", None, 1, 0) high = lab.BM25Index(b=1.0) low = lab.BM25Index(b=0.0) for index in (high, low): index.add(short) index.add(long) high_hits = {h.chunk.chunk_id: h.score for h in high.search("sanctions", namespace="t")} low_hits = {h.chunk.chunk_id: h.score for h in low.search("sanctions", namespace="t")} assert high_hits["short"] > high_hits["long"] assert low_hits["short"] == pytest.approx(low_hits["long"]) def test_bm25_idf_is_never_negative(): """A term present in every document must contribute 0, not a negative score.""" index = lab.BM25Index() for i in range(3): index.add(lab.Chunk(f"c{i}", "d", "t", "hold", "s", 0, 1, "internal", None, 1, 0)) for hit in index.search("hold", namespace="t"): assert hit.score >= 0.0 def test_bm25_rejects_bad_parameters(): with pytest.raises(ValueError): lab.BM25Index(k1=-1) with pytest.raises(ValueError): lab.BM25Index(b=1.5) def test_bm25_results_are_deterministic_under_ties(): index = lab.BM25Index() for name in ("zeta", "alpha", "mid"): index.add(lab.Chunk(name, "d", "t", "hold payment", "s", 0, 1, "internal", None, 1, 0)) ids = [h.chunk.chunk_id for h in index.search("hold", namespace="t")] assert ids == sorted(ids) # ====================================================================================== # 5. Vector index # ====================================================================================== def build_vectors(*documents): index = lab.VectorIndex() for document in documents: for chunk in lab.chunk_document(document): index.add(chunk) return index def test_vector_index_is_namespaced(): index = build_vectors(doc("pmt", "wholesale", PAYMENT_TEXT), doc("col", "retail", RETAIL_TEXT)) assert index.namespaces() == ["retail", "wholesale"] hits = index.search("goodwill reversal", namespace="wholesale") assert all(h.chunk.tenant == "wholesale" for h in hits) def test_vector_search_in_an_empty_namespace(): assert lab.VectorIndex().search("anything", namespace="nope") == [] def test_vector_results_are_ordered_by_similarity(): index = build_vectors(doc("pmt", text=PAYMENT_TEXT)) hits = index.search("Acme Trading FZE beneficiary", namespace="wholesale") scores = [h.score for h in hits] assert scores == sorted(scores, reverse=True) # ====================================================================================== # 6. Fusion and reranking # ====================================================================================== def scored(chunk_id, score=1.0): return lab.Scored(lab.Chunk(chunk_id, "d", "t", chunk_id, "s", 0, 1, "internal", None, 1, 0), score) def test_rrf_is_score_free(): """Wildly different score scales must not change the fused order.""" a = [scored("x", 1000.0), scored("y", 999.0)] b = [scored("x", 0.001), scored("y", 0.0009)] fused = lab.reciprocal_rank_fusion([a, b]) assert [f.chunk.chunk_id for f in fused] == ["x", "y"] def test_rrf_rewards_consistent_agreement(): """1st + 10th = 0.03068 ; 3rd + 3rd = 0.03175 — agreement wins.""" strong = [scored("strong")] + [scored(f"p{i}") for i in range(9)] other = [scored(f"q{i}") for i in range(9)] + [scored("strong")] consistent_a = [scored("a"), scored("b"), scored("consistent")] consistent_b = [scored("c"), scored("d"), scored("consistent")] one = lab.reciprocal_rank_fusion([strong, other]) two = lab.reciprocal_rank_fusion([consistent_a, consistent_b]) strong_score = next(s.score for s in one if s.chunk.chunk_id == "strong") consistent_score = next(s.score for s in two if s.chunk.chunk_id == "consistent") assert consistent_score > strong_score def test_rrf_is_order_independent_across_rankings(): a = [scored("x"), scored("y")] b = [scored("y"), scored("z")] assert ([s.chunk.chunk_id for s in lab.reciprocal_rank_fusion([a, b])] == [s.chunk.chunk_id for s in lab.reciprocal_rank_fusion([b, a])]) def test_rrf_handles_a_single_ranking_and_an_empty_one(): assert len(lab.reciprocal_rank_fusion([[scored("x")]])) == 1 assert lab.reciprocal_rank_fusion([]) == [] assert lab.reciprocal_rank_fusion([[]]) == [] def test_rrf_rejects_a_bad_k(): with pytest.raises(ValueError): lab.reciprocal_rank_fusion([[scored("x")]], k=0) def test_reranker_scores_the_pair(): match = lab.Chunk("m", "d", "t", "dual control is required for release", "s", 0, 1, "internal", None, 1, 0) other = lab.Chunk("o", "d", "t", "the weather in March was mild", "s", 0, 1, "internal", None, 1, 0) q = "dual control release" assert lab.lexical_overlap_reranker(q, match) > lab.lexical_overlap_reranker(q, other) def test_reranker_handles_an_empty_query(): chunk = lab.Chunk("c", "d", "t", "text", "s", 0, 1, "internal", None, 1, 0) assert lab.lexical_overlap_reranker("", chunk) == 0.0 def test_rerank_only_reorders_the_candidates_given(): candidates = [scored("a"), scored("b"), scored("c")] out = lab.rerank("a", candidates, reranker=lambda q, c: {"c": 3.0, "b": 2.0, "a": 1.0}[c.chunk_id], limit=2) assert [s.chunk.chunk_id for s in out] == ["c", "b"] def test_rerank_applies_a_relevance_floor(): """First-stage retrieval always returns SOMETHING — an ANN index has a nearest neighbour even for a query about nothing in the corpus. Without a floor, 'nothing is relevant' is indistinguishable from 'here are the least-irrelevant chunks'.""" candidates = [scored("a"), scored("b"), scored("c")] out = lab.rerank("a", candidates, reranker=lambda q, c: 1.0 if c.chunk_id == "c" else 0.0) assert [s.chunk.chunk_id for s in out] == ["c"] def test_rerank_floor_is_configurable(): candidates = [scored("a"), scored("b")] out = lab.rerank("a", candidates, reranker=lambda q, c: {"a": 0.4, "b": 0.9}[c.chunk_id], min_score=0.5) assert [s.chunk.chunk_id for s in out] == ["b"] # ====================================================================================== # 7. The authorized retriever # ====================================================================================== def build_retriever(): retriever = lab.AuthorizedRetriever(bm25=lab.BM25Index(), vectors=lab.VectorIndex()) retriever.ingest(doc("pol", "wholesale", POLICY_TEXT, updated_tick=100)) retriever.ingest(doc("pmt", "wholesale", PAYMENT_TEXT, classification="confidential", updated_tick=140)) retriever.ingest(doc("col", "retail", RETAIL_TEXT, updated_tick=90)) retriever.ingest(doc("deal", "wholesale", MNPI_TEXT, classification="restricted", barrier="advisory", updated_tick=150)) return retriever def principal(**kwargs): kwargs.setdefault("user_id", "u") kwargs.setdefault("tenant", "wholesale") kwargs.setdefault("max_classification", "confidential") return lab.Principal(**kwargs) def test_retrieval_returns_chunks_with_citations(): result = build_retriever().retrieve("dual control release", principal()) assert result.chunks assert all("@v1#" in c for c in result.citations) def test_a_cross_tenant_query_is_never_a_candidate(): """Not filtered out — never retrieved. `considered` counts only the namespace searched.""" retriever = build_retriever() result = retriever.retrieve("goodwill reversal late fee", principal()) assert all(s.chunk.tenant == "wholesale" for s in result.chunks) assert all(s.chunk.doc_id != "col" for s in result.chunks) def test_the_same_query_works_for_the_owning_tenant(): retriever = build_retriever() result = retriever.retrieve("goodwill reversal late fee", principal(tenant="retail")) assert [s.chunk.doc_id for s in result.chunks] == ["col"] def test_classification_is_enforced_before_ranking(): retriever = build_retriever() low = retriever.retrieve("PMT-771 held amount", principal(max_classification="internal")) assert all(s.chunk.classification != "confidential" for s in low.chunks) assert low.excluded_by_entitlement > 0 def test_a_barrier_hides_content_from_outside_it(): retriever = build_retriever() outside = retriever.retrieve("Project Falcon acquisition", principal(max_classification="restricted")) inside = retriever.retrieve("Project Falcon acquisition", principal(max_classification="restricted", barriers=("advisory",))) assert all(s.chunk.doc_id != "deal" for s in outside.chunks) assert any(s.chunk.doc_id == "deal" for s in inside.chunks) def test_visibility_is_a_conjunction(): retriever = build_retriever() chunk = lab.Chunk("c", "d", "wholesale", "text", "s", 0, 1, "restricted", "advisory", 1, 0) assert not retriever.visible(principal(max_classification="restricted"), chunk) assert not retriever.visible(principal(max_classification="internal", barriers=("advisory",)), chunk) assert not retriever.visible(principal(tenant="retail", max_classification="restricted", barriers=("advisory",)), chunk) assert retriever.visible(principal(max_classification="restricted", barriers=("advisory",)), chunk) def test_freshness_contract_drops_stale_chunks(): retriever = build_retriever() policy = lab.RetrievalPolicy(max_stale_ticks=30) fresh = retriever.retrieve("PMT-771 held", principal(), now_tick=160, policy=policy) stale = retriever.retrieve("PMT-771 held", principal(), now_tick=400, policy=policy) assert fresh.chunks assert stale.chunks == () assert stale.excluded_by_freshness > 0 def test_no_freshness_contract_means_no_exclusions(): retriever = build_retriever() result = retriever.retrieve("PMT-771 held", principal(), now_tick=10_000) assert result.excluded_by_freshness == 0 assert result.chunks def test_disabling_rerank_is_recorded_as_degradation(): retriever = build_retriever() result = retriever.retrieve("dual control", principal(), policy=lab.RetrievalPolicy(enable_rerank=False)) assert "rerank" in result.degraded assert result.chunks # degraded in QUALITY, still an answer def test_a_query_matching_nothing_returns_empty_not_an_error(): retriever = build_retriever() result = retriever.retrieve("zzzz nonexistent terminology", principal()) assert result.chunks == () def test_result_limit_is_respected(): retriever = build_retriever() result = retriever.retrieve("payment hold release control", principal(), policy=lab.RetrievalPolicy(result_limit=2)) assert len(result.chunks) <= 2 def test_retrieval_is_deterministic(): a = build_retriever().retrieve("dual control release", principal()) b = build_retriever().retrieve("dual control release", principal()) assert [s.chunk.chunk_id for s in a.chunks] == [s.chunk.chunk_id for s in b.chunks] # ====================================================================================== # 8. Grounding # ====================================================================================== def evidence(): return build_retriever().retrieve("PMT-771 hold reason and release rules", principal()).chunks def test_a_supported_claim_gets_a_citation(): report = lab.check_grounding( [lab.Claim("Releases above AED 100,000 require dual control by two authorised officers")], evidence()) assert report.is_grounded assert report.supported[0].citation is not None def test_an_invented_claim_is_reported_unsupported(): report = lab.check_grounding( [lab.Claim("The customer has been notified by email and telephone")], evidence()) assert not report.is_grounded assert report.unsupported[0].citation is None def test_coverage_is_the_supported_fraction(): report = lab.check_grounding([ lab.Claim("Releases above AED 100,000 require dual control by two authorised officers"), lab.Claim("The customer has been notified by email and telephone"), ], evidence()) assert report.coverage == pytest.approx(0.5) def test_a_claim_of_only_stopwords_is_trivially_supported(): report = lab.check_grounding([lab.Claim("it is the")], evidence()) assert report.is_grounded def test_no_evidence_means_nothing_is_supported(): report = lab.check_grounding([lab.Claim("payment PMT-771 is held")], []) assert not report.is_grounded def test_empty_claims_are_grounded(): report = lab.check_grounding([], evidence()) assert report.is_grounded assert report.coverage == 1.0 def test_grounding_threshold_is_validated(): with pytest.raises(ValueError): lab.check_grounding([], [], min_overlap=0.0) with pytest.raises(ValueError): lab.check_grounding([], [], min_overlap=1.5) def test_a_stricter_threshold_supports_no_more_claims(): claims = [lab.Claim("Payment PMT-771 for AED 250,000 to Acme Trading FZE is HELD")] loose = lab.check_grounding(claims, evidence(), min_overlap=0.4) strict = lab.check_grounding(claims, evidence(), min_overlap=0.99) assert len(strict.supported) <= len(loose.supported) # ====================================================================================== # 9. Context assembly # ====================================================================================== def assemble(**kwargs): kwargs.setdefault("budget_tokens", 400) kwargs.setdefault("instructions", "You are an investigator.") kwargs.setdefault("question", "Why is PMT-771 held?") return lab.assemble_context(**kwargs) def test_segments_are_emitted_in_cache_friendly_order(): context = assemble(tool_schemas="lookup()", policy="no MNPI", memory="RM for Acme", retrieved=evidence()) headers = [line for line in context.text.split("\n") if line.startswith("[") and line.endswith("]")] assert headers == ["[instructions]", "[tool_schemas]", "[policy]", "[memory]", "[retrieved]", "[question]"] def test_the_question_is_always_last(): context = assemble(retrieved=evidence()) assert context.text.rstrip().endswith("Why is PMT-771 held?") def test_the_budget_is_respected(): context = assemble(budget_tokens=120, retrieved=evidence()) assert context.tokens <= 120 def test_lowest_ranked_chunks_are_dropped_first(): chunks = evidence() assert len(chunks) >= 2 context = assemble(budget_tokens=90, retrieved=chunks) assert context.dropped_chunks assert chunks[0].chunk.chunk_id not in context.dropped_chunks def test_dropping_is_reported_not_silent(): context = assemble(budget_tokens=90, retrieved=evidence()) assert set(context.included_chunks) & set(context.dropped_chunks) == set() assert len(context.included_chunks) + len(context.dropped_chunks) == len(evidence()) def test_an_impossible_budget_raises_rather_than_truncating_the_question(): with pytest.raises(ValueError): assemble(budget_tokens=3, instructions="x" * 4000) def test_budget_must_be_positive(): with pytest.raises(ValueError): assemble(budget_tokens=0) def test_stable_prefix_counts_only_stable_segments(): context = assemble(tool_schemas="lookup()", policy="no MNPI", memory="varies per session", retrieved=evidence()) assert context.stable_prefix_tokens > 0 assert 0.0 < context.cacheable_fraction < 1.0 def test_empty_retrieval_omits_the_segment(): context = assemble(retrieved=()) assert "[retrieved]" not in context.text def test_assembly_is_deterministic(): a = assemble(retrieved=evidence()) b = assemble(retrieved=evidence()) assert a == b