backend/rag_engine.py

from __future__ import annotations

import os
from collections.abc import Iterable

from dotenv import load_dotenv
from pydantic_ai import Agent
from pydantic_ai.models.groq import GroqModel

from backend.document_store import store


load_dotenv()

SYSTEM_PROMPT = (
    "Tu es Hianatra, un assistant pédagogique. "
    "Réponds uniquement à partir du contexte fourni. "
    "Si la réponse n'est pas dans le contexte, dis-le clairement."
)

FREE_GROQ_MODELS = [
    "llama-3.1-8b-instant",
    "llama-3.3-70b-versatile",
    "qwen/qwen3-32b",
]


def _get_groq_model_name() -> str:
    return os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")


def _dedupe(values: Iterable[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for value in values:
        if value in seen:
            continue
        seen.add(value)
        out.append(value)
    return out


def get_available_models() -> list[str]:
    configured = os.getenv("GROQ_FREE_MODELS", "")
    configured_list = [m.strip() for m in configured.split(",") if m.strip()]
    return _dedupe([_get_groq_model_name(), *configured_list, *FREE_GROQ_MODELS])


def get_default_model() -> str:
    return _get_groq_model_name()


def _build_model_settings(
    temperature: float | None,
    top_p: float | None,
    max_tokens: int | None,
) -> dict[str, float | int]:
    settings: dict[str, float | int] = {}
    if temperature is not None:
        settings["temperature"] = float(max(0.0, min(2.0, temperature)))
    if top_p is not None:
        settings["top_p"] = float(max(0.0, min(1.0, top_p)))
    if max_tokens is not None:
        settings["max_tokens"] = int(max(64, min(4096, max_tokens)))
    return settings


def _build_agent(model_name: str) -> Agent:
    return Agent(
        GroqModel(model_name),
        system_prompt=SYSTEM_PROMPT,
    )

agent = _build_agent(_get_groq_model_name())


def _extract_agent_text(result: object) -> str:
    """Compat pydantic-ai: .data (ancien) -> .output (récent)."""
    for attr in ("output", "data", "result"):
        if hasattr(result, attr):
            value = getattr(result, attr)
            if value is None:
                continue
            return value if isinstance(value, str) else str(value)
    return str(result)


async def ask_with_rag(
    question: str,
    k: int = 4,
    model_name: str | None = None,
    temperature: float | None = 0.2,
    top_p: float | None = 0.95,
    max_tokens: int | None = 1024,
) -> tuple[str, list[str], str]:
    relevant_chunks = store.search(question, k=k)

    if not relevant_chunks:
        return (
            "Aucun document chargé. Veuillez uploader un document d'abord.",
            [],
            model_name or _get_groq_model_name(),
        )

    context = "\n\n---\n\n".join(relevant_chunks)
    prompt = (
        f"Contexte du document :\n{context}\n\n"
        f"Question : {question}\n\n"
        "Réponds de manière claire et précise en te basant sur le contexte ci-dessus."
    )

    selected_model = (model_name or _get_groq_model_name()).strip()
    selected_agent = agent if selected_model == _get_groq_model_name() else _build_agent(selected_model)
    settings = _build_model_settings(temperature, top_p, max_tokens)

    if settings:
        try:
            result = await selected_agent.run(prompt, model_settings=settings)
            return _extract_agent_text(result), relevant_chunks, selected_model
        except TypeError:
            # Compat anciennes versions de pydantic-ai
            pass

    result = await selected_agent.run(prompt)
    return _extract_agent_text(result), relevant_chunks, selected_model