factory-rag / scripts /check_retrieval.py
Hypn0's picture
fix: plain-dict chunks from Qdrant ScoredPoints; doc_name/score in chunk payload; retrieval smoke 3/3
c2dcd30
Raw History Blame Contribute Delete
1.74 kB
#!/usr/bin/env python3
"""Fresh-process retrieval smoke test for the KX-40 manual.
Runs the three ground-truth KX-40 queries through the full Retriever
(hybrid RRF + rerank, exactly what app.py uses) and prints the TOP-3
per query with page and score. PASS iff the expected marker string
appears in one of the top-3 chunk texts for each query. Exit code is
nonzero on any FAIL.
"""
from __future__ import annotations
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
QUERIES = [
("What is the maximum torque of the KX-40?", "45 Nm", 1),
("What is the part number in the KX-40 calibration instructions?", "KX40-8817", 2),
("What is the pressure setpoint for the KX-40?", "2.4 bar", 2),
]
def main() -> int:
import retrieve
retr = retrieve.Retriever()
failed = False
try:
for query, marker, want_page in QUERIES:
res = retr.retrieve(query, top_n=3)
chunks = res["chunks"]
print(f"\nQ: {query}")
for i, c in enumerate(chunks, 1):
text = c.get("text", "")
score = c.get("score")
print(f" [{i}] page={c.get('page')} score={score} "
f"text={text[:80]!r}")
hit = any(marker in c.get("text", "") for c in chunks)
print(f" -> {'PASS' if hit else 'FAIL'} "
f"(expected {marker!r} on page {want_page}, "
f"got pages {[c.get('page') for c in chunks]})")
if not hit:
failed = True
finally:
retr.close()
print("\nRESULT:", "FAIL" if failed else "PASS 3/3")
return 1 if failed else 0
if __name__ == "__main__":
sys.exit(main())