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",
    "awas yojana": "awas yojana housing scheme pmay subsidy",
    "awas": "housing pmay",

    # 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 and key Hindi terms from normalized text for retrieval."""
    # Keep English words (alpha) and common Hindi keywords that are in KB tags
    words = [str(w) for w in normalized_text.split() if any(c.isalpha() for c in w)]
    if not words:
        # Fallback to the original text if no normalization happened
        return str(normalized_text)
    # Standard slicing for list of strings
    result_words = words[0:20]
    return " ".join(result_words)


# ─────────────────────────────────────────────
# 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)
            
            # Substantial bonus for exact phrase matching in document
            if query.lower() in self.documents[i]:
                score += 1.0
            
            # Bonus for longer exact word matches
            for qw in query_words:
                if len(qw) > 3 and qw in self.documents[i]:
                    score += 0.2
            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()



# ─────────────────────────────────────────────
# EMBEDDING RETRIEVER — multilingual, Hindi-aware
# Uses paraphrase-multilingual-MiniLM-L12-v2
# Falls back to SimpleRetriever if ChromaDB not ready
# ─────────────────────────────────────────────

try:
    from sentence_transformers import SentenceTransformer
    from FlagEmbedding import FlagReranker
    import chromadb

    print("Loading embedding model (multilingual MiniLM)...")
    embedder = SentenceTransformer(
        'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'
    )
    print("Loading reranker (bge-reranker-v2-m3)...")
    reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=False)
    
    chroma_client = chromadb.PersistentClient(path="./chroma_db")
    EMBEDDING_READY = True
    print("Embedding retriever ready.")

except Exception as e:
    print(f"Embedding retriever not available: {e}. Using keyword retriever.")
    EMBEDDING_READY = False


def retrieve_with_embeddings(query: str, top_k: int = 3) -> list:
    """
    Two-stage retrieval:
    Stage 1 - ChromaDB embedding search (top 10)
    Stage 2 - bge-reranker picks best 3
    Falls back to SimpleRetriever if embeddings not ready.
    Critical for Hinglish queries like 
    'home loan ka interest kitna hai for salaried?'
    """
    if not EMBEDDING_READY:
        print("DEBUG RAG: falling back to keyword retriever")
        retriever = get_retriever()
        return retriever.retrieve(query, top_k=top_k)
    
    try:
        collection = chroma_client.get_collection("banking_hindi")
    except Exception:
        print("DEBUG RAG: ChromaDB collection not found, run ingest.py first")
        print("DEBUG RAG: falling back to keyword retriever")
        retriever = get_retriever()
        return retriever.retrieve(query, top_k=top_k)
    
    try:
        # Stage 1: embedding similarity search
        query_embedding = embedder.encode([query]).tolist()
        results = collection.query(
            query_embeddings=query_embedding,
            n_results=min(10, collection.count())
        )
        candidates = results['documents'][0]
        metadatas = results['metadatas'][0]
        
        print(f"DEBUG RAG: {len(candidates)} candidates from ChromaDB")
        
        # Stage 2: rerank
        pairs = [[query, doc] for doc in candidates]
        scores = reranker.compute_score(pairs)
        ranked = sorted(
            zip(scores, candidates, metadatas), 
            reverse=True
        )
        
        top_results = [
            {
                "id": meta.get("id", ""),
                "title": meta.get("title", ""),
                "content": doc,
                "category": meta.get("category", "")
            }
            for _, doc, meta in ranked[:top_k]
        ]
        print(f"DEBUG RAG: top result = {top_results[0]['title'] if top_results else 'none'}")
        return top_results
        
    except Exception as e:
        print(f"DEBUG RAG: embedding retrieval error: {e}")
        retriever = get_retriever()
        return retriever.retrieve(query, top_k=top_k)