app.py

"""
Hindi Banking Voice Assistant
CPU Basic safe — uses HF InferenceClient with providers that actually work.

STT  : openai/whisper-large-v3        (hf-inference, ASR)
LLM  : mistralai/Mistral-7B-Instruct-v0.3  (featherless provider, free tier)
TTS  : facebook/mms-tts-hin           (hf-inference, Hindi TTS, small model)
"""

import os
import io
import numpy as np
import soundfile as sf
import gradio as gr
from huggingface_hub import InferenceClient

from rag_pipeline import (
    normalize_jargon,
    translate_to_retrieval_query,
    build_rag_prompt,
    format_response_for_tts,
    get_retriever,
)

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", "")

STT_MODEL = "openai/whisper-large-v3"
LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
TTS_MODEL = "facebook/mms-tts-hin"

retriever = get_retriever()
print(f"Knowledge base ready: {len(retriever.doc_ids)} documents")


# ─────────────────────────────────────────────
# INFERENCE FUNCTIONS
# ─────────────────────────────────────────────

def get_client(provider: str) -> InferenceClient:
    return InferenceClient(provider=provider, api_key=HF_TOKEN)


def stt_whisper(audio_array: np.ndarray, sample_rate: int) -> str:
    """Convert Hindi audio to text using Whisper via hf-inference."""
    buf = io.BytesIO()
    sf.write(buf, audio_array, sample_rate, format="WAV", subtype="PCM_16")
    buf.seek(0)

    client = get_client("hf-inference")
    result = client.automatic_speech_recognition(
        audio=buf.read(),
        model=STT_MODEL,
    )
    # result is an ASROutput object with .text
    return result.text.strip() if hasattr(result, "text") else str(result).strip()


def llm_generate(prompt: str) -> str:
    """Generate Hindi answer using Mistral via featherless (free tier)."""
    client = get_client("featherless")

    # Build messages in chat format
    messages = [
        {
            "role": "system",
            "content": (
                "आप एक सहायक बैंकिंग सहायक हैं। केवल दी गई जानकारी के आधार पर "
                "सरल हिंदी में 3-4 वाक्यों में उत्तर दें। "
                "यदि जानकारी नहीं है तो कहें: 'यह जानकारी मेरे पास नहीं है। "
                "कृपया अपने बैंक से संपर्क करें।' हमेशा हिंदी में उत्तर दें।"
            ),
        },
        {"role": "user", "content": prompt},
    ]

    response = client.chat.completions.create(
        model=LLM_MODEL,
        messages=messages,
        max_tokens=350,
        temperature=0.7,
    )
    return response.choices[0].message.content.strip()


def tts_hindi(text: str):
    """Convert Hindi text to speech using MMS-TTS-HIN via hf-inference."""
    client = get_client("hf-inference")
    try:
        audio_bytes = client.text_to_speech(text=text, model=TTS_MODEL)
        buf = io.BytesIO(audio_bytes)
        audio_array, sample_rate = sf.read(buf)
        return int(sample_rate), audio_array.astype(np.float32)
    except Exception as e:
        print(f"TTS error: {e}")
        return None


# ─────────────────────────────────────────────
# PIPELINE GENERATORS
# ─────────────────────────────────────────────

def run_voice_pipeline(audio_input):
    """Voice mode: mic audio → transcript, answer, citations, audio."""
    if audio_input is None:
        yield "", "कृपया माइक्रोफोन बटन दबाकर अपना प्रश्न पूछें।", "", None
        return
    if not HF_TOKEN:
        yield "", "⚠️ HF_TOKEN Secret नहीं मिला। Space Settings → Secrets में जोड़ें।", "", None
        return

    try:
        sample_rate, audio_array = audio_input
        audio_float = audio_array.astype(np.float32)
        if np.abs(audio_float).max() > 1.0:
            audio_float /= 32768.0
        if audio_float.ndim > 1:
            audio_float = audio_float.mean(axis=1)

        yield "⏳ आवाज़ पहचाना जा रहा है...", "", "", None

        hindi_text = stt_whisper(audio_float, sample_rate)
        if not hindi_text:
            yield "", "आवाज़ स्पष्ट नहीं सुनाई दी। कृपया दोबारा कोशिश करें।", "", None
            return

        yield hindi_text, "⏳ जानकारी खोजी जा रही है...", "", None

        normalized = normalize_jargon(hindi_text)
        eng_query  = translate_to_retrieval_query(normalized)
        docs       = retriever.retrieve(eng_query, top_k=3) or retriever.retrieve(hindi_text, top_k=2)
        citations  = "\n".join(f"• {d['title']}" for d in docs) if docs else "कोई स्रोत नहीं मिला।"

        yield hindi_text, "⏳ उत्तर तैयार किया जा रहा है...", citations, None

        # Build a concise context-augmented prompt for the chat model
        context = "\n\n".join(
            f"[{i+1}] {d['title']}\n{d['content'][:500]}"
            for i, d in enumerate(docs)
        ) if docs else "कोई प्रासंगिक जानकारी नहीं मिली।"

        full_prompt = (
            f"नीचे दी गई जानकारी के आधार पर प्रश्न का उत्तर दें:\n\n"
            f"{context}\n\n"
            f"प्रश्न: {hindi_text}"
        )

        answer = llm_generate(full_prompt)
        clean_answer = format_response_for_tts(answer) or \
            "यह जानकारी मेरे पास नहीं है। कृपया अपने बैंक से संपर्क करें।"

        yield hindi_text, clean_answer, citations, None

        yield hindi_text, clean_answer, "⏳ ऑडियो बनाया जा रहा है...", None
        audio_out = tts_hindi(clean_answer)
        yield hindi_text, clean_answer, citations, audio_out

    except Exception as e:
        import traceback; traceback.print_exc()
        yield "", f"⚠️ त्रुटि: {e}", "", None


def run_text_pipeline(text_input: str):
    """Text mode: typed question → answer, citations, audio."""
    if not text_input or not text_input.strip():
        yield text_input, "कृपया एक प्रश्न लिखें।", "", None
        return
    if not HF_TOKEN:
        yield text_input, "⚠️ HF_TOKEN Secret नहीं मिला।", "", None
        return

    try:
        yield text_input, "⏳ जानकारी खोजी जा रही है...", "", None

        normalized = normalize_jargon(text_input)
        eng_query  = translate_to_retrieval_query(normalized)
        docs       = retriever.retrieve(eng_query, top_k=3) or retriever.retrieve(text_input, top_k=2)
        citations  = "\n".join(f"• {d['title']}" for d in docs) if docs else "कोई स्रोत नहीं मिला।"

        yield text_input, "⏳ उत्तर तैयार किया जा रहा है...", citations, None

        context = "\n\n".join(
            f"[{i+1}] {d['title']}\n{d['content'][:500]}"
            for i, d in enumerate(docs)
        ) if docs else "कोई प्रासंगिक जानकारी नहीं मिली।"

        full_prompt = (
            f"नीचे दी गई जानकारी के आधार पर प्रश्न का उत्तर दें:\n\n"
            f"{context}\n\n"
            f"प्रश्न: {text_input}"
        )

        answer = llm_generate(full_prompt)
        clean_answer = format_response_for_tts(answer) or \
            "यह जानकारी मेरे पास नहीं है। कृपया अपने बैंक से संपर्क करें।"

        yield text_input, clean_answer, citations, None

        yield text_input, clean_answer, "⏳ ऑडियो बनाया जा रहा है...", None
        audio_out = tts_hindi(clean_answer)
        yield text_input, clean_answer, citations, audio_out

    except Exception as e:
        import traceback; traceback.print_exc()
        yield text_input, f"⚠️ त्रुटि: {e}", "", None


# ─────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────

EXAMPLES = [
    "मुझे होम लोन के लिए कौन से दस्तावेज़ चाहिए?",
    "EMI क्या होती है और कैसे calculate होती है?",
    "मुद्रा लोन कैसे मिलेगा?",
    "CIBIL score क्या होता है?",
    "PM Awas Yojana में subsidy कैसे मिलती है?",
    "जन धन खाता कैसे खोलें?",
    "किसान क्रेडिट कार्ड क्या है?",
    "Atal Pension Yojana में कैसे जुड़ें?",
    "SBI personal loan की interest rate क्या है?",
    "महिला स्वयं सहायता समूह को loan कैसे मिलेगा?",
    "Street vendor ko SVANidhi loan kaise milega?",
    "Bank ke khilaf complaint kahan karein?",
]

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Noto+Sans+Devanagari:wght@300;400;500;700&family=Sora:wght@300;400;600;700&display=swap');

:root {
    --saffron: #FF6B00;
    --saffron-glow: rgba(255,107,0,0.15);
    --green: #0A7A3E;
    --navy: #0F1E3D;
    --navy-mid: #162848;
    --gold: #D4A017;
    --gold-light: #F0C042;
    --text-primary: #F4EFE6;
    --text-secondary: #B8C4D8;
    --border: rgba(212,160,23,0.3);
}
* { box-sizing: border-box; }
body, .gradio-container {
    font-family: 'Sora','Noto Sans Devanagari',sans-serif !important;
    background: var(--navy) !important;
    color: var(--text-primary) !important;
}
.gradio-container { max-width: 1100px !important; margin: 0 auto !important; }

.tricolor { height:4px; background:linear-gradient(90deg,#FF6B00 33%,white 33% 66%,#0A7A3E 66%); }

.app-header {
    background: linear-gradient(135deg,var(--navy) 0%,var(--navy-mid) 60%,#0D2847 100%);
    border-bottom: 2px solid var(--gold);
    padding: 28px 40px 22px; position: relative; overflow: hidden;
}
.app-header::before {
    content:''; position:absolute; top:-60%; right:-5%;
    width:350px; height:350px;
    background:radial-gradient(circle,var(--saffron-glow) 0%,transparent 70%);
    pointer-events:none;
}
.app-header::after { content:'🏦'; position:absolute; right:36px; top:50%; transform:translateY(-50%); font-size:72px; opacity:0.09; }
.header-badge { display:inline-flex; align-items:center; gap:8px; background:rgba(212,160,23,0.1); border:1px solid var(--gold); border-radius:100px; padding:4px 14px; font-size:11px; font-weight:600; letter-spacing:1.5px; text-transform:uppercase; color:var(--gold-light); margin-bottom:10px; }
.header-title { font-size:28px; font-weight:700; color:var(--text-primary); margin-bottom:5px; font-family:'Noto Sans Devanagari','Sora',sans-serif !important; }
.header-subtitle { font-size:13px; color:var(--text-secondary); font-weight:300; }

.notice { background:rgba(212,160,23,0.07); border:1px solid rgba(212,160,23,0.28); border-radius:10px; padding:10px 18px; font-size:12px; color:var(--gold-light); margin:12px 40px; display:flex; align-items:flex-start; gap:10px; line-height:1.6; }

.steps { display:flex; gap:6px; flex-wrap:wrap; padding:12px 40px; background:rgba(10,25,50,0.55); border-bottom:1px solid var(--border); }
.step { display:flex; align-items:center; gap:5px; background:rgba(255,255,255,0.03); border:1px solid var(--border); border-radius:100px; padding:4px 11px; font-size:11px; color:var(--text-secondary); }
.sn { width:16px; height:16px; background:var(--saffron); border-radius:50%; display:flex; align-items:center; justify-content:center; font-size:9px; font-weight:700; color:white; flex-shrink:0; }

label { color:var(--gold-light) !important; font-size:11px !important; font-weight:600 !important; letter-spacing:0.8px !important; text-transform:uppercase !important; }

textarea { background:rgba(10,20,45,0.85) !important; border:1px solid var(--border) !important; border-radius:10px !important; color:var(--text-primary) !important; font-family:'Noto Sans Devanagari','Sora',sans-serif !important; font-size:14px !important; line-height:1.7 !important; padding:13px !important; }
textarea:focus { border-color:var(--saffron) !important; outline:none !important; box-shadow:0 0 0 3px var(--saffron-glow) !important; }

.answer-area textarea { background:linear-gradient(135deg,rgba(10,122,62,0.07) 0%,rgba(10,20,45,0.9) 100%) !important; border:1px solid rgba(10,122,62,0.3) !important; color:#DFF0E6 !important; font-size:15px !important; line-height:1.8 !important; min-height:130px !important; }
.sources-area textarea { background:rgba(212,160,23,0.04) !important; border:1px solid rgba(212,160,23,0.18) !important; font-size:12px !important; color:var(--text-secondary) !important; }

.panel { background:var(--navy-mid) !important; border:1px solid var(--border) !important; border-radius:14px !important; padding:20px !important; margin:8px !important; }

button.primary { background:linear-gradient(135deg,var(--saffron) 0%,#D55000 100%) !important; border:none !important; border-radius:10px !important; color:white !important; font-weight:700 !important; font-size:13px !important; padding:12px 28px !important; cursor:pointer !important; transition:all 0.25s !important; box-shadow:0 4px 20px rgba(255,107,0,0.35) !important; letter-spacing:0.5px !important; text-transform:uppercase !important; }
button.primary:hover { transform:translateY(-2px) !important; box-shadow:0 8px 30px rgba(255,107,0,0.5) !important; }

.tab-nav button { background:transparent !important; border:none !important; border-bottom:2px solid transparent !important; color:var(--text-secondary) !important; font-size:13px !important; font-weight:500 !important; padding:10px 22px !important; border-radius:0 !important; transition:all 0.2s !important; }
.tab-nav button.selected { border-bottom-color:var(--saffron) !important; color:#FF8C33 !important; font-weight:600 !important; }

.footer { padding:16px 40px; text-align:center; font-size:11px; color:rgba(184,196,216,0.4); border-top:1px solid var(--border); background:rgba(8,15,35,0.5); }
::-webkit-scrollbar { width:5px; } ::-webkit-scrollbar-thumb { background:var(--gold); border-radius:3px; }
"""


def build_ui():
    with gr.Blocks(css=CSS, title="बैंकिंग सहायक", theme=gr.themes.Base()) as demo:

        gr.HTML('<div class="tricolor"></div>')
        gr.HTML("""
        <div class="app-header">
            <div class="header-badge">🎙️ Voice · Hindi · Free Tier</div>
            <div class="header-title">बैंकिंग सहायक</div>
            <div class="header-subtitle">AI-powered Hindi Banking & Finance Assistant — Loans · Schemes · EMI · Accounts · Social Security</div>
        </div>
        """)
        gr.HTML("""
        <div class="notice">
            <span>🔑</span>
            <span><b>One-time setup:</b> Space <b>Settings → Secrets</b> → add
            <code style="background:rgba(0,0,0,0.3);padding:1px 6px;border-radius:4px;">HF_TOKEN</code>
            = your free HF token from hf.co/settings/tokens</span>
        </div>
        """)
        gr.HTML("""
        <div class="steps">
            <div class="step"><span class="sn">1</span>हिंदी बोलें</div>
            <div class="step"><span class="sn">2</span>Whisper STT</div>
            <div class="step"><span class="sn">3</span>Jargon Norm</div>
            <div class="step"><span class="sn">4</span>RAG Retrieval</div>
            <div class="step"><span class="sn">5</span>Mistral-7B LLM</div>
            <div class="step"><span class="sn">6</span>MMS Hindi TTS</div>
            <div class="step"><span class="sn">7</span>हिंदी उत्तर</div>
        </div>
        """)

        with gr.Tabs(elem_classes="tab-nav"):

            with gr.Tab("🎙️ Voice Mode"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes="panel"):
                        gr.HTML('<p style="color:#B8C4D8;font-size:13px;margin:0 0 12px;">माइक बटन दबाएं → हिंदी में बोलें → Submit</p>')
                        v_audio_in  = gr.Audio(sources=["microphone"], type="numpy", label="🎤 अपना प्रश्न बोलें")
                        v_btn       = gr.Button("🔍 उत्तर खोजें", variant="primary")
                        v_transcript= gr.Textbox(label="📝 Transcript", interactive=False, lines=2, placeholder="आपकी बात यहाँ दिखेगी...")
                    with gr.Column(scale=1, elem_classes="panel"):
                        v_answer    = gr.Textbox(label="💬 उत्तर", interactive=False, lines=5, placeholder="उत्तर यहाँ आएगा...", elem_classes="answer-area")
                        v_audio_out = gr.Audio(label="🔊 उत्तर सुनें", type="numpy", autoplay=True)
                        v_citations = gr.Textbox(label="📚 स्रोत", interactive=False, lines=3, elem_classes="sources-area")

                v_btn.click(fn=run_voice_pipeline, inputs=[v_audio_in],
                            outputs=[v_transcript, v_answer, v_citations, v_audio_out])

            with gr.Tab("⌨️ Text Mode"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes="panel"):
                        gr.HTML('<p style="color:#B8C4D8;font-size:13px;margin:0 0 12px;">हिंदी या Hinglish में टाइप करें</p>')
                        t_input     = gr.Textbox(label="✏️ अपना प्रश्न लिखें", lines=3, placeholder="जैसे: होम लोन के लिए कौन से दस्तावेज़ चाहिए?")
                        t_btn       = gr.Button("🔍 उत्तर खोजें", variant="primary")
                        t_transcript= gr.Textbox(label="📝 आपका प्रश्न", interactive=False, lines=2)
                    with gr.Column(scale=1, elem_classes="panel"):
                        t_answer    = gr.Textbox(label="💬 उत्तर", interactive=False, lines=5, elem_classes="answer-area")
                        t_audio_out = gr.Audio(label="🔊 उत्तर सुनें", type="numpy", autoplay=True)
                        t_citations = gr.Textbox(label="📚 स्रोत", interactive=False, lines=3, elem_classes="sources-area")

                t_btn.click(fn=run_text_pipeline, inputs=[t_input],
                            outputs=[t_transcript, t_answer, t_citations, t_audio_out])
                t_input.submit(fn=run_text_pipeline, inputs=[t_input],
                               outputs=[t_transcript, t_answer, t_citations, t_audio_out])

        gr.HTML('<div style="padding:6px 10px 0;"><p style="color:var(--gold);font-size:11px;font-weight:600;letter-spacing:1px;text-transform:uppercase;margin:8px 0 4px 32px;">📌 उदाहरण प्रश्न</p></div>')
        gr.Examples(examples=[[q] for q in EXAMPLES], inputs=[t_input], label="", cache_examples=False)

        with gr.Accordion("ℹ️ About / जानकारी", open=False):
            gr.HTML("""
            <div style="padding:18px;color:#B8C4D8;font-size:13px;line-height:1.8;">
              <div style="display:grid;grid-template-columns:1fr 1fr;gap:20px;margin-bottom:14px;">
                <div>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:8px;">🤖 Models</div>
                  <div>• <b>STT:</b> openai/whisper-large-v3 (hf-inference)</div>
                  <div>• <b>LLM:</b> mistralai/Mistral-7B-Instruct-v0.3 (featherless)</div>
                  <div>• <b>TTS:</b> facebook/mms-tts-hin (hf-inference)</div>
                  <br/>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:6px;">⏱️ Response Time</div>
                  <div>• STT: ~5–10s &nbsp;• LLM: ~5–15s &nbsp;• TTS: ~3–8s</div>
                  <div><b>Total: ~15–35s per query</b></div>
                </div>
                <div>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:8px;">📖 Knowledge Base</div>
                  <div>• 12 Major Indian Banks</div>
                  <div>• 11 Govt Credit Schemes (Mudra, PMAY, KCC...)</div>
                  <div>• 7 Social Security Schemes (APY, PMJJBY, EPFO...)</div>
                  <div>• Financial Terms (EMI, CIBIL, KYC, MCLR...)</div>
                  <div>• MFI, NBFC, RRBs, Co-op Banks</div>
                  <div>• Consumer Rights & RBI Ombudsman</div>
                </div>
              </div>
              <div style="background:rgba(255,107,0,0.07);border:1px solid rgba(255,107,0,0.2);border-radius:8px;padding:10px 14px;font-size:12px;color:#FF8C33;">
                ⚠️ यह assistant केवल सामान्य जानकारी देता है। वित्तीय निर्णयों के लिए अपने बैंक या certified advisor से सलाह लें।
              </div>
            </div>
            """)

        gr.HTML('<div class="footer">बैंकिंग सहायक · Whisper + Mistral-7B + MMS-TTS · CPU Basic Safe · HF Inference Providers</div>')

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)