RICS / backend /tests /test_rag_store.py
StormShadow308's picture
Add demo documentation and Docker setup for v2 report generation system
aad7814
Raw
History Blame Contribute Delete
2.99 kB
from __future__ import annotations
from backend.core.rag_store import TIER_MASTER, TIER_REFERENCE, Chunk, get_rag_store
def test_master_not_scrubbed_reference_scrubbed():
store = get_rag_store()
store.ingest_document(
"t1", "master:x",
[Chunk(text="Roof covering is slate at SW1A 1AA.", section_id="D2", tier=TIER_MASTER)],
tier=TIER_MASTER,
)
store.ingest_document(
"t1", "ref:y",
[Chunk(text="Roof covering is slate at SW1A 1AA.", tier=TIER_REFERENCE)],
tier=TIER_REFERENCE,
)
master_hits = store.search("t1", "roof slate", tier=TIER_MASTER, top_k=1)
assert "SW1A 1AA" in master_hits[0].text # master kept verbatim
ref_hits = store.search("t1", "roof slate", tier=TIER_REFERENCE, top_k=1)
assert "SW1A 1AA" not in ref_hits[0].text # reference was scrubbed at ingest
def test_master_ranked_ahead_of_reference():
store = get_rag_store()
store.ingest_document(
"t2", "m", [Chunk(text="central heating gas boiler", section_id="G4", tier=TIER_MASTER)],
tier=TIER_MASTER,
)
store.ingest_document(
"t2", "r", [Chunk(text="central heating gas boiler", tier=TIER_REFERENCE)],
tier=TIER_REFERENCE,
)
hits = store.search("t2", "central heating gas boiler", top_k=2)
assert hits[0].tier == TIER_MASTER
def test_counts_and_clear():
store = get_rag_store()
store.ingest_document(
"t3", "m", [Chunk(text="alpha", tier=TIER_MASTER)], tier=TIER_MASTER
)
assert store.count("t3", TIER_MASTER) == 1
store.clear_tier("t3", TIER_MASTER)
assert store.count("t3", TIER_MASTER) == 0
def test_search_for_generation_never_returns_reference_tier():
"""Past reports must not enter the mapping pipeline."""
store = get_rag_store()
store.ingest_document(
"t5",
"master:m",
[Chunk(text="Chimney stack inspected from ground level.", section_id="E1", tier=TIER_MASTER)],
tier=TIER_MASTER,
)
store.ingest_document(
"t5",
"ref:r",
[Chunk(text="Previous client chimney at 14 Oak Street had severe cracking.", tier=TIER_REFERENCE)],
tier=TIER_REFERENCE,
)
hits = store.search_for_generation("t5", "chimney stack", top_k=5)
assert len(hits) >= 1
assert all(h.tier == TIER_MASTER for h in hits)
assert all("Oak Street" not in h.text for h in hits)
def test_unscrubbed_reference_excluded_from_search():
"""Reference chunks with is_scrubbed=False must never appear in search results."""
store = get_rag_store()
store.ingest_document(
"t4",
"ref:ok",
[Chunk(text="Roof covering comprises slate tiles.", tier=TIER_REFERENCE)],
tier=TIER_REFERENCE,
)
ti = store._get("t4", TIER_REFERENCE)
ti.meta[0]["is_scrubbed"] = False
ti.meta[0]["text"] = "Secret client data at SW1A 1AA."
hits = store.search("t4", "secret client", tier=TIER_REFERENCE, top_k=5)
assert hits == []