ask.py

"""
ask.py — Etape 2 : poser une question a ses documents.

    python ask.py "quelle est la duree de la garantie ?"
    python ask.py "..." --k 6 --show-context
    python ask.py                       (mode interactif)
"""

from __future__ import annotations

import argparse
import sys
import time

from llm import DEFAULT_MODEL, ask_ollama, build_prompt
from rag_core import Embedder, VectorStore


def answer(question: str, store: VectorStore, embedder: Embedder,
           k: int = 4, model: str = DEFAULT_MODEL, show_context: bool = False) -> None:
    t0 = time.time()

    # 1. la question devient un vecteur, dans le meme espace que les documents
    qvec = embedder.encode([question])[0]

    # 2. recherche des k morceaux les plus proches
    passages = store.search(qvec, k=k)

    print(f"\nSources retenues ({len(passages)}) :")
    for i, (chunk, score) in enumerate(passages, start=1):
        loc = chunk.source + (f" p.{chunk.page}" if chunk.page else "")
        apercu = chunk.text[:90].replace("\n", " ")
        print(f"  {i}. [{score:.3f}] {loc} — {apercu}...")

    if show_context:
        print("\n--- Contexte envoye au modele ---")
        for i, (chunk, _) in enumerate(passages, start=1):
            print(f"\n[Extrait {i} | {chunk.source}]\n{chunk.text}")
        print("--- fin du contexte ---")

    # 3. generation ancree sur ces extraits
    prompt = build_prompt(question, passages)
    print("\nGeneration en cours...\n")
    reponse = ask_ollama(prompt, model=model)

    print("=" * 70)
    print(reponse)
    print("=" * 70)
    print(f"({time.time() - t0:.1f} s — modele {model}, k={k})")


def main() -> None:
    ap = argparse.ArgumentParser(description="Interroge tes documents.")
    ap.add_argument("question", nargs="*", help="la question (vide = mode interactif)")
    ap.add_argument("--k", type=int, default=4, help="nombre d'extraits a recuperer")
    ap.add_argument("--model", default=DEFAULT_MODEL, help="modele Ollama")
    ap.add_argument("--index", default="index")
    ap.add_argument("--show-context", action="store_true",
                    help="afficher les extraits envoyes au modele")
    args = ap.parse_args()

    print("Chargement de l'index et du modele d'embeddings...")
    try:
        store = VectorStore.load(args.index)
    except FileNotFoundError as e:
        print(f"\n{e}")
        sys.exit(1)
    embedder = Embedder(store.model_name)
    print(f"   {len(store.chunks)} morceaux indexes.")

    if args.question:
        answer(" ".join(args.question), store, embedder,
               args.k, args.model, args.show_context)
        return

    print("\nMode interactif — 'quit' pour sortir.")
    while True:
        try:
            q = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if q.lower() in {"quit", "exit", "q", ""}:
            break
        answer(q, store, embedder, args.k, args.model, args.show_context)


if __name__ == "__main__":
    main()