#!/usr/bin/env python3 """Original LockedIn Labs synthetic context-building exercise, Python 3.11+. Run: python3 fde-retrieval-context-lab.py test python3 fde-retrieval-context-lab.py demo python3 fde-retrieval-context-lab.py evaluate --candidate scoped Standard library only. Lexical ranking, not embeddings, inference or a service. """ import argparse from collections import Counter from dataclasses import asdict, dataclass, replace import hashlib import json import re import socket import unittest from unittest.mock import patch def digest(text): return hashlib.sha256(text.encode("utf-8")).hexdigest() def tokens(text): return re.findall(r"[a-z0-9]+", text.lower()) @dataclass(frozen=True) class Principal: # Trusted host context in a real adapter. The CLI supplies synthetic values. subject: str tenant: str allowed: frozenset @dataclass class Source: tenant: str id: str version: int text: str active: bool = True @dataclass(frozen=True) class Chunk: tenant: str source: str version: int locator: str text: str text_digest: str def rank(sources, query): """Simple term-frequency baseline; stable paragraph locators, no model call.""" terms = set(tokens(query)) candidates = [] for source in sources: for index, paragraph in enumerate(source.text.split("\n\n"), 1): counts = Counter(tokens(paragraph)) score = sum(counts[term] for term in terms) if score: candidates.append((score, Chunk(source.tenant, source.id, source.version, f"paragraph:{index}", paragraph, digest(paragraph)))) return [chunk for _, chunk in sorted(candidates, key=lambda item: (-item[0], item[1].source, item[1].locator))] class ContextBuilder: def __init__(self, sources): if len({(source.tenant, source.id) for source in sources}) != len(sources): raise ValueError("duplicate_source_identity") self.sources = sources self.cache = {} def build(self, principal, query, top_k=2, max_chars=400): if not principal.subject or not principal.tenant: raise ValueError("trusted_identity_required") if not query.strip() or len(query) > 200: raise ValueError("bounded_query_required") if not 1 <= top_k <= 3 or not 1 <= max_chars <= 400: raise ValueError("invalid_context_budget") # Authorize before reading text, chunking or ranking. This lab reads a # local synthetic list; real storage must enforce the same scope. eligible = [source for source in self.sources if source.tenant == principal.tenant and source.id in principal.allowed and source.active] # All currently eligible sources participate, including new sources and # previously empty lookups. Rechecking only cached hits misses additions. snapshot = tuple(sorted((source.id, source.version, digest(source.text)) for source in eligible)) key = (principal.subject, principal.tenant, query, top_k, max_chars) prior = self.cache.get(key) hit = prior is not None and prior[0] == snapshot if hit: selected, omitted = prior[1], prior[2] else: selected, used, omitted = [], 0, 0 for chunk in rank(eligible, query): if len(selected) == top_k: break if used + len(chunk.text) > max_chars: omitted += 1 continue selected.append(chunk) used += len(chunk.text) selected = tuple(selected) self.cache[key] = (snapshot, selected, omitted) return {"state": "context_ready" if selected else "withheld_budget" if omitted else "no_matching_context", "cache": "hit" if hit else "rebuilt", "evidence": [asdict(chunk) for chunk in selected], "text_chars": sum(len(chunk.text) for chunk in selected), "omitted_for_budget": omitted, "handling": "Untrusted source evidence, never tool permission or system instructions.", "limits": "Selected lexical matches only; not complete corpus coverage, answer support or model quality."} def fixtures(): return [Source("north", "P1", 1, "Pump reset requires local isolation.\n\nPump inspection uses the approved checklist."), Source("north", "P2", 1, "Calendar handover requires an operator."), Source("south", "P9", 1, "pump " * 8 + "OTHER_TENANT_MARKER")] def postfilter_example(sources, principal, query): """Deliberately flawed teaching comparison: globally rank, take one, filter. Never use this for context assembly. All records here are synthetic. """ best = rank(sources, query)[:1] allowed = {(source.tenant, source.id) for source in sources if source.tenant == principal.tenant and source.id in principal.allowed and source.active} return [asdict(chunk) for chunk in best if (chunk.tenant, chunk.source) in allowed] class ContextTests(unittest.TestCase): def setUp(self): self.sources = fixtures() self.principal = Principal("learner-north", "north", frozenset({"P1", "P2"})) self.builder = ContextBuilder(self.sources) def test_authorized_context_preserves_paragraph_provenance(self): result = self.builder.build(self.principal, "pump", top_k=1) self.assertEqual(result["state"], "context_ready") self.assertEqual(result["evidence"][0], asdict(Chunk("north", "P1", 1, "paragraph:1", self.sources[0].text.split("\n\n")[0], digest(self.sources[0].text.split("\n\n")[0])))) self.assertNotIn("OTHER_TENANT_MARKER", json.dumps(result)) def test_prefilter_preserves_useful_work_when_foreign_result_dominates(self): self.assertEqual(postfilter_example(self.sources, self.principal, "pump"), []) self.assertEqual(self.builder.build(self.principal, "pump", top_k=1)["evidence"][0]["source"], "P1") def test_revoked_grant_cannot_reuse_cached_text(self): self.builder.build(self.principal, "pump") revoked = replace(self.principal, allowed=frozenset({"P2"})) result = self.builder.build(revoked, "pump") self.assertEqual((result["cache"], result["evidence"]), ("rebuilt", [])) def test_withdrawal_and_unchanged_version_edit_invalidate_cache(self): first = self.builder.build(self.principal, "pump") self.sources[0].text = "Pump reset now needs a revised synthetic checklist." changed = self.builder.build(self.principal, "pump") self.assertNotEqual(first["evidence"][0]["text_digest"], changed["evidence"][0]["text_digest"]) self.assertEqual(changed["cache"], "rebuilt") self.sources[0].active = False self.assertEqual(self.builder.build(self.principal, "pump")["evidence"], []) def test_new_source_invalidates_a_cached_empty_lookup(self): self.assertEqual(self.builder.build(self.principal, "valve")["evidence"], []) self.sources.append(Source("north", "P3", 1, "Valve isolation uses a synthetic permit.")) granted = replace(self.principal, allowed=self.principal.allowed | {"P3"}) result = self.builder.build(granted, "valve") self.assertEqual(result["evidence"][0]["source"], "P3") self.assertEqual(result["cache"], "rebuilt") def test_budget_keeps_whole_chunks_and_never_calls_empty_a_supported_answer(self): result = self.builder.build(self.principal, "pump", max_chars=5) self.assertEqual((result["state"], result["text_chars"]), ("withheld_budget", 0)) self.assertEqual(self.builder.build(self.principal, "absent")["state"], "no_matching_context") with self.assertRaisesRegex(ValueError, "invalid_context_budget"): self.builder.build(self.principal, "pump", top_k=0) def test_cache_is_identity_scoped_and_caller_cannot_mutate_it(self): original = self.builder.build(self.principal, "pump") original["evidence"][0]["text"] = "caller mutation" self.assertEqual(self.builder.build(self.principal, "pump")["cache"], "hit") self.assertNotIn("caller mutation", json.dumps(self.builder.build(self.principal, "pump"))) other = Principal("learner-south", "south", frozenset({"P9"})) result = self.builder.build(other, "pump") self.assertEqual(result["cache"], "rebuilt") self.assertEqual({chunk["source"] for chunk in result["evidence"]}, {"P9"}) def test_duplicate_identity_is_rejected(self): with self.assertRaisesRegex(ValueError, "duplicate_source_identity"): ContextBuilder(self.sources + [self.sources[0]]) @dataclass class RetrievalCase: id: str split: str query: str sources: list principal: Principal relevant: frozenset max_chars: int = 400 def evidence_key(item): return tuple(item[field] for field in ("tenant", "source", "version", "locator", "text_digest")) def source_chunks(sources): for source in sources: for index, paragraph in enumerate(source.text.split("\n\n"), 1): yield asdict(Chunk(source.tenant, source.id, source.version, f"paragraph:{index}", paragraph, digest(paragraph))) def evaluation_cases(): """Authored gold paragraph labels; never inferred from a candidate's ranking. Transfer is visible synthetic practice, not a secret or independent holdout. """ north = Principal("learner-north", "north", frozenset({"P1", "P2", "P4", "P5"})) def case(id, split, query, sources, relevant, principal=north, max_chars=400): labels = frozenset(evidence_key(item) for item in source_chunks(sources) if (item["source"], item["locator"]) in relevant) if len(labels) != len(relevant): raise ValueError("invalid_authored_label") return RetrievalCase(id, split, query, sources, principal, labels, max_chars) return [ case("D1-foreign-dominates", "development", "pump reset", fixtures(), {("P1", "paragraph:1")}), case("D2-two-required-passages", "development", "check", [ Source("north", "P4", 1, "Isolation check comes before opening."), Source("north", "P5", 1, "Restart check follows reassembly."), ], {("P4", "paragraph:1"), ("P5", "paragraph:1")}), case("D3-withdrawn", "development", "pump", [ Source("north", "P1", 1, "Pump reset procedure.", active=False), ], set()), case("D4-revoked", "development", "pump", fixtures(), set(), replace(north, allowed=frozenset())), case("T1-different-terminology", "transfer", "overheating threshold", [ Source("north", "P1", 1, "Thermal limit is forty units."), ], {("P1", "paragraph:1")}), case("T2-withdrawn-distractor", "transfer", "valve", [ Source("north", "P1", 1, "valve " * 12, active=False), Source("north", "P2", 2, "Valve release requires a checklist."), ], {("P2", "paragraph:1")}), case("T3-foreign-only", "transfer", "pump", [ Source("south", "P9", 1, "Pump reset procedure."), ], set(), replace(north, allowed=frozenset({"P9"}))), case("T4-evidence-does-not-fit", "transfer", "permit", [ Source("north", "P1", 1, "Permit review requires the complete synthetic instruction before any release."), ], {("P1", "paragraph:1")}, max_chars=20), ] def retrieve_case(case, candidate): if candidate == "scoped": return ContextBuilder(case.sources).build(case.principal, case.query, top_k=2, max_chars=case.max_chars)["evidence"] if candidate == "refuse-all": return [] if candidate != "postfilter": raise ValueError("unknown_candidate") # Keep the same whole-chunk context budget around the deliberately flawed ranker. items = postfilter_example(case.sources, case.principal, case.query) selected, used = [], 0 for item in items: if used + len(item["text"]) <= case.max_chars: selected.append(item) used += len(item["text"]) return selected def score_case(case, evidence): eligible = [source for source in case.sources if source.tenant == case.principal.tenant and source.id in case.principal.allowed and source.active] permitted = {evidence_key(item) for item in source_chunks(eligible)} keys = [evidence_key(item) for item in evidence] valid = [key in permitted and digest(item["text"]) == item["text_digest"] for item, key in zip(evidence, keys)] hits = {key for key, ok in zip(keys[:2], valid[:2]) if ok and key in case.relevant} first = next((index for index, (key, ok) in enumerate(zip(keys[:2], valid[:2]), 1) if ok and key in case.relevant), None) return {"id": case.id, "split": case.split, "query": case.query, "context_budget_chars": case.max_chars, "required": [list(key) for key in sorted(case.relevant)], "retrieved": [list(key) for key in keys], "relevant_found": len(hits), "recall_at_2": len(hits) / len(case.relevant) if case.relevant else None, "reciprocal_rank_at_2": 1 / first if first else 0 if case.relevant else None, "correct_empty": not evidence if not case.relevant else None, "invalid_evidence": valid.count(False), "context_budget_violation": len(evidence) > 2 or sum(len(item["text"]) for item in evidence) > case.max_chars} def evaluate(candidate="scoped", split="all", omitted=()): if candidate not in ("scoped", "postfilter", "refuse-all") or split not in ("all", "development", "transfer"): raise ValueError("unknown_candidate_or_split") planned = [case for case in evaluation_cases() if split == "all" or case.split == split] omitted = set(omitted) if omitted - {case.id for case in planned}: raise ValueError("omitted_case_not_in_selected_packet") rows = [score_case(case, retrieve_case(case, candidate)) for case in planned if case.id not in omitted] answerable = [row for row in rows if row["required"]] empty = [row for row in rows if not row["required"]] recall_sum = sum(row["recall_at_2"] for row in answerable) rr_sum = sum(row["reciprocal_rank_at_2"] for row in answerable) return {"candidate": candidate, "split": split, "top_k": 2, "planned": len(planned), "executed": len(rows), "unrun": sorted(omitted), "complete": not omitted, "answerable_planned": sum(bool(case.relevant) for case in planned), "unanswerable_planned": sum(not case.relevant for case in planned), "answerable_executed": len(answerable), "unanswerable_executed": len(empty), "mean_recall_at_2": {"sum": recall_sum, "queries": len(answerable), "value": recall_sum / len(answerable) if answerable else None}, "mrr_at_2": {"sum": rr_sum, "queries": len(answerable), "value": rr_sum / len(answerable) if answerable else None}, "correct_empty": sum(row["correct_empty"] for row in empty), "invalid_evidence": sum(row["invalid_evidence"] for row in rows), "context_budget_violations": sum(row["context_budget_violation"] for row in rows), "cases": rows, "limits": "Eight visible synthetic cases. Macro query averages on executed answerable cases only; no universal pass score, model answers, embeddings, independent holdout or production outcome."} class RetrievalEvaluationTests(unittest.TestCase): def test_known_recall_and_rank_denominators_retain_lexical_misses(self): report = evaluate() self.assertEqual((report["planned"], report["executed"], report["answerable_executed"]), (8, 8, 5)) self.assertEqual(report["mean_recall_at_2"], {"sum": 3.0, "queries": 5, "value": .6}) self.assertEqual(report["mrr_at_2"]["value"], .6) misses = [row["id"] for row in report["cases"] if row["recall_at_2"] == 0] self.assertEqual(misses, ["T1-different-terminology", "T4-evidence-does-not-fit"]) self.assertEqual((report["correct_empty"], report["unanswerable_executed"], report["invalid_evidence"]), (3, 3, 0)) def test_second_rank_and_duplicate_hits_keep_different_denominators(self): case = evaluation_cases()[0] first, second = retrieve_case(case, "scoped") row = score_case(case, [second, first]) self.assertEqual((row["recall_at_2"], row["reciprocal_rank_at_2"]), (1, .5)) def test_refusing_everything_does_not_earn_retrieval_quality(self): report = evaluate("refuse-all") self.assertEqual(report["correct_empty"], 3) self.assertEqual(report["mean_recall_at_2"]["value"], 0) self.assertEqual(evaluate("postfilter")["mean_recall_at_2"]["value"], .1) self.assertEqual(evaluate("postfilter")["mrr_at_2"]["value"], .2) def test_visible_transfer_packet_is_separate_from_development(self): self.assertEqual(evaluate(split="development")["mean_recall_at_2"]["value"], 1) self.assertEqual(evaluate(split="transfer")["mean_recall_at_2"]["value"], 1 / 3) def test_omission_is_unrun_and_zero_denominator_is_not_perfect(self): report = evaluate(omitted=["T1-different-terminology"]) self.assertFalse(report["complete"]) self.assertEqual(report["unrun"], ["T1-different-terminology"]) self.assertEqual(report["mean_recall_at_2"], {"sum": 3.0, "queries": 4, "value": .75}) none = evaluate(omitted=[case.id for case in evaluation_cases()]) self.assertIsNone(none["mrr_at_2"]["value"]) with self.assertRaisesRegex(ValueError, "omitted_case_not_in_selected_packet"): evaluate(split="development", omitted=["T1-different-terminology"]) def test_foreign_stale_and_changed_payload_cannot_be_counted_as_valid_evidence(self): case = evaluation_cases()[0] item = retrieve_case(case, "scoped")[0] for field, value in (("tenant", "south"), ("version", 0), ("text", "changed text")): with self.subTest(field=field): row = score_case(case, [{**item, field: value}]) self.assertEqual((row["invalid_evidence"], row["recall_at_2"]), (1, 0)) two = evaluation_cases()[1] duplicate = retrieve_case(two, "scoped")[0] self.assertEqual(score_case(two, [duplicate, duplicate])["recall_at_2"], .5) def demo(): sources = fixtures() principal = Principal("learner-north", "north", frozenset({"P1", "P2"})) builder = ContextBuilder(sources) reports = {"postfilter_top_one": postfilter_example(sources, principal, "pump"), "authorized": builder.build(principal, "pump"), "unchanged": builder.build(principal, "pump")} sources[0].text = "Pump reset now requires a revised synthetic checklist." reports["edited_same_version"] = builder.build(principal, "pump") reports["revoked"] = builder.build(replace(principal, allowed=frozenset()), "pump") print(json.dumps(reports, indent=2)) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("command", choices=["test", "demo", "evaluate"]) parser.add_argument("--candidate", choices=["scoped", "postfilter", "refuse-all"], default="scoped") parser.add_argument("--split", choices=["all", "development", "transfer"], default="all") parser.add_argument("--omit-case", action="append", default=[]) args = parser.parse_args() with patch.object(socket.socket, "connect", side_effect=AssertionError("network_forbidden")): if args.command == "test": suite = unittest.TestSuite(unittest.defaultTestLoader.loadTestsFromTestCase(case) for case in (ContextTests, RetrievalEvaluationTests)) result = unittest.TextTestRunner(verbosity=2).run(suite) raise SystemExit(0 if result.wasSuccessful() else 1) elif args.command == "evaluate": try: report = evaluate(args.candidate, args.split, args.omit_case) except ValueError as error: parser.error(str(error)) print(json.dumps(report, indent=2)) raise SystemExit(0 if report["complete"] else 2) else: demo()