scripts/mintoak/chat_server.py

"""
Mintoak Website Assistant — Local Chat Server
================================================
Run with:  mlx-env/bin/python scripts/mintoak/chat_server.py
Open:      http://localhost:5001
"""

import os
import re
import json
import time
import threading

from flask import Flask, request, jsonify, render_template, Response, stream_with_context
from flask_cors import CORS

import chromadb
from chromadb.utils import embedding_functions
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler

# ─────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────
BASE_DIR       = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
CHROMA_DB_PATH = os.path.join(BASE_DIR, "data/mintoak/chroma_db")
COLLECTION_NAME = "mintoak_content"
CHUNKS_JSON    = os.path.join(BASE_DIR, "data/mintoak/mintoak_chunks.json")
MODEL_PATH     = "mlx-community/Qwen2.5-1.5B-Instruct-4bit"
ADAPTER_PATH   = os.path.join(BASE_DIR, "adapters/mintoak")
PORT           = 5001

# ─────────────────────────────────────────────
# GUARDRAIL WORD LISTS
# ─────────────────────────────────────────────
PROFANE_PATTERNS = [
    "fuck", "shit", "asshole", "bitch", "bastard", "cunt", "dick", "pussy",
    "idiot", "stupid", "dumbass", "kill yourself", "kys", "retard",
    "wanker", "motherfucker"
]
INJECTION_PATTERNS = [
    "ignore previous", "system prompt", "developer mode", "you are now",
    "pretend to be", "reveal instructions", "system instruction", "dan mode",
    "hypothetically speaking", "override rules", "jailbreak"
]
GREETING_RE  = re.compile(r"^(hi+|hello+|hey+|yo+|greetings|good\s+(morning|afternoon|evening))\b", re.I)
IDENTITY_RE  = re.compile(r"\b(who are (you|u)|what is your name|what'?s your name|what do you do|introduce yourself|tell me about yourself)\b", re.I)
PREAMBLE_RES = [
    re.compile(r"^(?:as per|according to|based on)\s+(?:the\s+)?(?:context|provided context|mintoak documentation|documentation)(?:,\s*)?", re.I),
    re.compile(r"^based on what is provided in the context(?:,\s*)?", re.I),
    re.compile(r"^from the context(?:,\s*)?", re.I),
    re.compile(r"^Based on the Mintoak documentation,\s*(?:here is what we know about\s+[^:\n]+:)?\s*", re.I),
]

SYSTEM_PROMPT = (
    "You are the Mintoak Website Assistant, a senior, business-aware digital advisor on the Mintoak website. "
    "Your tone is professional, clear, confident, technology-forward, and benefit-led. Do not use hype, casualness, or over-enthusiasm. "
    "Do not use emojis by default; use at most one emoji only where it adds clarity (e.g. 💡, 📈) and never in pricing, credibility, or trust contexts. "
    "Answer strictly using the provided context. Never invent, assume, or pull from outside knowledge. "
    "Mintoak is strictly a bank-led B2B platform for merchants; Mintoak does NOT issue loyalty points or rewards to end-consumers. If asked how customers or end-consumers earn loyalty points, clarify that Mintoak does not directly run consumer loyalty programs, but enables merchants to set up their own custom campaigns and offers for their customers. Do not confuse B2B merchant app users with B2C end-consumers. "
    "Start your response directly and naturally—do NOT use robotic preambles like 'As per the context' or 'Based on the provided context'. "
    "Frame Mintoak as a white-labeled, modular, cloud-native, API-first SaaS platform/solution provider (never as a 'software company', 'software', 'tool', or 'app'). "
    "Avoid these banned phrases: 'seamless', 'empowering', 'leverage' (as a verb), 'game-changer', 'synergy', 'unlocking', 'in today\'s fast-paced world', 'as we move forward'. "
    "Follow this exact structure: Context (direct answer) -> Insight (translate features into merchant/business outcomes like activation, cost-to-acquire, time-to-market) -> Direction (nudge to the next best page/demo and close with the source URL). "
    "Keep responses highly concise (maximum 2-3 sentences or short bullets) unless explicitly asked for detail. "
    "If the context does not contain the answer, reply exactly: 'The requested information does not currently exist on www.mintoak.com. You can get in touch with our team at https://www.mintoak.com/contact-us.' "
    "IMPORTANT - Lead Capture Flow: If the user requests a demo, pricing, partnership, or contact details, give a brief helpful response and politely ask for their name and email. "
    "If and ONLY if the user has already provided their name and email in the conversation history, confirm follow-up details and append exactly '[CAPTURE_LEAD]' at the very end of your response."
)

# ─────────────────────────────────────────────
# STARTUP: Load model & DB once
# ─────────────────────────────────────────────
print("🔄  Loading ChromaDB …")
_chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
_emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
_collection = _chroma_client.get_or_create_collection(name=COLLECTION_NAME, embedding_function=_emb_fn)

# Populate database on start if it's empty
if _collection.count() == 0 and os.path.exists(CHUNKS_JSON):
    print("Vector database is empty. Populating chunks on startup...")
    with open(CHUNKS_JSON, "r", encoding="utf-8") as f:
        chunks = json.load(f)
    
    ids = [c["id"] for c in chunks]
    documents = [f"Title: {c['title']}\nContent: {c['content']}" for c in chunks]
    metadatas = [{"url": c["url"], "title": c["title"], "category": c["category"]} for c in chunks]
    
    # ChromaDB batch insertion
    batch_size = 500
    for i in range(0, len(chunks), batch_size):
        _collection.add(
            ids=ids[i:i+batch_size],
            documents=documents[i:i+batch_size],
            metadatas=metadatas[i:i+batch_size]
        )
print(f"✅  Vector DB ready — {_collection.count()} passages indexed.")

print("🔄  Loading LLM weights (this takes ~3 s the first time) …")
if os.path.exists(ADAPTER_PATH):
    _model, _tokenizer = load(MODEL_PATH, adapter_path=ADAPTER_PATH)
else:
    _model, _tokenizer = load(MODEL_PATH)
print("✅  Model ready. Starting chat server …\n")

# ─────────────────────────────────────────────
# FLASK APP
# ─────────────────────────────────────────────
app = Flask(__name__, template_folder="templates", static_folder="static")
CORS(app)
_model_lock = threading.Lock()   # serialize concurrent generation calls


# ─────────────────────────────────────────────
# RAG HELPERS
# ─────────────────────────────────────────────
# CONFIG-driven QUERY ENHANCEMENT RULES
# ─────────────────────────────────────────────
QUERY_ENHANCEMENT_RULES = [
    {
        "keywords": ["gpv", "gross payment", "volume", "process", "transaction", "festive"],
        "expansion": "During India's October 2025 festive period, Mintoak's platform recorded 44% year-on-year growth in digital payment value across its merchant base",
        "prompt_note": "State that Mintoak processes payments for over 4 million merchants, recording a 44% YoY growth in digital payment value and 42% YoY growth in digital payment volume during the October 2025 festive period.",
        "top_k": 3
    },
    {
        "keywords": ["document", "documents", "checklist", "paperwork", "certificates"],
        "expansion": "What documents are needed for merchant onboarding? DigiOnboard KYC KYB",
        "prompt_note": "Ensure you list the specific documents needed: registration certificates, GST certificates, cancelled cheques, and board resolutions.",
        "top_k": 3
    },
    {
        "keywords": ["acquire merchant", "merchant acquisition", "onboard", "onboarding", "kyc", "kyb"],
        "expansion": "Mintoak DigiOnboard merchant onboarding KYC KYB acquisition banks",
        "prompt_note": None,
        "top_k": 2
    },
    {
        "keywords": ["rewardrun", "reward run", "gamified campaign", "gamified campaigns", "loyalty campaign", "loyalty campaigns"],
        "expansion": "Mintoak RewardRun gamified loyalty and merchant engagement platform for acquirers",
        "override": True,
        "prompt_note": "Ensure you mention Mintoak RewardRun as the gamified loyalty and merchant engagement platform.",
        "top_k": 3
    },
    {
        "keywords": ["job", "jobs", "career", "careers", "hiring", "opening", "openings", "join us", "work at mintoak"],
        "expansion": "Mintoak careers job openings hiring opportunities join the team",
        "prompt_note": "Direct the user to the careers page at https://www.mintoak.com/career.",
        "top_k": 2
    },
    {
        "keywords": ["address", "office", "offices", "headquarters", "location", "locations", "located", "contact us", "contact-us"],
        "expansion": "Mintoak office locations and addresses Mumbai Bangalore Dubai contact-us",
        "override": True,
        "prompt_note": "List all Mintoak office locations (Mumbai, Bangalore, Dubai) with their Google Maps links.",
        "top_k": 2
    },
    {
        "keywords": ["cto", "technology head", "tech head", "head of technology", "head of engineering", "engineering head", "who is your technology head", "who is the technology head"],
        "expansion": "Kabeer Jain is the Chief Technology Officer (CTO) and Co-founder of Mintoak.",
        "override": True,
        "prompt_note": "State that Kabeer Jain is Mintoak's Chief Technology Officer (CTO) and Co-founder.",
        "top_k": 2
    },
    {
        "keywords": ["culture", "work culture", "life at mintoak", "work environment", "mintoak's culture", "values", "core values", "team culture"],
        "expansion": "About Us | Mintoak | SaaS | FinTech. Mintoak values: Curiosity (Fostering learning), Diversity (Embracing differences), Excellence (Striving towards the highest standards), Empowerment (Fostering confidence, autonomy, and growth), Integrity (Honoring commitments), Innovation (Challenging status quo), Impact (Driving meaningful outcomes). Collaboration is at the heart of our culture. We encourage networking, emotional intelligence, and cross-functional teamwork to drive growth and success together.",
        "override": True,
        "prompt_note": "Describe Mintoak's collaborative culture, its team culture, and list its core values (Curiosity, Diversity, Excellence, Empowerment, Integrity, Innovation, Impact).",
        "top_k": 2
    },
    {
        "keywords": ["ceo", "founder", "founders", "founded", "cpo", "leadership", "team", "management", "raman", "khanduja", "sanjay", "nazareth", "rama", "tadepalli", "kabeer", "jain", "rohit", "ramana"],
        "expansion": "Mintoak co-founders and executive leadership team: Raman Khanduja (CEO), Sanjay Nazareth (COO), Rama Tadepalli (CPO), Kabeer Jain (CTO), and Rohit Ramana (CFO). Nilesh Lonkar is the VP of Engineering.",
        "override": True,
        "prompt_note": "Ensure you mention Mintoak's co-founders: Raman Khanduja (CEO), Sanjay Nazareth (COO), Rama Tadepalli (CPO), Kabeer Jain (CTO), and Rohit Ramana (CFO).",
        "top_k": 2
    }
]

MAX_CHUNK_CHARS = 1200   # Truncate each chunk to cap prompt size & speed up prefill


def has_lead_intent(query: str) -> bool:
    q = query.lower()
    # Check for demo requests
    if any(kw in q for kw in ["book a demo", "schedule a demo", "request a demo", "book demo", "schedule demo", "request demo", "get a demo", "live demo"]):
        return True
    # Check for pricing/cost queries
    if any(kw in q for kw in ["pricing", "price list", "pricing plans", "subscription cost", "licensing cost", "what are the charges", "setup fees"]):
        return True
    if "cost" in q and any(w in q for w in ["how much", "what is", "what's"]):
        return True
    if "how much" in q and any(w in q for w in ["cost", "charge", "pay", "fee"]):
        return True
    # Check for partnership queries
    if any(kw in q for kw in ["become a partner", "partnering with you", "partnership opportunities", "partner with mintoak"]):
        return True
    # Check for contact/sales queries
    if any(kw in q for kw in ["talk to sales", "sales team", "contact sales", "sales department", "sales contact"]):
        return True
    if any(kw in q for kw in ["get in touch", "how to contact", "contact us", "contact details", "callback", "call back"]):
        return True
    # Check for request to contact
    if "call me" in q or "phone number" in q or "email address" in q:
        return True
    return False

def retrieve_context(query: str, history: list = None, top_k: int = 2):
    """Return (context_text, url_to_title_dict, prompt_note) or ('OUT_OF_SCOPE_REFUSAL', {}, None)."""
    query_lower = query.lower()
    search_query = query
    prompt_note = None

    # Reformulate query if it is a conversational check-again feedback query
    if history:
        feedback_phrases = [
            "check again", "please check", "wrong", "incorrect", "not true", 
            "this wont", "this won't", "that is not", "explain further", "details", 
            "tell me more", "don't agree", "dont agree", "not helpful", "what else",
            "other way", "more details"
        ]
        is_feedback = any(p in query_lower for p in feedback_phrases)
        if is_feedback:
            # Find the last user query
            last_user_query = None
            for msg in reversed(history):
                if msg.get("role") == "user":
                    last_user_query = msg.get("content")
                    break
            if last_user_query:
                # If the last query was about business growth, route to alternative/general growth topics
                if any(w in last_user_query.lower() for w in ["grow", "business", "help", "merchants"]):
                    search_query = f"{query} Mintoak payment acceptance SmartPayments loyalty rewards SellSmart lending business growth"
                    top_k = 3
                else:
                    search_query = last_user_query

    # Apply configuration-driven expansion rules dynamically (without wiping out original query keywords)
    for rule in QUERY_ENHANCEMENT_RULES:
        if any(kw in query_lower for kw in rule["keywords"]):
            if rule.get("override"):
                search_query = rule["expansion"]
            else:
                search_query = f"{query} {rule['expansion']}"
            top_k = rule["top_k"]
            prompt_note = rule.get("prompt_note")
            break

    results = _collection.query(query_texts=[search_query], n_results=top_k)

    best_distance = 999.0
    if results and results.get("distances") and results["distances"][0]:
        best_distance = results["distances"][0][0]

    # Brand-relevant context terms (in-scope even if Mintoak is not explicitly named)
    in_scope_keywords = [
        # 1. Brand Names
        "mintoak", "digiledge",
        # 2. Product Names
        "smartpayments", "digionboard", "rewardrun", "business360", "sellsmart", "staffaccess", "soundhub", "soundbox",
        # 3. Leadership Names
        "raman", "khanduja", "sanjay", "nazareth", "rama", "tadepalli", "kabeer", "jain", "rohit", "ramana",
        # 4. FinTech Domain
        "payment", "payments", "transaction", "transactions", "merchant", "merchants", "bank", "banks", "card", "cards", 
        "gpv", "volume", "acquiring", "acquirer", "acquirers", "credit", "lending", "loan", "loans", "invoicing", 
        "reconciliation", "settlement", "settlements",
        # 5. MSME/SME
        "sme", "smes", "msme", "msmes", "business", "businesses", "retail", "retailer", "retailers",
        # 6. Onboarding & KYC
        "onboard", "onboarding", "kyc", "kyb", "register", "registration", "account", "accounts", "signup", "signin", 
        "login", "log", "logout",
        # 7. Security & Credentials
        "password", "passwords", "reset", "otp", "credential", "credentials", "auth", "authentication", "security",
        # 8. Contact & Offices
        "address", "addresses", "office", "offices", "contact", "location", "locations", "headquarters", "registered", 
        "corporate", "callback", "call", "support", "help", "demo",
        # 9. Careers
        "job", "jobs", "career", "careers", "hiring", "opening", "openings", "join", "recruit", "recruitment",
        # 10. Reports & Docs
        "report", "reports", "statement", "statements", "download", "downloads", "dashboard",
        # 11. Developer Resources
        "api", "apis", "sdk", "sdks", "developer", "developers"
    ]
    import re
    is_brand_context = any(re.search(rf"\b{re.escape(w)}\b", query_lower) for w in in_scope_keywords)
    # Strict cut-off: if distance is extremely high, it's always out-of-scope
    if best_distance > 0.82:
        return "OUT_OF_SCOPE_REFUSAL", {}, None
    # Moderate cut-off: allow bypass only if it has brand context
    if best_distance > 0.60 and not is_brand_context:
        return "OUT_OF_SCOPE_REFUSAL", {}, None

    context_parts = []
    url_to_title = {}

    # Always inject product catalog for product-related queries
    search_lower = search_query.lower() if search_query else ""
    if any(pw in query_lower or pw in search_lower for pw in ["product", "products", "offering", "offerings", "catalog", "offer", "offers"]):
        try:
            cat = _collection.get(ids=["synthetic_master_catalog"])
            if cat and cat.get("documents"):
                context_parts.append(f"Source: {cat['metadatas'][0]['url']}\nContent: {cat['documents'][0]}")
                url_to_title[cat['metadatas'][0]['url']] = cat['metadatas'][0].get("title", "Mintoak Product Catalog")
        except Exception:
            pass

    if results and results.get("documents") and results["documents"][0]:
        docs = results["documents"][0]
        metas = results["metadatas"][0]
        dists = results["distances"][0] if results.get("distances") else [0.0] * len(docs)
        best_dist = dists[0]
        for i, (doc, meta, dist) in enumerate(zip(docs, metas, dists)):
            if dist > best_dist + 0.15 and dist > 0.45:
                continue
            if meta.get("title") == "Mintoak Product List Catalog" and context_parts:
                continue
            # Truncate chunk to MAX_CHUNK_CHARS to keep prompt short → faster prefill
            doc_trimmed = doc[:MAX_CHUNK_CHARS] + ("…" if len(doc) > MAX_CHUNK_CHARS else "")
            context_parts.append(f"Source: {meta['url']}\nContent: {doc_trimmed}")
            url_to_title[meta["url"]] = meta.get("title", "Learn more")

    return "\n\n---\n\n".join(context_parts), url_to_title, prompt_note


def clean_response(text: str, allow_lead_capture: bool = True) -> str:
    """Strip robotic preambles and fix capitalisation."""
    for rx in PREAMBLE_RES:
        text = rx.sub("", text).strip()
    # Strip any model-generated citation lines like:
    # "👉 For more details, visit: ..." or "For more details, visit: ..."
    text = re.sub(r'(?i)(?:\n\s*)*👉?\s*For\s+more\s+details,?\s+visit:.*$', '', text).strip()
    
    # Handle LLM fallback non-compliance phrases and override with the brand-approved fallback
    has_lead = "[CAPTURE_LEAD]" in text and allow_lead_capture
    text_no_lead = text.replace("[CAPTURE_LEAD]", "").strip()
    text_lower = text_no_lead.lower()
    
    fallback_triggers = [
        "not fully answered", "not mentioned", "not provided", "unable to answer", 
        "cannot be determined", "no mention", "not find any information", 
        "does not contain", "no information", "not in the context", "not explicitly mentioned",
        "isn't mentioned", "not found", "does not provide", "not provide", "text does not",
        "context does not", "not mention", "does not specify", "not specified", "do not have",
        "not have information", "not explicitly stated", "does not appear to contain"
    ]
    if any(trigger in text_lower for trigger in fallback_triggers):
        text_no_lead = (
            "The requested information does not currently exist on www.mintoak.com. "
            "You can get in touch with our team at https://www.mintoak.com/contact-us."
        )
    else:
        # Programmatic Sanitization Filters (Zero-Tolerance Enforcement)
        # Brand name correction
        text_no_lead = re.sub(r'(?<![\./])\b(MintOak|MINTOAK|mintoak)\b(?![\./])', 'Mintoak', text_no_lead)
        
        # Banned words replacement
        text_no_lead = re.sub(r'\bseamlessly\b', 'efficiently', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bseamless\b', 'integrated', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bempowers\b', 'enables', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bempowering\b', 'enabling', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bempower\b', 'enable', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bleveraged\b', 'utilized', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bleverages\b', 'utilizes', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bleveraging\b', 'utilizing', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bleverage\b', 'use', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bgame-changer\b', 'significant advancement', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bgame changer\b', 'significant advancement', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bsynergy\b', 'alignment', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bsynergies\b', 'alignments', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bunlocking\b', 'releasing', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\bunlock\b', 'access', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r"\bin today's fast-paced world,?\s*", "", text_no_lead, flags=re.I)
        text_no_lead = re.sub(r"\bas we move forward,?\s*", "", text_no_lead, flags=re.I)
        
        # Software/tool/app -> platform/module (in the context of Mintoak)
        text_no_lead = re.sub(r'\bMintoak (\w+ )?(?:software|tool|app)\b', r'Mintoak \1platform', text_no_lead, flags=re.I)
        text_no_lead = re.sub(r'\b(?:software|tool|app) Mintoak\b', 'Mintoak platform', text_no_lead, flags=re.I)

        # Emoji Limiter (maximum of 1 emoji)
        emoji_pattern = re.compile(r'[\U00010000-\U0010ffff\u2600-\u27bf]', flags=re.UNICODE)
        matches = list(emoji_pattern.finditer(text_no_lead))
        if len(matches) > 1:
            cleaned_chars = []
            for i, char in enumerate(text_no_lead):
                is_subsequent_emoji = False
                for match in matches[1:]:
                    m_start, m_end = match.span()
                    if m_start <= i < m_end:
                        is_subsequent_emoji = True
                        break
                if not is_subsequent_emoji:
                    cleaned_chars.append(char)
            text_no_lead = "".join(cleaned_chars)

    if has_lead:
        text = text_no_lead + " [CAPTURE_LEAD]"
    else:
        text = text_no_lead

    if text and text[0].islower():
        text = text[0].upper() + text[1:]
    return text


def build_sources(context: str, url_to_title: dict, query: str = "") -> list:
    """Extract unique source URLs from context and pair with titles, adding matching product links dynamically."""
    sources = []
    seen = set()
    
    # 1. Parse retrieved sources
    for line in context.split("\n"):
        if line.startswith("Source: "):
            url = line.replace("Source: ", "").strip()
            if url not in seen:
                seen.add(url)
                sources.append({"url": url, "title": url_to_title.get(url, "Learn more")})
                
    # 2. Add relevant product offerings if matching keyword occurs in query
    if query:
        q = query.lower()
        product_mappings = [
            (["grow", "business", "sme", "merchant", "increase", "sales"], "https://www.mintoak.com/products/mintoak-business360", "Mintoak Business 360"),
            (["loyalty", "reward", "rewards", "points", "campaign"], "https://www.mintoak.com/products/mintoak-marketinghub", "Mintoak Marketing Hub"),
            (["onboard", "kyc", "kyb", "registration", "document"], "https://www.mintoak.com/products/mintoak-digionboard", "Mintoak DigiOnboard"),
            (["pos", "payment", "soundbox", "soundhub", "terminal"], "https://www.mintoak.com/products/mintoak-soundhub", "Mintoak SoundHub"),
            (["payment acceptance", "merchant payments", "transactions"], "https://www.mintoak.com/products/mintoak-smartpayments", "Mintoak SmartPayments")
        ]
        
        for kws, p_url, p_title in product_mappings:
            if any(kw in q for kw in kws):
                if p_url not in seen:
                    seen.add(p_url)
                    sources.append({"url": p_url, "title": p_title})
                    
    # Sort sources: product pages on top, then blog pages, then others
    def get_source_priority(s):
        url = s.get("url", "").lower()
        if "/products" in url:
            return 0
        elif "/blog" in url:
            return 1
        return 2

    sources.sort(key=get_source_priority)
    if len(sources) > 1:
        if "/products" in sources[0]["url"].lower() and "/products" in sources[1]["url"].lower():
            sources = sources[:2]
        else:
            sources = sources[:1]
    return sources


def generate_suggestions(query: str) -> list:
    """
    Instant keyword-based follow-up suggestions — zero LLM calls, zero latency.
    Maps detected topics to curated follow-up questions.
    """
    q = query.lower()

    TOPIC_MAP = [
        (["upi", "unified payment"],  [
            "What are the transaction fees for UPI payments?",
            "How does Mintoak support UPI for small merchants?",
            "Can UPI payments be tracked in real time?"
        ]),
        (["product", "offering", "catalog", "what do you offer"],  [
            "Tell me more about Mintoak's merchant engagement platform",
            "How does merchant onboarding work with Mintoak?",
            "What payment methods does Mintoak support?"
        ]),
        (["loyalty", "reward", "cashback", "points"],  [
            "How do customers earn loyalty points?",
            "Can merchants customize their loyalty offers?",
            "How does Mintoak's merchant engagement platform benefit banks?"
        ]),
        (["merchant", "onboard", "kyc", "kyb"],  [
            "What documents are needed for merchant onboarding?",
            "How long does merchant onboarding take with Mintoak?",
            "What is the difference between KYC and KYB?"
        ]),
        (["bank", "banking", "financial", "payment"],  [
            "How does Mintoak help banks acquire merchants?",
            "What analytics does Mintoak provide to banks?",
            "How does Mintoak's white-label solution work?"
        ]),
        (["credit card", "debit card", "card"],  [
            "What card payment types does Mintoak support?",
            "How do business credit cards help merchants?",
            "What are the benefits of card acceptance for SMEs?"
        ]),
        (["sme", "small business", "merchant account"],  [
            "How can SMEs benefit from digital payments?",
            "What is a merchant account and why do I need one?",
            "How does Mintoak help SMEs grow?"
        ]),
        (["africa", "african", "india", "indian"],  [
            "How does Mintoak operate across different markets?",
            "What is Mintoak's approach to emerging markets?",
            "How does digital merchant onboarding work in Africa?"
        ]),
    ]

    for keywords, suggestions in TOPIC_MAP:
        if any(kw in q for kw in keywords):
            return suggestions

    # Generic fallback suggestions
    return [
        "What products does Mintoak offer?",
        "How can Mintoak help my business grow?",
        "How do I get in touch with the Mintoak team?"
    ]


# ─────────────────────────────────────────────
# ROUTES
# ─────────────────────────────────────────────
@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/chat", methods=["POST"])
def chat():
    data = request.get_json(force=True)
    query = (data.get("query") or "").strip()
    history = data.get("history") or []
    if not query:
        return Response("data: " + json.dumps({"error": "Empty query"}) + "\n\n", mimetype="text/event-stream")

    query_lower = query.lower()

    # Helper to yield static event-stream response
    def stream_static(response_text, sources, suggestions, type_name):
        yield f"data: {json.dumps({'event': 'sources', 'sources': sources})}\n\n"
        yield f"data: {json.dumps({'event': 'token', 'token': response_text})}\n\n"
        yield f"data: {json.dumps({'event': 'done', 'suggestions': suggestions, 'latency': 0.0, 'type': type_name})}\n\n"

    # ── Guardrail: profanity ──────────────────
    # ── Guardrail: profanity check ─────────────
    if any(w in query_lower for w in PROFANE_PATTERNS):
        return Response(stream_static(
            "I can assist with questions about Mintoak's payment, merchant engagement, and value-added services platform. How can I guide you on those topics today?",
            [], [], "guardrail"
        ), mimetype="text/event-stream")

    # ── Guardrail: prompt injection ───────────
    if any(p in query_lower for p in INJECTION_PATTERNS):
        return Response(stream_static(
            "I can assist with questions about Mintoak's payment, merchant engagement, and value-added services platform. How can I guide you on those topics today?",
            [], [], "guardrail"
        ), mimetype="text/event-stream")

    # ── Greeting ──────────────────────────────
    query_clean = query.strip().lower().rstrip("?").strip()
    if GREETING_RE.match(query_clean):
        return Response(stream_static(
            "Hello! I am the Mintoak Website Assistant. How can I help you with Mintoak's platform or services today?",
            [],
            [
                "What products does Mintoak offer?",
                "What is UPI and how does Mintoak support it?",
                "How can Mintoak help my business?"
            ],
            "greeting"
        ), mimetype="text/event-stream")

    # ── Identity ──────────────────────────────
    if IDENTITY_RE.search(query_clean):
        return Response(stream_static(
            "I am the Mintoak Website Assistant. I can answer questions about Mintoak's white-labeled SaaS platform, product modules, and services. How can I help you today?",
            [],
            [
                "What products does Mintoak offer?",
                "Tell me about Mintoak's merchant engagement platform",
                "How does merchant onboarding work?"
            ],
            "identity"
        ), mimetype="text/event-stream")

    # ── RAG retrieval ─────────────────────────
    context, url_to_title, prompt_note = retrieve_context(query, history=history)
    if context == "OUT_OF_SCOPE_REFUSAL":
        return Response(stream_static(
            "I can assist with questions about Mintoak's payment, merchant engagement, and value-added services platform. How can I guide you on those topics today?",
            [],
            [
                "What products does Mintoak offer?",
                "What is UPI and how does Mintoak support it?",
                "How does Mintoak help merchants?"
            ],
            "out_of_scope"
        ), mimetype="text/event-stream")

    # ── LLM inference ─────────────────────────
    sources = build_sources(context, url_to_title, query=query)

    def stream_rag():
        # First send sources so UI can render them instantly
        yield f"data: {json.dumps({'event': 'sources', 'sources': sources})}\n\n"

        # Check if user has provided name + email in history or query
        history_text = " ".join(msg.get("content", "") for msg in history) + " " + query
        has_provided_info = "@" in history_text
        allow_lead_capture = has_provided_info and has_lead_intent(query)

        # Rewrite the query passed to the LLM if it is a conversational feedback query
        llm_query = query
        if history:
            feedback_phrases = [
                "check again", "please check", "wrong", "incorrect", "not true", 
                "this wont", "this won't", "that is not", "explain further", "details", 
                "tell me more", "don't agree", "dont agree", "not helpful", "what else",
                "other way", "more details"
            ]
            if any(p in query_lower for p in feedback_phrases):
                # Find the last user query
                last_user_query = None
                for msg in reversed(history):
                    if msg.get("role") == "user":
                        last_user_query = msg.get("content")
                        break
                if last_user_query and any(w in last_user_query.lower() for w in ["grow", "business", "help", "merchants"]):
                    llm_query = "What other ways can Mintoak help my business grow, such as through payment acceptance, credit/lending, or analytics?"

        user_prompt = f"Context:\n{context}\n\nQuestion: {llm_query}"
        if prompt_note:
            user_prompt += f"\n\nNote: {prompt_note}"

        # Build messages from system prompt + conversation history + current user query
        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        for msg in history:
            role = "user" if msg.get("role") == "user" else "assistant"
            content = msg.get("content", "").strip()
            content = content.replace("[CAPTURE_LEAD]", "").strip()
            messages.append({"role": role, "content": content})
        
        messages.append({"role": "user", "content": user_prompt})
        prompt = _tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        t0 = time.time()
        full_response = ""

        with _model_lock:
            gen = stream_generate(_model, _tokenizer, prompt=prompt, max_tokens=256, sampler=make_sampler(temp=0.0))
            citation_started = False
            lead_tag_sent = False
            for tok in gen:
                token_text = tok.text
                full_response += token_text
                
                if citation_started:
                    if "[CAPTURE_LEAD]" in full_response and not lead_tag_sent and allow_lead_capture:
                        lead_tag_sent = True
                        yield f"data: {json.dumps({'event': 'token', 'token': '[CAPTURE_LEAD]'})}\n\n"
                    continue
                
                last_snippet = full_response[-40:].lower()
                if "👉" in token_text or "for more details" in last_snippet or "details, visit" in last_snippet:
                    citation_started = True
                    continue
                
                if "[CAPTURE_LEAD]" in token_text:
                    if allow_lead_capture and not lead_tag_sent:
                        lead_tag_sent = True
                        yield f"data: {json.dumps({'event': 'token', 'token': '[CAPTURE_LEAD]'})}\n\n"
                    continue
                
                yield f"data: {json.dumps({'event': 'token', 'token': token_text})}\n\n"
            
            # Ensure lead tag is sent if generated
            should_send_lead = allow_lead_capture and "[CAPTURE_LEAD]" in full_response
            if should_send_lead and not lead_tag_sent:
                yield f"data: {json.dumps({'event': 'token', 'token': ' [CAPTURE_LEAD]'})}\n\n"
            
            # Auto-save lead in backend if details are present and capture lead was triggered
            if "[CAPTURE_LEAD]" in full_response and has_provided_info:
                email_match = re.search(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', history_text)
                if email_match:
                    email = email_match.group(0)
                    name_match = re.search(r'\b(?:my name is|i am|this is)\s+([A-Za-z\s]{2,30})', history_text, re.I)
                    name = name_match.group(1).strip() if name_match else "Interested Acquirer"
                    try:
                        leads_file = os.path.join(BASE_DIR, "data/mintoak/leads.jsonl")
                        os.makedirs(os.path.dirname(leads_file), exist_ok=True)
                        lead_entry = {
                            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                            "name": name,
                            "email": email,
                            "query": query,
                            "conversation_id": data.get("conversation_id", "direct")
                        }
                        with open(leads_file, "a", encoding="utf-8") as f:
                            f.write(json.dumps(lead_entry) + "\n")
                        print(f"👥 Lead Captured (Auto-Save): {name} ({email})")
                    except Exception as e:
                        print(f"Error auto-saving lead in chat route: {e}")

        latency = round(time.time() - t0, 2)
        cleaned_response = clean_response(full_response.strip(), allow_lead_capture=allow_lead_capture)

        # If references were not generated in the response body, append them dynamically
        if "Source:" not in cleaned_response and "details, visit:" not in cleaned_response and sources:
            links = ", ".join(f"[{s['title']}]({s['url']})" for s in sources)
            appended_text = f"\n\n👉 For more details, visit: {links}"
            yield f"data: {json.dumps({'event': 'token', 'token': appended_text})}\n\n"

        suggestions = generate_suggestions(query)
        yield f"data: {json.dumps({'event': 'done', 'suggestions': suggestions, 'text': cleaned_response, 'latency': latency, 'type': 'rag'})}\n\n"
    return Response(stream_with_context(stream_rag()), mimetype="text/event-stream")


@app.route("/api/health", methods=["GET"])
def health():
    return jsonify({
        "status": "ok",
        "model": MODEL_PATH,
        "collection_size": _collection.count()
    })


@app.route("/api/title", methods=["POST"])
def generate_title():
    try:
        data = request.get_json(force=True)
        message = data.get("message", "").strip()
        if not message:
            return jsonify({"title": "New Chat"})

        # Prompt the LLM to summarize the first message into a short title
        prompt = (
            f"<|im_start|>system\nYou are a helpful assistant. Summarize the user's query into a short conversation title of 3-5 words. Do not use quotes, punctuation, or any meta-text. Return only the title.<|im_end|>\n"
            f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        )
        with _model_lock:
            gen = stream_generate(_model, _tokenizer, prompt=prompt, max_tokens=12, sampler=make_sampler(temp=0.0))
            title = "".join(tok.text for tok in gen).strip()

        # Clean title
        title = re.sub(r'["\'\-\.\!\?]', '', title).strip()
        if not title or len(title) > 40:
            title = message[:35] + ("..." if len(message) > 35 else "")
        return jsonify({"title": title})
    except Exception:
        title = message[:35] + ("..." if len(message) > 35 else "")
        return jsonify({"title": title or "New Chat"})


@app.route("/api/lead", methods=["POST"])
def save_lead():
    try:
        data = request.get_json(force=True)
        name = data.get("name", "").strip()
        email = data.get("email", "").strip()
        query = data.get("query", "").strip()
        conv_id = data.get("conversation_id", "").strip()
        
        leads_file = os.path.join(BASE_DIR, "data/mintoak/leads.jsonl")
        os.makedirs(os.path.dirname(leads_file), exist_ok=True)
        
        lead_entry = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "name": name,
            "email": email,
            "query": query,
            "conversation_id": conv_id
        }
        with open(leads_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(lead_entry) + "\n")
            
        print(f"👥 Lead Captured: {name} ({email}) for query '{query}'")
        return jsonify({"status": "success", "message": "Lead saved successfully"})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


# ─────────────────────────────────────────────
# ENTRY POINT
# ─────────────────────────────────────────────
if __name__ == "__main__":
    print(f"🚀  Mintoak Chat Server running → http://localhost:{PORT}\n")
    app.run(host="0.0.0.0", port=PORT, debug=False, threaded=True)