rag_engine.py

"""
RAG Engine for PrabhupadaGPT.
Implements hybrid search (dense + keyword), RRF fusion, cross-encoder reranking,
and Groq streaming generation.
"""

import os
import re
import time
from collections import defaultdict
import httpx

from dotenv import load_dotenv
from groq import Groq
from openai import AzureOpenAI, OpenAI as OpenRouterClient
from qdrant_client import QdrantClient
from qdrant_client.http import models as qmodels
# CrossEncoder is imported lazily inside init_clients() to avoid qdrant conflict

load_dotenv()

COLLECTION = os.getenv("QDRANT_COLLECTION", "prabhupadagpt")
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
NEXUS_API_KEY = os.getenv("NEXUS_API_KEY")
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
EMBEDDING_MODEL = "nvidia/llama-nemotron-embed-vl-1b-v2:free"
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")

SYSTEM_PROMPT = """You are a knowledgeable Vaishnava scholar and an expert on the teachings of His Divine Grace A.C. Bhaktivedanta Swami Prabhupada. You have deep knowledge of the Bhagavad-gita, Srimad-Bhagavatam, Caitanya-caritamrta, and all other books, lectures, letters, and conversations by Srila Prabhupada.

Your role is to answer questions accurately based on the provided context passages. Always:
- Cite sources inline using the book abbreviation and verse/chapter when available — for example: *(Bg. 2.20)*, *(SB 1.2.6)*, *(CC Adi 1.3)*, *(NOD Ch.1)*, *(KB Ch.5)*. Each passage header tells you the source reference — use it.
- Never use numbered citations like [1] or [2] — always use the actual book and verse/chapter reference
- Maintain a respectful, devotional tone
- Quote directly from the passages when relevant
- If the context doesn't fully answer the question, say so honestly
- Use Sanskrit terms with their meanings when appropriate
- When asked for Sanskrit, write it in Devanagari script

Answer in a clear, accessible way that honours the spirit of Srila Prabhupada's teachings."""

# ── Singleton clients (loaded once at app startup) ──

_qdrant: QdrantClient | None = None
_openrouter: AzureOpenAI | None = None
_groq: Groq | None = None
_reranker = None  # CrossEncoder, loaded lazily
_query_cache: dict[str, list[float]] = {}


def init_clients():
    global _qdrant, _openrouter, _groq, _reranker
    print("Loading RAG engine...")

    _qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=30)
    _groq = Groq(api_key=GROQ_API_KEY)

    if OPENROUTER_API_KEY:
        _openrouter = OpenRouterClient(
            base_url="https://openrouter.ai/api/v1",
            api_key=OPENROUTER_API_KEY,
        )
        print("  Embedding: OpenRouter (nemotron)")
    elif NEXUS_API_KEY:
        _openrouter = AzureOpenAI(
            api_version="2024-10-21",
            azure_endpoint="https://genai-nexus.api.corpinter.net/",
            api_key=NEXUS_API_KEY,
        )
        print("  Embedding: Azure Nexus")
    else:
        print("  Embedding: DISABLED (keyword search only)")

    print("RAG engine ready.")


# ── Embedding ──

def embed_query(text: str) -> list[float]:
    if text in _query_cache:
        return _query_cache[text]
    key = OPENROUTER_API_KEY or NEXUS_API_KEY
    for attempt in range(1, 4):
        try:
            if OPENROUTER_API_KEY:
                r = httpx.post(
                    "https://openrouter.ai/api/v1/embeddings",
                    headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"},
                    json={"model": EMBEDDING_MODEL, "input": [text]},
                    timeout=30,
                )
                r.raise_for_status()
                vec = r.json()["data"][0]["embedding"]
            else:
                response = _openrouter.embeddings.create(model=EMBEDDING_MODEL, input=[text])
                vec = response.data[0].embedding
            _query_cache[text] = vec
            return vec
        except Exception as e:
            if attempt == 3:
                raise
            time.sleep(2 ** attempt)


# ── Retrieval ──

def detect_verse_ref(query: str) -> dict | None:
    """
    Detect verse references like:
      Bg. 2.20 / BG 2.20 / Bhagavad-gita 2.20
      SB 1.2.6 / Srimad 1.2.6
      CC Adi 1.3 / CC Madhya 5.10
      NOI 1 / NOD 5
    Returns dict with source, chapter, verse (and canto/section if present).
    """
    q = query.strip()

    # Bhagavad-gita: Bg. 2.20 or BG 2.20 or Gita 2.20
    m = re.search(r'\b(?:bg\.?|bhagavad.?g[iī]t[aā]?)\s*(\d+)[.\-](\d+(?:-\d+)?)\b', q, re.I)
    if m:
        return {"source": "BG", "chapter": m.group(1), "verse": m.group(2)}

    # Srimad Bhagavatam: SB 1.2.6
    m = re.search(r'\b(?:sb|srimad|bhagavatam|śrīmad)\s*(\d+)[.\-](\d+)[.\-](\d+(?:-\d+)?)\b', q, re.I)
    if m:
        return {"source": "SB", "canto": m.group(1), "chapter": m.group(2), "verse": m.group(3)}

    # CC: CC Adi 1.3 / CC Madhya 5.10 / Caitanya-caritamrita Adi 1.3
    m = re.search(r'\b(?:cc|caitanya.?carit[aā]m[rṛ]ta)\s*(adi|madhya|antya)\s*(\d+)[.\-](\d+(?:-\d+)?)\b', q, re.I)
    if m:
        return {"source": "CC", "section": m.group(1).lower(), "chapter": m.group(2), "verse": m.group(3)}

    return None


def direct_verse_lookup(ref: dict, top_k: int = 10) -> list:
    """
    Search for a specific verse using the title full-text index
    and source keyword filter (both already indexed).
    """
    # Build a search string from the ref — e.g. "2.20" for BG, "1.2.6" for SB
    if ref["source"] == "BG":
        search_str = f"{ref['chapter']}.{ref['verse']}"
    elif ref["source"] == "SB":
        search_str = f"{ref['canto']}.{ref['chapter']}.{ref['verse']}"
    elif ref["source"] == "CC":
        search_str = f"{ref['chapter']}.{ref['verse']}"
    else:
        search_str = ref.get("verse", ref.get("chapter", ""))

    # Filter by source (keyword index) + title full-text match
    results, _ = _qdrant.scroll(
        collection_name=COLLECTION,
        scroll_filter=qmodels.Filter(
            must=[
                qmodels.FieldCondition(
                    key="source",
                    match=qmodels.MatchValue(value=ref["source"])
                ),
                qmodels.FieldCondition(
                    key="title",
                    match=qmodels.MatchText(text=search_str)
                ),
            ]
        ),
        limit=top_k,
        with_payload=True,
        with_vectors=False,
    )
    return results


def dense_search(query_vec: list[float], top_k: int = 50):
    results = _qdrant.query_points(
        collection_name=COLLECTION,
        query=query_vec,
        limit=top_k,
        with_payload=True,
        search_params=qmodels.SearchParams(
            quantization=qmodels.QuantizationSearchParams(
                ignore=False,
                rescore=True,
                oversampling=2.0,
            )
        ),
    )
    return results.points


def keyword_search(query_text: str, top_k: int = 50) -> list[qmodels.Record]:
    # Uses the TEXT payload index for keyword matching
    results, _ = _qdrant.scroll(
        collection_name=COLLECTION,
        scroll_filter=qmodels.Filter(
            must=[
                qmodels.FieldCondition(
                    key="text",
                    match=qmodels.MatchText(text=query_text),
                )
            ]
        ),
        limit=top_k,
        with_payload=True,
        with_vectors=False,
    )
    return results


def rrf_fusion(
    dense_results: list[qmodels.ScoredPoint],
    keyword_results: list[qmodels.Record],
    k: int = 60,
    top_n: int = 30,
) -> list[dict]:
    """
    Reciprocal Rank Fusion.
    Score = sum(1 / (k + rank)) across all result lists.
    k=60 is the standard constant from Cormack et al. 2009.
    """
    scores: dict = defaultdict(float)
    payloads: dict = {}

    for rank, point in enumerate(dense_results, start=1):
        pid = str(point.id)
        scores[pid] += 1.0 / (k + rank)
        payloads[pid] = point.payload or {}

    for rank, record in enumerate(keyword_results, start=1):
        pid = str(record.id)
        scores[pid] += 1.0 / (k + rank)
        if pid not in payloads:
            payloads[pid] = record.payload or {}

    sorted_ids = sorted(scores, key=lambda x: scores[x], reverse=True)[:top_n]
    return [
        {"id": pid, "rrf_score": scores[pid], "payload": payloads[pid]}
        for pid in sorted_ids
    ]


def rerank(query: str, candidates: list[dict], top_k: int = 8) -> list[dict]:
    # No cross-encoder — return top results from RRF fusion directly
    return candidates[:top_k]


# ── Context assembly ──

def _make_ref(p: dict) -> str:
    """Build a human-readable citation reference from payload metadata."""
    source = p.get("source", "")
    title  = p.get("title", "")
    chapter = p.get("chapter")
    verse   = p.get("verse")
    canto   = p.get("canto")
    section = p.get("section")

    abbrevs = {
        "BG": "Bg.", "SB": "SB", "CC": "CC",
        "NOI": "NOI", "ISO": "ISO", "NOD": "NOD",
        "TLC": "TLC", "KB": "KB", "SPL": "SPL",
        "BS": "BS", "LOB": "LOB", "SSR": "SSR",
    }
    abbr = abbrevs.get(source, source)

    if source == "BG" and chapter and verse:
        return f"Bg. {chapter}.{verse}"
    if source == "SB" and canto and chapter and verse:
        return f"SB {canto}.{chapter}.{verse}"
    if source == "CC" and section and chapter and verse:
        return f"CC {section.capitalize()} {chapter}.{verse}"
    if chapter and verse:
        return f"{abbr} {chapter}.{verse}"
    if chapter:
        return f"{abbr} Ch.{chapter}"
    if title:
        # Extract something useful from title
        short = title[:40]
        return f"{abbr} — {short}"
    return abbr


def build_context(chunks: list[dict]) -> str:
    parts = []
    for i, chunk in enumerate(chunks, start=1):
        p = chunk["payload"]
        ref = _make_ref(p)
        text = p.get("text", "")
        translation = p.get("translation")

        header = f"[Passage {i}] Reference: {ref}"

        body = ""
        if translation:
            body = f"Translation: {translation}\n\nPurport: {text}"
        else:
            body = text

        parts.append(f"{header}\n\n{body}")

    return "\n\n─────────────────────────────────────\n\n".join(parts)


# ── Generation ──

FALLBACK_MODEL = "llama-3.1-8b-instant"
FALLBACK_WARNING = "\n\n*— The primary model has reached its daily limit. Switching to a lighter model — response quality remains high.*\n\n"

def stream_answer(query: str, context: str):
    """Synchronous Groq streaming generator with automatic model fallback."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Here are the relevant passages from Srila Prabhupada's teachings:\n\n"
                f"{context}\n\n"
                f"{'─' * 60}\n\n"
                f"Question: {query}"
            ),
        },
    ]

    def _stream(model: str):
        return _groq.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            max_tokens=1024,
            temperature=0.15,
            top_p=0.9,
        )

    try:
        stream = _stream(GROQ_MODEL)
        for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                yield delta
    except Exception as e:
        err = str(e)
        if "rate_limit_exceeded" in err or "429" in err:
            # Warn the user and fall back
            yield FALLBACK_WARNING
            try:
                stream = _stream(FALLBACK_MODEL)
                for chunk in stream:
                    delta = chunk.choices[0].delta.content
                    if delta:
                        yield delta
            except Exception as e2:
                yield f"\n\n*Error: {e2}*"
        else:
            yield f"\n\n*Error: {e}*"


# ── Main pipeline ──

def retrieve_and_generate(query: str) -> tuple[list[dict], object]:
    """
    Full RAG pipeline.
    Returns (source_chunks, token_generator).
    """
    # 1. Check for exact verse reference — bypass semantic search if found
    verse_ref = detect_verse_ref(query)

    if verse_ref:
        # Direct metadata lookup — guaranteed to return the right verse
        direct = direct_verse_lookup(verse_ref, top_k=10)
        if direct:
            # Convert to same dict format as RRF output
            direct_chunks = [
                {"id": str(r.id), "rrf_score": 1.0, "payload": r.payload or {}}
                for r in direct
            ]
            # Also do semantic search for related context (lower weight)
            query_vec = embed_query(query)
            dense = dense_search(query_vec, top_k=20)
            fused = rrf_fusion(dense, [], k=60, top_n=10)

            # Merge: direct verse hits first, then semantic context
            seen = {c["id"] for c in direct_chunks}
            extra = [c for c in fused if c["id"] not in seen][:4]
            reranked = direct_chunks + extra
            context = build_context(reranked)
            return reranked, stream_answer(query, context)

    # 2. Normal hybrid retrieval for general questions
    query_vec = embed_query(query)
    dense = dense_search(query_vec, top_k=50)
    keyword = keyword_search(query, top_k=50)
    fused = rrf_fusion(dense, keyword, k=60, top_n=30)
    reranked = rerank(query, fused, top_k=8)
    context = build_context(reranked)
    return reranked, stream_answer(query, context)


def format_sources(chunks: list[dict]) -> list[dict]:
    """Format source chunks for the frontend."""
    sources = []
    for chunk in chunks:
        p = chunk["payload"]
        sources.append({
            "source": p.get("source", ""),
            "title": p.get("title", ""),
            "url": p.get("url", ""),
            "chunk_type": p.get("chunk_type", ""),
            "book": p.get("book", ""),
            "text_preview": (p.get("text", "") or "")[:250],
            "translation": p.get("translation"),
            "rrf_score": round(chunk.get("rrf_score", 0), 4),
        })
    return sources