rag-from-scratch/retrieve.py

#!/usr/bin/env python3
"""
Retrieval puro (senza generazione LLM)

Loop interattivo: inserisci una query, ottieni i chunk più simili dalla
collection ChromaDB tramite embedding semantico — senza chiamare Ollama
per la generation.

Utile per:
  - verificare la qualità del retrieval prima di diagnosticare risposte sbagliate
  - controllare che i chunk giusti vengano recuperati per una query
  - usare la pipeline come motore di ricerca semantica

Input:  chroma_db/<stem> (collection ChromaDB)
Output: lista chunk con score di similarità

Uso:
    python retrieve.py --stem <nome>

Nel loop interattivo:
    Query: <testo>      → chunk più simili con score
    Query: <testo> -f   → testo completo dei chunk
    Query: exit         → uscita
"""

import argparse
import json
import sys
import urllib.error
import urllib.request
from pathlib import Path

import chromadb

# ─── Configurazione ───────────────────────────────────────────────────────────

sys.path.insert(0, str(Path(__file__).parent))
import config as _cfg

project_root = Path(__file__).parent
CHROMA_DIR   = project_root / "chroma_db"

OLLAMA_URL  = _cfg.OLLAMA_URL
EMBED_MODEL = _cfg.EMBED_MODEL
TOP_K       = _cfg.TOP_K


# ─── Embedding ────────────────────────────────────────────────────────────────

def embed(text: str) -> list[float]:
    """Genera il vettore della query tramite Ollama."""
    payload = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/embeddings",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())["embedding"]


# ─── Retrieval ────────────────────────────────────────────────────────────────

def retrieve(collection: chromadb.Collection, query: str, top_k: int) -> list[dict]:
    """
    Genera l'embedding della query e recupera i top_k chunk più simili.
    Ritorna lista di dict con chiavi: rank, similarity, sezione, titolo, text.
    """
    vector = embed(query)
    results = collection.query(
        query_embeddings=[vector],
        n_results=top_k,
        include=["documents", "metadatas", "distances"],
    )
    chunks = []
    for rank, (text, meta, dist) in enumerate(
        zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        ),
        start=1,
    ):
        chunks.append({
            "rank":       rank,
            "similarity": round(1 - dist, 4),
            "source":     meta.get("source", ""),
            "sezione":    meta.get("sezione", ""),
            "titolo":     meta.get("titolo", ""),
            "text":       text,
        })
    return chunks


# ─── Output ───────────────────────────────────────────────────────────────────

def print_results(chunks: list[dict], full: bool = False) -> None:
    print(f"── {len(chunks)} chunk recuperati ─────────────────────────────────\n")
    for c in chunks:
        src = f"[{c['source']}] " if c.get("source") else ""
        loc = c["sezione"]
        if c["titolo"]:
            loc += f" > {c['titolo']}"
        print(f"  [{c['rank']}] similarità: {c['similarity']:.4f}  |  {src}{loc}")
        if full:
            print()
            print(c["text"])
        else:
            print(f"      {c['text'][:200].replace(chr(10), ' ')}")
            if len(c["text"]) > 200:
                print(f"      … ({len(c['text'])} caratteri totali)")
        print()


# ─── Loop interattivo ─────────────────────────────────────────────────────────

def run_loop(collection: chromadb.Collection, top_k: int) -> None:
    print("── Loop retrieval ──────────────────────── (exit per uscire, -f per testo completo)\n")
    while True:
        try:
            raw = input("Query: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nUscita.")
            break

        if not raw:
            continue
        if raw.lower() == "exit":
            break

        full = raw.endswith(" -f")
        query = raw[:-3].strip() if full else raw

        try:
            chunks = retrieve(collection, query, top_k)
        except (urllib.error.URLError, OSError) as e:
            print(f"❌ Errore embedding (Ollama raggiungibile?): {e}\n")
            continue

        print()
        print_results(chunks, full=full)


# ─── Entry point ──────────────────────────────────────────────────────────────

def _build_epilog() -> str:
    lines = [
        "Uso:",
        "  python retrieve.py --stem <nome>",
        "",
        "Nel loop interattivo:",
        "  <query>       chunk più simili con score (testo troncato)",
        "  <query> -f    testo completo dei chunk",
        "  exit          termina",
    ]
    if CHROMA_DIR.exists():
        try:
            client = chromadb.PersistentClient(path=str(CHROMA_DIR))
            names = [c.name for c in client.list_collections()]
            if names:
                lines += ["", f"Collection disponibili: {', '.join(names)}"]
            else:
                lines += ["", "Nessuna collection trovata — eseguire prima: python ingestion/ingest.py"]
        except Exception:
            pass
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Retrieval puro (senza LLM)\n\n"
            "Loop interattivo: inserisci una query e ottieni i chunk più simili\n"
            "tramite embedding semantico, senza generazione LLM."
        ),
        epilog=_build_epilog(),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--stem",
        help="Collection di un singolo documento (retrocompatibile)",
    )
    parser.add_argument(
        "--collection",
        help="Collection multi-documento creata con: ingest.py --collection <nome> --stems ...",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=TOP_K,
        metavar="N",
        help=f"Numero di chunk da restituire per query (default: {TOP_K} da config.py).",
    )
    args = parser.parse_args()

    collection_name = args.collection or args.stem
    if not collection_name:
        parser.error("specifica --stem <nome> oppure --collection <nome>")

    print("─── Retrieval puro ──────────────────────────────────────────\n")
    print(f"  Collection   : {collection_name}")
    print(f"  Embed model  : {EMBED_MODEL}")
    print(f"  Top-K        : {args.top_k}")
    print()

    if not CHROMA_DIR.exists():
        print("❌ chroma_db/ non trovata — esegui prima ingestion", file=sys.stderr)
        return 1

    client = chromadb.PersistentClient(path=str(CHROMA_DIR))
    collections = [c.name for c in client.list_collections()]
    if collection_name not in collections:
        print(f"❌ Collection '{collection_name}' non trovata in chroma_db/", file=sys.stderr)
        if args.stem:
            print(f"   → python ingestion/ingest.py --stem {collection_name}", file=sys.stderr)
        else:
            print(f"   → python ingestion/ingest.py --collection {collection_name} --stems doc1 doc2 ...", file=sys.stderr)
        return 1

    collection = client.get_collection(collection_name)
    print(f"✅ Collection '{collection_name}' caricata ({collection.count()} chunk)\n")

    run_loop(collection, args.top_k)
    return 0


if __name__ == "__main__":
    sys.exit(main())