evo_inference.py

"""
evo_inference.py
Step 8: Evo-ready synthesis with plugin support + safe fallback.

(Objective)
- Try to import your real Evo plugin: `evo_plugin.load_model()`.
- If not present, fall back to `evo_plugin_example.load_model()`.
- Provide `synthesize_with_evo(...)` that builds a grounded prompt from retrieved hits
  and either:
    a) calls the generator (generative mode), or
    b) returns a clean extractive answer (safe mode).
"""

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

# Try to load your real Evo plugin first; else example plugin
_GENERATOR = None
try:
    from evo_plugin import load_model as _load_real  # your 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  # fallback to extractive-only

MAX_SNIPPET_CHARS = 400  # (Objective) keep context concise


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


def _build_grounded_prompt(question: str, lang: str, hits: List[Dict]) -> str:
    """
    (Objective) Construct a compact prompt that includes:
    - role + response style instruction (language-aware),
    - the user's question,
    - the top retrieved chunks as "Context #i".
    """
    lang = normalize_lang(lang)
    if lang == "fr":
        system = (
            "Tu es le Copilote Gouvernemental de Maurice. Réponds dans la langue demandée, "
            "clairement et étape par étape, en te basant STRICTEMENT sur le contexte. "
            "Inclure: documents requis, frais, où postuler, délais. Si une info manque, dis-le."
        )
    elif lang == "mfe":
        system = (
            "To enn Copilot Gouv Moris. Reponn dan langaz itilizater, kler ek pas-a-pas, "
            "bas lor KI SUIVAN. Met: ki dokiman bizin, fre, kot pou al, delai. "
            "Si pa ase info, dir li."
        )
    else:
        system = (
            "You are the Mauritius Government Copilot. Answer in the user's language, "
            "clearly and step-by-step, using ONLY the provided context. Include: required documents, "
            "fees, where to apply, processing time. If something is missing, say so."
        )

    ctx_lines = []
    for i, h in enumerate(hits[:6], 1):
        ctx_lines.append(f"[Context #{i}] { _snippet(h['text']) }")

    ctx_block = "\n".join(ctx_lines) if ctx_lines else "[Context] (none)"
    prompt = (
        f"{system}\n\n"
        f"[Question]\n{question}\n\n"
        f"{ctx_block}\n\n"
        f"[Instructions]\n"
        f"- Be concise (6–10 lines).\n"
        f"- Use bullet steps.\n"
        f"- Do not invent links or fees; mention if unknown.\n"
        f"- Answer in language code: {lang}.\n"
        f"[Answer]\n"
    )
    return prompt


def _extractive_answer(user_query: str, lang: str, hits: List[Dict]) -> str:
    """
    (Objective) The safe fallback: bullet points from hits + standard steps.
    """
    if not hits:
        return L(lang, "intro_err")

    bullets = []
    for h in hits[:4]:
        bullets.append(f"- {_snippet(h['text'])}")

    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)]

    md = (
        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)
    )
    return md


def synthesize_with_evo(
    user_query: str,
    lang: str,
    hits: List[Dict],
    mode: str = "extractive",          # "extractive" | "generative"
    max_new_tokens: int = 192,
    temperature: float = 0.4,
) -> str:
    """
    (Objective)
    - If mode == 'generative' and a generator is available, build a grounded prompt and generate.
    - Else, return the safe extractive answer.
    """
    lang = normalize_lang(lang)

    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=prompt,
            max_new_tokens=int(max_new_tokens),
            temperature=float(temperature),
        )
        # In case the generator echos or misses structure, still return something readable
        text = text.strip()
        if not text:
            return _extractive_answer(user_query, lang, hits)
        return text
    except Exception:
        # Any runtime issue falls back to safe mode
        return _extractive_answer(user_query, lang, hits)