rag-from-scratch/rag/rag.py

#!/usr/bin/env python3
"""
Pipeline RAG interattiva

Riceve una domanda, recupera i chunk più rilevanti da ChromaDB (retrieval)
e genera una risposta tramite Ollama (generation).

Input:  chroma_db/<stem> (collection ChromaDB)
Output: risposta a schermo

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

Nel loop interattivo:
    Domanda: <testo>       → risposta
    Domanda: <testo> -v    → risposta + chunk recuperati
    Domanda: exit          → uscita
"""

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

import chromadb

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

_here        = Path(__file__).parent
project_root = _here.parent
sys.path.insert(0, str(_here))
sys.path.insert(0, str(project_root))
import config as _cfg
from ingestion.config import EMBED_MODEL

CHROMA_DIR = project_root / "chroma_db"

OLLAMA_URL    = _cfg.OLLAMA_URL
LLM_MODEL     = _cfg.OLLAMA_MODEL
TOP_K         = _cfg.TOP_K
TEMPERATURE   = _cfg.TEMPERATURE
NO_THINK      = _cfg.NO_THINK
SYSTEM_PROMPT = _cfg.SYSTEM_PROMPT


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

def embed(text: str) -> list[float]:
    """Genera il vettore della domanda 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"]


# ─── Generazione ──────────────────────────────────────────────────────────────

def call_ollama(prompt: str, system: str = "") -> str:
    """Chiama Ollama /api/generate e ritorna la risposta."""
    payload = json.dumps({
        "model": LLM_MODEL,
        "system": system,
        "prompt": prompt,
        "stream": False,
        "think": not NO_THINK,
        "options": {"temperature": TEMPERATURE},
    }).encode()
    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/generate",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        return json.loads(resp.read())["response"].strip()


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

def retrieve(collection: chromadb.Collection, question: str) -> list[dict]:
    """
    Genera l'embedding della domanda e recupera i TOP_K chunk più simili.
    Ritorna lista di dict con chiavi: text, source, header_path, distance.
    """
    vector = embed(question)
    results = collection.query(
        query_embeddings=[vector],
        n_results=TOP_K,
        include=["documents", "metadatas", "distances"],
    )
    chunks = []
    for text, meta, dist in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0],
    ):
        chunks.append({
            "text":        text,
            "source":      meta.get("source", ""),
            "header_path": meta.get("header_path", ""),
            "distance":    dist,
        })
    return chunks


# ─── Prompt ───────────────────────────────────────────────────────────────────

def build_prompt(question: str, chunks: list[dict]) -> str:
    """Ritorna (system, user_prompt) separati per l'API Ollama."""
    context_parts = []
    for i, c in enumerate(chunks, start=1):
        header = f"[Contesto {i}"
        if c["header_path"]:
            header += f" — {c['header_path']}"
        header += "]"
        context_parts.append(f"{header}\n{c['text']}")

    context = "\n\n".join(context_parts)
    user_prompt = f"{context}\n\nDomanda: {question}"
    return SYSTEM_PROMPT, user_prompt


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

def answer(question: str, collection: chromadb.Collection, verbose: bool) -> None:
    try:
        chunks = retrieve(collection, question)
    except (urllib.error.URLError, OSError) as e:
        print(f"ERRORE embedding: {e}")
        return

    if verbose:
        print("\n── Chunk recuperati ──────────────────────────────────────────")
        for i, c in enumerate(chunks, start=1):
            sim = 1 - c["distance"]
            src = f"[{c['source']}] " if c.get("source") else ""
            loc = c["header_path"] or ""
            print(f"  [{i}] {src}{loc}  (similarità: {sim:.3f})")
            print(f"      {c['text'][:120].replace(chr(10), ' ')}...")
        print("──────────────────────────────────────────────────────────────\n")

    system, prompt = build_prompt(question, chunks)

    try:
        response = call_ollama(prompt, system=system)
    except (urllib.error.URLError, OSError) as e:
        print(f"ERRORE generazione: {e}")
        return

    print(f"\n{response}\n")


def run_loop(collection: chromadb.Collection) -> None:
    print("── Loop RAG ─────────────────────────────────────── (exit per uscire)\n")
    while True:
        try:
            raw = input("Domanda: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nUscita.")
            break

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

        verbose = raw.endswith(" -v")
        question = raw[:-3].strip() if verbose else raw

        answer(question, collection, verbose)


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

def _build_epilog() -> str:
    lines = [
        "Uso:",
        "  python rag.py --stem <nome>",
        "",
        "Loop interattivo:",
        "  <domanda>       risposta basata sul documento",
        "  <domanda> -v    risposta + chunk recuperati con score di similarità",
        "  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=(
            "Pipeline RAG interattiva\n\n"
            "Risponde a domande in linguaggio naturale su un documento\n"
            "indicizzato in ChromaDB da ingestion/ingest.py."
        ),
        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 ...",
    )
    args = parser.parse_args()

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

    print("─── Pipeline RAG ────────────────────────────────────────────\n")
    print(f"  Collection : {collection_name}")
    print(f"  Modello    : {LLM_MODEL}")
    print(f"  Top-K      : {TOP_K}")
    print(f"  Thinking   : {'off' if NO_THINK else 'on'}")
    print()

    if not CHROMA_DIR.exists():
        print("ERRORE: chroma_db/ non trovata — esegui prima ingestion")
        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"ERRORE: Collection '{collection_name}' non trovata in chroma_db/")
        if args.stem:
            print(f"   → python ingestion/ingest.py --stem {collection_name}")
        else:
            print(f"   → python ingestion/ingest.py --collection {collection_name} --stems doc1 doc2 ...")
        return 1

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

    run_loop(collection)
    return 0


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