"""
test_retrieval.py — STEP 3: Prove retrieval works (NO generation yet).

This is the most important test. Type a note, and it shows the most
similar real findings from your data. If these look relevant, RAG's
foundation is working — independent of any generation model.

Run:
    source venv/bin/activate
    python test_retrieval.py
"""

import chromadb
from sentence_transformers import SentenceTransformer

DB_PATH = "./nc_vectors"
COLLECTION = "findings"

embedder = SentenceTransformer("all-MiniLM-L6-v2")
client = chromadb.PersistentClient(path=DB_PATH)
col = client.get_collection(COLLECTION)
print(f"Loaded {col.count()} findings from the vector DB.\n")

while True:
    note = input("Type a short note (or 'quit'): ").strip()
    if note.lower() in ("quit", "exit", ""):
        break
    qvec = embedder.encode(note).tolist()
    hits = col.query(query_embeddings=[qvec], n_results=3)
    docs = hits["documents"][0]
    dists = hits["distances"][0]
    print("\n--- Top 3 similar findings from YOUR data ---")
    for i, (d, dist) in enumerate(zip(docs, dists), 1):
        preview = " ".join(d.split())[:300]
        print(f"\n[{i}] (similarity score: {1 - dist:.3f})")
        print(preview)
    print("\n" + "=" * 60 + "\n")