evo_inference.py

"""
evo_inference.py — Step 8 (refined)
Adds a GENERATIVE path using a small plugin (FLAN-T5 stand-in) while keeping the
old EXTRACTIVE fallback (bullet points) if generation isn't available.

What's new in this refinement:
- Answers are explicitly labeled **[Generative]** or **[Extractive]** so you
  can tell which path ran at a glance.

How it works:
- We try to import your real evo plugin (evo_plugin.py). If not found, we load
  evo_plugin_example.py instead. If both fail, we stay in extractive mode.
- synthesize_with_evo(...) accepts mode/temp/max_tokens from the UI.
"""

from typing import List, Dict
from utils_lang import L, normalize_lang

# Try to load your real Evo plugin first; else use the example; else None.
_GENERATOR = None
try:
    from evo_plugin import load_model as _load_real   # <- your future file (optional)
    _GENERATOR = _load_real()
except Exception:
    try:
        from evo_plugin_example import load_model as _load_example
        _GENERATOR = _load_example()
    except Exception:
        _GENERATOR = None  # no generator available

MAX_SNIPPET_CHARS = 400


def _snippet(text: str) -> str:
    text = " ".join(text.split())
    return text[:MAX_SNIPPET_CHARS] + ("..." if len(text) > MAX_SNIPPET_CHARS else "")


def _extractive_answer(user_query: str, lang: str, hits: List[Dict]) -> str:
    """Old safe mode: show top snippets + standard steps, now labeled."""
    if not hits:
        return "**[Extractive]**\n\n" + L(lang, "intro_err")

    bullets = [f"- {_snippet(h['text'])}" for h in hits[:4]]
    steps = {
        "en": [
            "• Step 1: Check eligibility & gather required documents.",
            "• Step 2: Confirm fees & payment options.",
            "• Step 3: Apply online or at the indicated office.",
            "• Step 4: Keep reference/receipt; track processing time.",
        ],
        "fr": [
            "• Étape 1 : Vérifiez l’éligibilité et rassemblez les documents requis.",
            "• Étape 2 : Confirmez les frais et les moyens de paiement.",
            "• Étape 3 : Déposez la demande en ligne ou au bureau indiqué.",
            "• Étape 4 : Conservez le reçu/la référence et suivez le délai de traitement.",
        ],
        "mfe": [
            "• Step 1: Get dokiman neseser ek verifie si to elegib.",
            "• Step 2: Konfirm fre ek manyer peyman.",
            "• Step 3: Fer demand online ouswa dan biro ki indike.",
            "• Step 4: Gard referans/reso; swiv letan tretman.",
        ],
    }[normalize_lang(lang)]

    return (
        "**[Extractive]**\n\n"
        f"**{L(lang, 'intro_ok')}**\n\n"
        f"**Q:** {user_query}\n\n"
        f"**Key information:**\n" + "\n".join(bullets) + "\n\n"
        f"**Suggested steps:**\n" + "\n".join(steps)
    )


def _build_grounded_prompt(question: str, lang: str, hits: List[Dict]) -> str:
    """Create a compact prompt that includes the question + top retrieved snippets."""
    lang = normalize_lang(lang)
    if lang == "fr":
        system = ("Tu es le Copilote Gouvernemental de Maurice. Réponds clairement, étape "
                  "par étape, en te basant UNIQUEMENT sur le contexte. Inclure: documents requis, "
                  "frais, où postuler, délais. Dire si une info manque.")
    elif lang == "mfe":
        system = ("To enn Copilot Gouv Moris. Reponn kler ek pas-a-pas, servi zis konteks ki donn. "
                  "Met: ki dokiman bizin, fre, kot pou al, delai. Dir si info manke.")
    else:
        system = ("You are the Mauritius Government Copilot. Answer clearly and step-by-step using "
                  "ONLY the provided context. Include: required documents, fees, where to apply, "
                  "processing time. State if anything is missing.")

    ctx = "\n".join([f"[Context #{i+1}] {_snippet(h['text'])}" for i, h in enumerate(hits[:6])]) or "[Context] (none)"
    return (
        f"{system}\n\n[Question]\n{question}\n\n{ctx}\n\n"
        f"[Instructions]\n- Be concise (6–10 lines)\n- Use bullet steps\n"
        f"- Do not invent links/fees\n- Answer in language code: {lang}\n[Answer]\n"
    )


def synthesize_with_evo(
    user_query: str,
    lang: str,
    hits: List[Dict],
    mode: str = "extractive",
    max_new_tokens: int = 192,
    temperature: float = 0.4,
) -> str:
    """
    If mode=='generative' and a generator exists, generate a grounded answer
    and label it **[Generative]**. Otherwise, return the labeled extractive fallback.
    """
    lang = normalize_lang(lang)

    # No retrieved context? Stay safe.
    if not hits:
        return _extractive_answer(user_query, lang, hits)

    if mode != "generative" or _GENERATOR is None:
        return _extractive_answer(user_query, lang, hits)

    prompt = _build_grounded_prompt(user_query, lang, hits)
    try:
        text = _GENERATOR.generate(
            prompt,
            max_new_tokens=int(max_new_tokens),
            temperature=float(temperature),
        ).strip()

        if not text:
            return _extractive_answer(user_query, lang, hits)

        return "**[Generative]**\n\n" + text

    except Exception:
        # Any runtime issue falls back to safe mode
        return _extractive_answer(user_query, lang, hits)