rag_pipeline.py

"""
RAG Pipeline — pure Python, zero external dependencies.
No torch, no numpy, no transformers. Safe on CPU Basic.
"""

import re
from knowledge_base import get_document_texts, KNOWLEDGE_BASE


# ─────────────────────────────────────────────
# FINANCE JARGON NORMALIZATION MAP
# Hindi/Hinglish → English canonical terms
# ─────────────────────────────────────────────
JARGON_MAP = {
    # EMI variants
    "kist": "EMI equated monthly installment",
    "maahik kist": "monthly installment EMI",
    "maasik bhugtan": "monthly payment EMI",
    "kisten": "installments EMI",
    "maahik bhugtan": "monthly installment EMI",

    # Loan
    "karz": "loan",
    "udhaar": "loan credit",
    "rin": "loan",
    "loan lena": "apply for loan",
    "loan milega": "loan eligibility",
    "paise chahiye": "need money loan",
    "paisa": "money funds",
    "raqam": "amount loan",

    # Interest
    "byaj": "interest rate",
    "sudh": "interest",
    "byaj dar": "interest rate",
    "faixed byaj": "fixed interest rate",
    "badlav wala byaj": "floating interest rate",

    # Bank account
    "khata": "bank account",
    "bachat khata": "savings account",
    "khata kholna": "open bank account",
    "bank mein khata": "bank account",

    # Collateral/Guarantee
    "zamanat": "collateral guarantee security",
    "zamanatdar": "guarantor",
    "girwi": "mortgage pledge",
    "girvi rakhna": "pledge collateral",

    # Documents
    "kaagaz": "documents",
    "dastavej": "documents",
    "pehchaan patra": "identity proof",
    "niwas praman": "address proof",
    "aay praman": "income proof",

    # Credit/CIBIL
    "saakh": "credit score CIBIL",
    "credit score kya hai": "what is credit score CIBIL",
    "score": "CIBIL credit score",

    # Principal/Tenure
    "mool rashi": "principal amount",
    "avadhi": "loan tenure duration",
    "muddat": "loan tenure period",
    "kitne saal": "how many years tenure",
    "kitne mahine": "how many months tenure",

    # Repayment
    "wapasi": "repayment",
    "bhugtan": "payment repayment",
    "chukana": "repay loan",
    "ada karna": "pay repay",

    # Government schemes
    "sarkar ki yojana": "government scheme",
    "yojana": "scheme",
    "sarkari loan": "government loan scheme",
    "subsidy": "subsidy government benefit",
    "anudan": "grant subsidy",

    # Specific schemes
    "mudra": "mudra loan PMMY",
    "kisaan": "farmer kisan",
    "kisan": "farmer kisan credit card",
    "jan dhan": "PMJDY jan dhan account",
    "bima": "insurance",
    "jeevan bima": "life insurance PMJJBY",
    "suraksha bima": "accident insurance PMSBY",
    "pension": "pension APY Atal Pension Yojana",
    "gramin bank": "rural bank RRB",
    "shg": "self help group SHG women loan",
    "samuh": "self help group SHG",
    "mahila samuh": "women self help group SHG microfinance",

    # Defaults/issues
    "default": "loan default NPA",
    "band ho gaya": "account closed loan default",
    "paise nahin de paya": "unable to repay loan default",
    "chhoot": "waiver loan waiver",

    # Property
    "ghar lena": "home purchase home loan",
    "makan": "house home property",
    "zameen": "land property",
    "ghar banana": "home construction loan",
    "flat": "apartment home loan",

    # Grievance
    "shikayat": "complaint grievance",
    "problem": "complaint issue grievance",
    "dhoka": "fraud complaint",
    "pareshan": "problem issue complaint",
}


def normalize_jargon(text: str) -> str:
    """Replace Hindi/Hinglish finance jargon with English equivalents."""
    text_lower = text.lower()
    for hindi_term, english_term in JARGON_MAP.items():
        if hindi_term in text_lower:
            text_lower = text_lower.replace(hindi_term, english_term)
    return text_lower


def translate_to_retrieval_query(normalized_text: str) -> str:
    """Extract English words from normalized text for retrieval."""
    words = [w for w in normalized_text.split() if any(c.isalpha() for c in w)]
    return " ".join(words[:20])


# ─────────────────────────────────────────────
# KEYWORD RETRIEVER — pure Python, no dependencies
# ─────────────────────────────────────────────

class SimpleRetriever:
    def __init__(self):
        self.doc_ids = []
        self.documents = []   # lowercased full text strings
        self.doc_words = []   # sets of words per doc
        self._build_index()

    def _build_index(self):
        for doc_id, text in get_document_texts():
            self.doc_ids.append(doc_id)
            lowered = text.lower()
            self.documents.append(lowered)
            self.doc_words.append(set(re.findall(r'\b\w+\b', lowered)))

        # Augment with tags
        for i, doc in enumerate(KNOWLEDGE_BASE):
            tags_text = " ".join(doc.get("tags", [])).lower()
            self.documents[i] += " " + tags_text
            self.doc_words[i].update(re.findall(r'\b\w+\b', tags_text))

    def retrieve(self, query: str, top_k: int = 3) -> list:
        query_words = set(re.findall(r'\b\w+\b', query.lower()))
        if not query_words:
            return []

        scores = []
        for i, doc_words in enumerate(self.doc_words):
            overlap = len(query_words & doc_words)
            score = overlap / (len(query_words) + 0.5)
            # Bonus for longer exact word matches
            for qw in query_words:
                if len(qw) > 4 and qw in self.documents[i]:
                    score += 0.3
            scores.append((score, i))

        scores.sort(reverse=True)

        results = []
        for score, idx in scores[:top_k]:
            if score <= 0:
                continue
            doc = KNOWLEDGE_BASE[idx]
            results.append({
                "id": doc["id"],
                "title": doc["title"],
                "content": doc["content"],
                "category": doc["category"],
            })
        return results


_retriever = None

def get_retriever() -> SimpleRetriever:
    global _retriever
    if _retriever is None:
        _retriever = SimpleRetriever()
    return _retriever


# ─────────────────────────────────────────────
# PROMPT BUILDER
# ─────────────────────────────────────────────

def build_rag_prompt(user_question: str, retrieved_docs: list) -> str:
    """Build Indic-Gemma prompt with retrieved context. Enforces Hindi output."""
    if retrieved_docs:
        context_parts = [
            f"[{i+1}] {doc['title']}\n{doc['content'][:600]}"
            for i, doc in enumerate(retrieved_docs)
        ]
        context = "\n\n".join(context_parts)
    else:
        context = "कोई प्रासंगिक जानकारी नहीं मिली।"

    return f"""<|system|>
आप एक सहायक बैंकिंग सहायक हैं जो भारतीय बैंकिंग, लोन, और सरकारी योजनाओं के बारे में सरल हिंदी में जानकारी देते हैं।

नियम:
1. केवल नीचे दी गई जानकारी के आधार पर उत्तर दें। अनुमान न लगाएं।
2. उत्तर छोटा, सरल और बोलने योग्य हो — 3-4 वाक्यों में।
3. यदि जानकारी उपलब्ध नहीं है, तो कहें: "यह जानकारी मेरे पास नहीं है। कृपया अपने बैंक से संपर्क करें।"
4. अंत में केवल एक जरूरी follow-up प्रश्न पूछें (यदि आवश्यक हो)।
5. हमेशा हिंदी में उत्तर दें।

संदर्भ जानकारी:
{context}
<|end|>
<|user|>
{user_question}
<|end|>
<|assistant|>"""


def format_response_for_tts(text: str) -> str:
    """Strip markdown and extra whitespace from LLM output before sending to TTS."""
    text = re.sub(r'\*+', '', text)
    text = re.sub(r'#+\s*', '', text)
    text = re.sub(r'\[[\d]+\]', '', text)
    text = re.sub(r'\n+', ' ', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()


def get_tts_description(text: str) -> str:
    """Speaker description for Indic-Parler-TTS."""
    return (
        "A calm, clear female voice speaking in Hindi. "
        "The speech is measured and helpful, like a bank customer service representative. "
        "Very clear pronunciation, moderate pace, friendly tone."
    )