File size: 1,739 Bytes
c2dcd30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
#!/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())