src/llm.py

import os
from google import genai
from google.genai import types

# -------------------------------------------------------------------------
# LLM: Google Gemini Flash
# - Excellent Hindi / Hinglish instruction following
# - Requires GOOGLE_API_KEY in Space secrets
# -------------------------------------------------------------------------
MODEL_NAME = "gemini-3-flash-preview"

_FALLBACK_NO_KEY    = "माफ़ करें, AI सेटअप नहीं हुआ। कृपया बाद में कोशिश करें।"
_FALLBACK_LLM_ERROR = "माफ़ करें, जानकारी लाने में समस्या हुई। कृपया दोबारा पूछें।"

_client = None


def _get_client():
    """Lazily initialises the Gemini client (once per process)."""
    global _client
    if _client is None:
        api_key = os.environ.get("GOOGLE_API_KEY")
        if not api_key:
            print("[LLM] Warning: GOOGLE_API_KEY not set.")
            return None
        _client = genai.Client(api_key=api_key)
        print(f"[LLM] Gemini {MODEL_NAME} ready.")
    return _client


def correct_hindi_query(raw_transcript: str) -> str:
    """
    Uses the LLM to quickly fix phonetic spelling mistakes from the Speech-to-Text engine
    before doing the vector search. For example, 'laan' -> 'loan', 'fiks dipojt' -> 'FD'.
    """
    client = _get_client()
    if not client:
        return raw_transcript

    prompt = (
        "You are an expert at fixing phonetic Hindi speech-to-text spelling mistakes.\n"
        "The user is a rural Indian farmer asking a banking question. The STT engine transcribed "
        "their voice with typos. For example, they might have said 'loan' but it transcribed as 'laan', "
        "or 'fixed deposit' as 'fiks dipojat'.\n\n"
        f"Original Transcript: '{raw_transcript}'\n\n"
        "Task:\n"
        "1. Fix any obvious banking term mispronunciations.\n"
        "2. Keep the query in Devanagari Hindi.\n"
        "3. Do NOT answer the question. Only output the corrected query string.\n"
        "4. If no correction is needed, just output the original string.\n\n"
        "Corrected Query:"
    )

    try:
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=prompt,
            config=types.GenerateContentConfig(
                max_output_tokens=100,
                temperature=0.0,
            ),
        )
        corrected = (response.text or "").strip()
        return corrected if corrected else raw_transcript
    except Exception as exc:
        print(f"[LLM] Query correction error: {exc}")
        return raw_transcript


def generate_hindi_response(hindi_question: str, english_context: str, history: list = None) -> str:
    """
    Generates a warm, conversational Hindi answer using FAQ context.
    Designed to sound like a helpful, friendly bank employee — not a robot.
    """
    client = _get_client()
    if not client:
        return _FALLBACK_NO_KEY

    prompt = (
        "You are 'सहायक', a helpful banking assistant for rural Indian farmers.\n"
        "Explain things in simple, direct Hindi (Devanagari script).\n\n"
        "RULES:\n"
        "1. GROUNDING: Base your answer ONLY on the 'FAQ Information' provided below. Do not invent details, ages, numbers, or rules not in the text.\n"
        "2. CONVERSATIONAL MEMORY: If the user says something conversational like 'Haa', 'Haan', 'Yes', 'No', 'Ok', or a greeting, look at the Previous Conversation History. "
        "They are likely answering your previous follow-up question. If they said yes, answer that previous topic. If the FAQ is empty, guide them back to banking.\n"
        "3. STT TYPOS: The user's input comes from Speech-to-Text. It might have slight typos like 'Haa' instead of 'Haan', or 'klonk' instead of 'loan'. Be smart and infer banking terms phonetically.\n"
        "4. UNKNOWN: If it's a completely new question and the FAQ is empty, say EXACTLY: 'माफ़ करें, मेरे पास इसकी जानकारी नहीं है। कृपया बैंक शाखा से संपर्क करें।'\n"
        "5. FORMAT: Give a clear, complete answer. DO NOT cut off mid-sentence. Write as much as needed to finish your thought.\n"
        "6. FOLLOW-UP: ALWAYS end your entire response with ONE relevant follow-up question on a new line, formatted exactly like this:\n"
        "क्या आप यह भी जानना चाहेंगे: [your question here]?\n\n"
        "SAFETY EXCEPTION: If the user asks about fraud, scams, or lost cards, you may advise them to contact the bank or police immediately.\n\n"
    )

    if history:
        prompt += "--- Previous Conversation History (For Context) ---\n"
        for user_msg, bot_msg, _ in history[-2:]:  # Only keep last 2 turns to save tokens
            prompt += f"User: {user_msg}\nYou: {bot_msg}\n\n"
        prompt += "--- End History ---\n\n"

    prompt += (
        f"--- FAQ Information ---\n{english_context}\n--- End FAQ ---\n\n"
        f"User's question: {hindi_question}\n\n"
        "Your response (in Hindi):"
    )

    try:
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=prompt,
            config=types.GenerateContentConfig(
                max_output_tokens=2048,
                temperature=0.1,
            ),
        )
        answer = (response.text or "").strip()
        return answer if answer else _FALLBACK_LLM_ERROR

    except Exception as exc:
        print(f"[LLM] Generation error: {exc}")
        return _FALLBACK_LLM_ERROR