src/rag.py

"""
rag.py
──────
The core RAG (Retrieval-Augmented Generation) logic for RoCodex.

Flow for every user question:
  1. Embed the question using the same model used to embed the articles
  2. Search the FAISS index for the top-k most relevant articles
  3. Build a prompt that includes the retrieved articles as context
  4. Send the prompt to Groq (Llama 3) and get back a cited answer
  5. Return the answer + the source articles used

Install:
  pip install groq faiss-cpu sentence-transformers
"""

import os
import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from groq import Groq

# ── Config ────────────────────────────────────────────────────────────────────

MODEL_NAME    = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
INDEX_FILE    = os.path.join(BASE_DIR, "scraper/data/faiss.index")
METADATA_FILE = os.path.join(BASE_DIR, "scraper/data/metadata.jsonl")
TOP_K         = 5  # how many articles to retrieve per question

# Groq model — llama-3.3-70b-versatile is free and handles Romanian very well
GROQ_MODEL = "llama-3.3-70b-versatile"

# System prompt — tells the LLM how to behave as a Romanian legal assistant
SYSTEM_PROMPT = """Ești RoCodex, un asistent juridic specializat în legislația română.

Răspunzi la întrebări despre legi românești bazându-te EXCLUSIV pe articolele de lege furnizate în context.

Reguli stricte:
1. Răspunde NUMAI pe baza articolelor furnizate. Nu inventa informații.
2. Citează întotdeauna sursa: legea și numărul articolului.
3. Dacă articolele furnizate nu conțin răspunsul, spune clar: "Nu am găsit informații relevante în baza de date pentru această întrebare."
4. Răspunde în română, clar și structurat.
5. Nu oferi sfaturi juridice personalizate — indică utilizatorul să consulte un avocat pentru situații specifice.
"""


# ── Singleton loader — loads model/index once, reuses across requests ─────────

_model    = None
_index    = None
_metadata = None


def _load_resources():
    """
    Load the embedding model, FAISS index, and metadata into memory.
    This happens once when the first query is made, then stays in memory.
    Loading takes ~10 seconds the first time (downloading the model).
    Subsequent calls return instantly.
    """
    global _model, _index, _metadata

    if _model is None:
        print("Loading embedding model...")
        _model = SentenceTransformer(MODEL_NAME)

    if _index is None:
        print("Loading FAISS index...")
        _index = faiss.read_index(INDEX_FILE)

    if _metadata is None:
        print("Loading metadata...")
        _metadata = []
        with open(METADATA_FILE, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    _metadata.append(json.loads(line))

    return _model, _index, _metadata


# ── Step 1+2: Retrieve relevant articles ─────────────────────────────────────

def retrieve(query: str, k: int = TOP_K) -> list[dict]:
    """
    Find the k most semantically similar articles to the query.

    Returns a list of dicts, each containing:
      - law_title: e.g. "CODUL MUNCII 24/01/2003"
      - article_number: e.g. "Articolul 145"
      - text: the article's cleaned text
      - score: cosine similarity score (0-1, higher = more relevant)
    """
    model, index, metadata = _load_resources()

    # Embed the query using the same model and settings as the articles
    query_vector = model.encode(
        [query],
        convert_to_numpy=True,
        normalize_embeddings=True,
    ).astype("float32")

    # Search the FAISS index
    scores, positions = index.search(query_vector, k)

    results = []
    for score, pos in zip(scores[0], positions[0]):
        if pos == -1:  # FAISS returns -1 for unfilled slots
            continue
        article = metadata[pos]
        results.append({
            "law_title":      article["law_title"],
            "article_number": article["article_number"],
            "text":           article["text"],
            "score":          float(score),
        })

    return results


# ── Step 3: Build the prompt ──────────────────────────────────────────────────

def build_prompt(query: str, articles: list[dict]) -> str:
    """
    Build the user message that combines the retrieved articles with the question.

    The articles are formatted clearly so the LLM can easily reference them.
    We include the law title and article number so the LLM can cite them.
    """
    context_parts = []
    for i, article in enumerate(articles, 1):
        context_parts.append(
            f"[Sursa {i}] {article['law_title']} — {article['article_number']}\n"
            f"{article['text']}"
        )

    context = "\n\n---\n\n".join(context_parts)

    return f"""Pe baza următoarelor articole din legislația română, răspunde la întrebarea de mai jos.

ARTICOLE RELEVANTE:
{context}

ÎNTREBARE:
{query}

Răspuns (citează sursele folosind [Sursa N]):"""


# ── Step 4+5: Generate answer with Groq ──────────────────────────────────────

def answer(query: str, groq_api_key: str) -> dict:
    """
    Full RAG pipeline: retrieve articles → build prompt → generate answer.

    Parameters:
      query        — the user's question in Romanian
      groq_api_key — your Groq API key (get one free at console.groq.com)

    Returns a dict:
      {
        "answer":   "Răspunsul generat de LLM...",
        "sources":  [list of retrieved articles with scores],
        "query":    "întrebarea originală"
      }
    """
    # Step 1+2: Retrieve relevant articles
    articles = retrieve(query, k=TOP_K)

    if not articles:
        return {
            "answer":  "Nu am putut găsi articole relevante în baza de date.",
            "sources": [],
            "query":   query,
        }

    # Step 3: Build the prompt
    prompt = build_prompt(query, articles)

    # Step 4: Call Groq API
    client = Groq(api_key=groq_api_key)

    chat_completion = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": prompt},
        ],
        temperature=0.1,   # low temperature = more factual, less creative
        max_tokens=1024,
    )

    generated_answer = chat_completion.choices[0].message.content

    # Step 5: Return answer + sources
    return {
        "answer":  generated_answer,
        "sources": articles,
        "query":   query,
    }


# ── Quick CLI test ────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys

    api_key = os.environ.get("GROQ_API_KEY") or (
        input("Introdu Groq API key: ").strip()
    )

    print("\nRoCodex RAG — test CLI")
    print("Scrie 'exit' pentru a ieși.\n")

    while True:
        try:
            query = input("Întrebare: ").strip()
        except (EOFError, KeyboardInterrupt):
            break

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

        print("\nCaut articole relevante și generez răspuns...\n")
        result = answer(query, api_key)

        print("=" * 60)
        print(result["answer"])
        print("\nSURSE FOLOSITE:")
        for i, src in enumerate(result["sources"], 1):
            print(f"  [{i}] scor={src['score']:.3f} | {src['law_title']} | {src['article_number']}")
        print("=" * 60)
        print()