src/rag_core.py

# src/rag_core.py
from __future__ import annotations

from typing import Dict, Any, List
import json
import os

from src import list as list_mode
from src import fulltext as fulltext_mode
from src import synthesis as synthesis_mode
from src import qa as qa_mode

from src.config import (
    CHUNKS_PATH,
    LIST_TRIGGERS,
    REFUSAL,
    SYNTHESIS_REFUSAL,
    QA_WARNING,
    QA_TOP_K_FINAL,
    QA_DOC_MAX_CHARS,
)
from src.utils import (
    normalize_article_id,
    extract_article_id,
    is_list_request,
    is_fulltext_request,
    is_synthesis_request,
)
from src.resources import get_vectorstore, groq_max_tokens_for


# ====================
# SUMMARY_AI (Groq-only, rapide)
# ====================

SUMMARY_TRIGGERS = [
    "résumé ia", "resume ia",
    "résume ia", "resume-ia",
    "summary ia", "ai summary",
]

SUMMARY_WARNING = (
    "Résumé IA : reformulation automatique (peut contenir des erreurs ou omissions). "
    "Vérifie toujours sur le texte officiel."
)


def is_summary_request(q: str) -> bool:
    ql = (q or "").lower()
    return any(t in ql for t in SUMMARY_TRIGGERS)


def _build_summary_context_from_extractive(article_id: str, full_text: str) -> str:
    """
    Construit un contexte court à partir de la synthèse extractive existante.
    On récupère 3–4 segments "- ..." pour alimenter le LLM avec très peu de texte.
    """
    extract = synthesis_mode.extractive_summary(article_id, full_text)

    lines: List[str] = []
    for line in extract.splitlines():
        line = line.strip()
        if line.startswith("- "):
            seg = line[2:].strip()
            if seg:
                lines.append(seg)

    lines = lines[:4]  # limite dure

    if not lines:
        # fallback ultra sûr
        return qa_mode.truncate_text(full_text, 400)

    return "\n".join(f"- {l}" for l in lines)


def _summary_ai(article_id: str) -> Dict[str, Any]:
    article_id = normalize_article_id(article_id)
    text = load_article_text(article_id)

    if not text:
        return {
            "mode": "SUMMARY_AI",
            "answer": f"Article {article_id} introuvable.",
            "articles": [],
        }

    # Contexte réduit (extraits) pour accélérer
    context = _build_summary_context_from_extractive(article_id, text)

    # Prompt strict FR + puces (défini dans qa.py)
    prompt = qa_mode.build_summary_prompt(article_id, context)

    # Paramètres Groq (via env vars)
    cfg = qa_mode.QAConfig(
        qa_top_k_final=1,
        qa_doc_max_chars=600,
        qa_max_tokens=groq_max_tokens_for("summary"),
        qa_temperature=float(os.environ.get("GROQ_TEMPERATURE", "0.1")),
    )

    ans = qa_mode.llm_generate_qa(prompt, cfg=cfg).strip()

    return {
        "mode": "SUMMARY_AI",
        "answer": f"{SUMMARY_WARNING}\n\n{ans}",
        "articles": [article_id],
    }


# ====================
# CHARGEMENT CORPUS (UNE FOIS)
# ====================

_ARTICLES: Dict[str, str] | None = None


def get_all_articles() -> Dict[str, str]:
    global _ARTICLES
    if _ARTICLES is None:
        articles: Dict[str, str] = {}
        with open(CHUNKS_PATH, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                obj = json.loads(line)
                aid = normalize_article_id(obj.get("article_id", ""))
                text = obj.get("text") or obj.get("page_content") or ""
                if aid and text:
                    articles[aid] = text
        _ARTICLES = articles
    return _ARTICLES


# ====================
# FULLTEXT
# ====================

def load_article_text(article_id: str) -> str | None:
    return fulltext_mode.load_article_text(
        normalize_article_id(article_id),
        CHUNKS_PATH,
    )


# ====================
# QA (Groq-only pour la génération)
# ====================

def _qa_answer(question: str) -> Dict[str, Any]:
    # Retrieval vectoriel (FAISS)
    vs = get_vectorstore()
    docs = vs.similarity_search(question, k=max(1, QA_TOP_K_FINAL))

    sources: List[str] = []
    for d in docs:
        aid = normalize_article_id(d.metadata.get("article_id", ""))
        if aid and aid not in sources:
            sources.append(aid)

    ctx_parts: List[str] = []
    for d in docs[:max(1, QA_TOP_K_FINAL)]:
        aid = normalize_article_id(d.metadata.get("article_id", "UNKNOWN"))
        txt = qa_mode.truncate_text(d.page_content or "", QA_DOC_MAX_CHARS)
        ctx_parts.append(f"[{aid}]\n{txt}")

    prompt = qa_mode.build_qa_prompt_fast(question, "\n\n".join(ctx_parts), sources)

    cfg = qa_mode.QAConfig(
        qa_top_k_final=QA_TOP_K_FINAL,
        qa_doc_max_chars=QA_DOC_MAX_CHARS,
        qa_max_tokens=groq_max_tokens_for("qa"),
        qa_temperature=float(os.environ.get("GROQ_TEMPERATURE", "0.1")),
    )

    ans = qa_mode.llm_generate_qa(prompt, cfg=cfg).strip()

    return {
        "mode": "QA",
        "answer": f"{QA_WARNING}\n\n{ans}",
        "articles": sources,
    }


# ====================
# ROUTEUR
# ====================

def _looks_like_question(q: str) -> bool:
    """
    Détecte une intention de question, même si la requête est courte.
    C'est crucial pour éviter que des questions tombent dans LIST par défaut.
    """
    ql = (q or "").strip().lower()
    if "?" in ql:
        return True
    starters = (
        "que ", "qu'", "quoi", "comment", "pourquoi", "quand", "où",
        "est-ce", "peux", "peut", "dois", "doit", "faut", "faudrait",
        "quelle", "quelles", "quel", "quels",
    )
    return ql.startswith(starters)


def answer_query(q: str) -> Dict[str, Any]:
    q = (q or "").strip()
    if not q:
        return {"mode": "QA", "answer": REFUSAL, "articles": []}

    article_id = extract_article_id(q)

    # 1) FULLTEXT
    if article_id and is_fulltext_request(q):
        article_id = normalize_article_id(article_id)
        text = load_article_text(article_id)
        return {
            "mode": "FULLTEXT",
            "answer": text or REFUSAL,
            "articles": [article_id],
        }

    # 2) SUMMARY_AI (Résumé IA)
    if article_id and is_summary_request(q):
        return _summary_ai(article_id)

    # 3) SYNTHESIS (extractif fiable)
    if is_synthesis_request(q):
        if not article_id:
            return {"mode": "SYNTHESIS", "answer": SYNTHESIS_REFUSAL, "articles": []}

        article_id = normalize_article_id(article_id)
        text = load_article_text(article_id)
        if not text:
            return {
                "mode": "SYNTHESIS",
                "answer": f"Article {article_id} introuvable.",
                "articles": [],
            }

        return {
            "mode": "SYNTHESIS",
            "answer": synthesis_mode.extractive_summary(article_id, text),
            "articles": [article_id],
        }

    # 4) LIST explicite
    if is_list_request(q):
        return list_mode.list_articles(
            q,
            articles=get_all_articles(),
            vs=None,  # important : LIST doit rester léger/explicable
            normalize_article_id=normalize_article_id,
            list_triggers=LIST_TRIGGERS,
            cfg=list_mode.ListConfig(),
        )

    # 5) Routage robuste : si c'est une QUESTION, on force QA
    if _looks_like_question(q):
        return _qa_answer(q)

    # 6) Si un article est mentionné et que ce n'est pas un mode dédié,
    # on privilégie QA (cas : "Que dit l'article D521-5" sans forcément de "?")
    if article_id:
        return _qa_answer(q)

    # 7) LIST par défaut si requête courte (mots-clés)
    if len(q.split()) <= 5:
        return list_mode.list_articles(
            q,
            articles=get_all_articles(),
            vs=None,
            normalize_article_id=normalize_article_id,
            list_triggers=LIST_TRIGGERS,
            cfg=list_mode.ListConfig(),
        )

    # 8) QA par défaut
    return _qa_answer(q)