205 lines
6.8 KiB
Python
205 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Self-contained tests for raglite.py.
|
|
|
|
Creates its own fixtures in temp dirs, exercises the library and the CLI,
|
|
and exits 0 only if every check passes. No network, no external deps.
|
|
|
|
Run: python3 /root/.hermes/wizard-outputs/null/test_raglite.py
|
|
"""
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(HERE))
|
|
|
|
import raglite # noqa: E402
|
|
|
|
FAILURES = []
|
|
|
|
|
|
def check(name, cond, detail=""):
|
|
if cond:
|
|
print(f" ok {name}")
|
|
else:
|
|
print(f" FAIL {name} {detail}")
|
|
FAILURES.append(name)
|
|
|
|
|
|
def doc_of(pid: str) -> str:
|
|
"""'path/to/codeforge.md::3' -> 'codeforge.md'"""
|
|
return pid.split("::")[0].rsplit("/", 1)[-1]
|
|
|
|
|
|
DOCS = {
|
|
"codeforge.md": (
|
|
"The wizard forged the spell in the codeforge.\n"
|
|
"\n"
|
|
"The codeforge hums with runes and ember-light.\n"
|
|
),
|
|
"kibana.md": (
|
|
"Kibana dashboards visualize logs from the elasticsearch cluster.\n"
|
|
"\n"
|
|
"Operators drill into slow queries here.\n"
|
|
),
|
|
"bm25.md": (
|
|
"BM25 ranks passages by term frequency and inverse document frequency.\n"
|
|
"\n"
|
|
"It is a bag-of-words retrieval function.\n"
|
|
),
|
|
"standup.md": (
|
|
"The troupe stands up at dawn.\n"
|
|
"\n"
|
|
"Each wizard reports their lane and blockers.\n"
|
|
),
|
|
}
|
|
|
|
|
|
def write_corpus(root):
|
|
for name, body in DOCS.items():
|
|
(root / name).write_text(body, encoding="utf-8")
|
|
|
|
|
|
def test_tokenize():
|
|
print("[tokenize]")
|
|
toks = raglite.tokenize("Hello, WORLD! It's a BM25-test.")
|
|
check("punctuation stripped, lowercased",
|
|
toks == ["hello", "world", "it's", "a", "bm25", "test"], repr(toks))
|
|
check("empty input -> empty list", raglite.tokenize("") == [])
|
|
|
|
|
|
def test_split_passages():
|
|
print("[split_passages]")
|
|
parts = raglite.split_passages("a\n\nb\n\n\nc\n")
|
|
check("blank-line split", parts == ["a", "b", "c"], repr(parts))
|
|
check("whitespace-only block dropped", raglite.split_passages("a\n \nb") == ["a", "b"])
|
|
|
|
|
|
def test_index_and_ranking(tmp):
|
|
print("[index + ranking]")
|
|
idx = raglite.build_index([str(tmp)])
|
|
check("4 docs x 2 paragraphs -> 8 passages", idx["n_docs"] == 8, idx["n_docs"])
|
|
check("avgdl is positive", idx["avgdl"] > 0, idx["avgdl"])
|
|
|
|
h = raglite.search(idx, "codeforge", k=3)
|
|
check("query 'codeforge' returns hits", len(h) > 0)
|
|
check("top hit is the codeforge passage",
|
|
bool(h) and doc_of(h[0]["id"]) == "codeforge.md"
|
|
and "codeforge" in h[0]["text"],
|
|
h[0]["id"] if h else "none")
|
|
|
|
h = raglite.search(idx, "elasticsearch logs", k=3)
|
|
check("top hit for 'elasticsearch logs' is kibana",
|
|
bool(h) and doc_of(h[0]["id"]) == "kibana.md", h[0]["id"] if h else "none")
|
|
|
|
h = raglite.search(idx, "bm25 ranking", k=3)
|
|
check("top hit for 'bm25 ranking' is bm25.md",
|
|
bool(h) and doc_of(h[0]["id"]) == "bm25.md", h[0]["id"] if h else "none")
|
|
|
|
h = raglite.search(idx, "nonexistentterm", k=5)
|
|
check("unknown term -> no hits", h == [], repr(h))
|
|
|
|
h = raglite.search(idx, "the", k=5)
|
|
check("scores strictly positive", all(x["score"] > 0 for x in h), repr(h))
|
|
check("scores sorted descending",
|
|
all(h[i]["score"] >= h[i + 1]["score"] for i in range(len(h) - 1)), repr(h))
|
|
|
|
# Negative control: a kibana query must not surface the codeforge doc.
|
|
h = raglite.search(idx, "kibana dashboards", k=5)
|
|
check("negative control: codeforge not ranked for kibana query",
|
|
all(doc_of(x["id"]) != "codeforge.md" for x in h),
|
|
repr([x["id"] for x in h]))
|
|
|
|
|
|
def test_roundtrip(tmp):
|
|
print("[json round-trip]")
|
|
idx = raglite.build_index([str(tmp)])
|
|
path = tmp / "index.json"
|
|
raglite.save_index(idx, path)
|
|
check("index file written", path.exists() and path.stat().st_size > 0)
|
|
|
|
loaded = raglite.load_index(path)
|
|
a = raglite.search(idx, "codeforge", k=5)
|
|
b = raglite.search(loaded, "codeforge", k=5)
|
|
check("loaded index gives identical hits",
|
|
[(x["id"], x["score"]) for x in a] == [(x["id"], x["score"]) for x in b],
|
|
f"{a} != {b}")
|
|
|
|
|
|
def test_exact_bm25():
|
|
"""Compare against BM25 values computed independently (bm25_probe.py)."""
|
|
print("[exact bm25 constants]")
|
|
with tempfile.TemporaryDirectory(prefix="raglite_tiny_") as d:
|
|
tiny = Path(d)
|
|
(tiny / "d1.md").write_text("a a b\n", encoding="utf-8")
|
|
(tiny / "d2.md").write_text("a c\n", encoding="utf-8")
|
|
idx = raglite.build_index([str(tiny)])
|
|
|
|
# Independently recomputed with the same Robertson formula (bm25_probe.py).
|
|
EXPECTED = {
|
|
("d1.md", "a"): 0.24472692187108,
|
|
("d2.md", "a"): 0.20035335911423582,
|
|
("d1.md", "b"): 0.6359148445504086,
|
|
("d2.md", "c"): 0.7617001984175223,
|
|
}
|
|
for (doc, term), want in EXPECTED.items():
|
|
hits = raglite.search(idx, term, k=10)
|
|
got = next((x["score"] for x in hits if doc_of(x["id"]) == doc), None)
|
|
check(f"score[{doc},{term}] == {want:.12f}",
|
|
got is not None and abs(got - want) < 1e-12,
|
|
f"got {got!r}")
|
|
|
|
# Higher term frequency must win for the same term.
|
|
hits = raglite.search(idx, "a", k=2)
|
|
check("tf ordering: d1 (tf=2) ranks above d2 (tf=1)",
|
|
len(hits) == 2 and doc_of(hits[0]["id"]) == "d1.md", repr(hits))
|
|
|
|
|
|
def test_cli(tmp):
|
|
print("[cli]")
|
|
out = tmp / "cli_index.json"
|
|
r = subprocess.run([sys.executable, str(HERE / "raglite.py"),
|
|
"index", str(tmp), "--out", str(out)],
|
|
capture_output=True, text=True)
|
|
check("cli index exits 0", r.returncode == 0, r.stderr[-400:])
|
|
check("cli index reports 8 passages", "8 passages" in r.stdout, r.stdout.strip())
|
|
|
|
r = subprocess.run([sys.executable, str(HERE / "raglite.py"),
|
|
"query", str(out), "elasticsearch logs", "-k", "2", "--json"],
|
|
capture_output=True, text=True)
|
|
check("cli query exits 0", r.returncode == 0, r.stderr[-400:])
|
|
try:
|
|
hits = json.loads(r.stdout)
|
|
except json.JSONDecodeError as e:
|
|
check("cli query emits valid JSON", False, f"{e}: {r.stdout[:200]}")
|
|
return
|
|
check("cli query emits valid JSON", True)
|
|
check("cli top hit is kibana",
|
|
bool(hits) and doc_of(hits[0]["id"]) == "kibana.md", repr(hits))
|
|
|
|
|
|
def main():
|
|
print("raglite test suite")
|
|
test_tokenize()
|
|
test_split_passages()
|
|
with tempfile.TemporaryDirectory(prefix="raglite_test_") as d:
|
|
tmp = Path(d)
|
|
write_corpus(tmp)
|
|
test_index_and_ranking(tmp)
|
|
test_roundtrip(tmp)
|
|
test_cli(tmp)
|
|
test_exact_bm25()
|
|
print()
|
|
if FAILURES:
|
|
print(f"FAILED ({len(FAILURES)}): " + ", ".join(FAILURES))
|
|
return 1
|
|
print("ALL CHECKS PASSED")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|