rag_engine.py

# -*- coding: utf-8 -*-
"""
Moteur RAG (Space edition) : recuperation + generation via l'API Hugging Face
(ou Ollama en repli local). Reponses fondees uniquement sur le contenu indexe
du site (public + non public selon les droits).
"""
import pickle
import functools

import requests

import config
import gdpr


@functools.lru_cache(maxsize=1)
def _get_embedder():
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer(config.EMBEDDING_MODEL)


@functools.lru_cache(maxsize=1)
def _get_index_and_chunks():
    import faiss
    index = faiss.read_index(str(config.INDEX_FILE))
    with open(config.CHUNKS_FILE, "rb") as f:
        chunks = pickle.load(f)
    return index, chunks


def reset_cache():
    _get_index_and_chunks.cache_clear()


# --------------------------------------------------------------------------
# Recuperation / retrieval
# --------------------------------------------------------------------------
def retrieve(query, top_k=None):
    top_k = top_k or config.TOP_K
    index, chunks = _get_index_and_chunks()
    embedder = _get_embedder()
    q_vec = embedder.encode([query], normalize_embeddings=True,
                            convert_to_numpy=True).astype("float32")
    scores, idx = index.search(q_vec, top_k)
    results = []
    for score, i in zip(scores[0], idx[0]):
        if i < 0 or score < config.MIN_SCORE:
            continue
        c = dict(chunks[i])
        c["score"] = float(score)
        results.append(c)
    return results


# --------------------------------------------------------------------------
# Prompt
# --------------------------------------------------------------------------
SYSTEM_FR = (
    "Tu es l'assistant interne du site Vizyon Ayiti 360. "
    "Tu reponds EXCLUSIVEMENT a partir du CONTEXTE fourni (extraits du site, "
    "publics et internes). Interdiction d'utiliser des connaissances externes. "
    "Si la reponse ne figure pas dans le contexte, dis-le clairement. Cite les "
    "titres/URL des sources. Reponds dans la langue de la question."
)
SYSTEM_EN = (
    "You are the internal assistant of the Vizyon Ayiti 360 website. "
    "You answer EXCLUSIVELY from the provided CONTEXT (public and internal site "
    "excerpts). Do not use external knowledge. If the answer is not in the "
    "context, say so clearly. Cite source titles/URLs. Reply in the language of "
    "the question."
)
SYSTEM_HT = (
    "Ou se asistan enten sit Vizyon Ayiti 360 a. Ou reponn SELMAN ak KONTEKS yo "
    "ba ou a (ekstre piblik ak enten nan sit la). Ou pa gen dwa itilize okenn "
    "lot konesans deyo. Si repons lan pa nan konteks la, di sa kle. Site tit ak "
    "adres (URL) sous yo. Reponn an kreyol ayisyen."
)
_SYSTEMS = {"fr": SYSTEM_FR, "en": SYSTEM_EN, "ht": SYSTEM_HT}


def build_prompt(query, passages, lang="fr"):
    blocks = []
    for i, p in enumerate(passages, 1):
        tag = "" if p.get("status", "publish") == "publish" else " [INTERNE]"
        blocks.append(f"[Source {i}]{tag} {p['title']} ({p['url']})\n{p['text']}")
    context = "\n\n".join(blocks) if blocks else "(aucun extrait pertinent)"
    system = _SYSTEMS.get(lang, SYSTEM_FR)
    user = (
        f"CONTEXTE:\n{context}\n\n"
        f"QUESTION: {query}\n\n"
        "REPONSE (fondee uniquement sur le contexte, avec citations) :"
    )
    return system, user


# --------------------------------------------------------------------------
# Generation
# --------------------------------------------------------------------------
def _generate_hf(system, user):
    if not config.HF_API_TOKEN:
        raise RuntimeError(
            "Token HF manquant. Ajoutez le secret HF_API_TOKEN (ou HF_TOKEN) "
            "dans les parametres du Space."
        )
    headers = {"Authorization": f"Bearer {config.HF_API_TOKEN}"}
    url = "https://router.huggingface.co/v1/chat/completions"
    payload = {
        "model": config.HF_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "temperature": config.TEMPERATURE,
        "max_tokens": config.MAX_TOKENS,
    }
    r = requests.post(url, headers=headers, json=payload, timeout=180)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"].strip()


def _generate_ollama(system, user):
    payload = {
        "model": config.OLLAMA_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "stream": False,
        "options": {"temperature": config.TEMPERATURE,
                    "num_predict": config.MAX_TOKENS},
    }
    r = requests.post(f"{config.OLLAMA_HOST}/api/chat", json=payload, timeout=180)
    r.raise_for_status()
    return r.json()["message"]["content"].strip()


def generate(system, user):
    if config.LLM_BACKEND == "ollama":
        return _generate_ollama(system, user)
    return _generate_hf(system, user)


# --------------------------------------------------------------------------
# Point d'entree / entry point
# --------------------------------------------------------------------------
def answer(query, lang="fr", extra_context=""):
    query = gdpr.safe_for_processing(query)
    passages = retrieve(query)
    if extra_context:
        passages = passages + [{
            "title": "Document televerse (session)",
            "url": "local://document",
            "text": extra_context[:4000],
            "score": 1.0,
            "type": "Document",
            "status": "publish",
        }]
    if not passages:
        _no = {
            "fr": ("Je n'ai trouve aucune information correspondante dans le "
                   "contenu indexe du site. Reformulez votre question."),
            "en": ("I could not find matching information in the indexed site "
                   "content. Try rephrasing your question."),
            "ht": ("Mwen pa jwenn okenn enfomasyon ki koresponn nan kontni sit "
                   "la. Eseye poze kesyon an yon lot jan."),
        }
        return _no.get(lang, _no["fr"]), []
    system, user = build_prompt(query, passages, lang)
    response = generate(system, user)
    sources = [{"title": p["title"], "url": p["url"],
                "score": p.get("score", 0),
                "status": p.get("status", "publish")}
               for p in passages if p["url"] != "local://document"]
    gdpr.log_interaction(query, response, sources)
    return response, sources