api/rag_engine.py

import os
import sys
import json
from dotenv import load_dotenv

load_dotenv()
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from embeddings.embedder import search
from sentence_transformers import CrossEncoder
from api.utils import detect_jurisdiction, auto_context_depth, tier_sources, strip_think_tags

# ─── Reranker ────────────────────────────────────────────────────────────────
_reranker = None
def get_reranker():
    global _reranker
    if _reranker is None:
        print("Initializing Reranker (Lazy)...")
        _reranker = CrossEncoder('BAAI/bge-reranker-base')
    return _reranker

# ─── Provider Config ──────────────────────────────────────────────────────────
OLLAMA_URL   = os.environ.get("OLLAMA_URL", "http://localhost:11434/api/chat")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "llama3:latest")

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
OPENROUTER_MODEL   = "meta-llama/llama-3.3-70b-instruct:free"

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
GROQ_MODEL   = "llama-3.3-70b-versatile"


LLM_PROVIDER = os.environ.get("LLM_PROVIDER", "openrouter")

# ─── System Prompt ────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are LexRAG, an expert AI counsel for UAE and Indian law, taxation, and accounting.

RULES:
1. If context documents are provided and relevant, answer ONLY from them. Cite source title and jurisdiction.
2. If context is insufficient or the question is off-topic, answer helpfully using your knowledge. Prefix ALL such paragraphs with [INDEPENDENT ANALYSIS].
3. Be concise and structured. Use bullet points for lists of rules or rates.
4. Always end your response with a one-line tag: JURISDICTION: India | UAE | Both | General
5. Never fabricate statute numbers or case names."""

# ─── Retrieval + Reranking ────────────────────────────────────────────────────
def search_and_rerank(question: str, jurisdiction: str = None, top_k: int = 5) -> list:
    filters = {}
    if jurisdiction and jurisdiction != "Both":
        filters["jurisdiction"] = [jurisdiction, "Both"]
    initial = search(question, top_k=20, filters=filters if filters else None)
    if not initial:
        return []
    try:
        pairs  = [[question, d["text"]] for d in initial]
        scores = get_reranker().predict(pairs)
        for i, s in enumerate(scores):
            initial[i]["rerank_score"] = float(s)
        ranked = sorted(initial, key=lambda x: x["rerank_score"], reverse=True)
        return ranked[:top_k]
    except Exception as e:
        print(f"Rerank error: {e}")
        return initial[:top_k]

# ─── Prompt Builder ───────────────────────────────────────────────────────────
def build_prompt(query: str, context_docs: list, history: list = None, confidence: str = "GROUNDED") -> str:
    if context_docs:
        ctx = "\n\n---\n\n".join([
            f"[Source: {d['source']} | Jurisdiction: {d['jurisdiction']} | Date: {d.get('date','')}]\n"
            f"Title: {d['doc_title']}\n\n{d['text']}"
            for d in context_docs
        ])
    else:
        ctx = "No relevant documents found."

    hist_str = ""
    if history:
        hist_str = "CONVERSATION HISTORY:\n" + "\n".join(
            f"{h['role'].upper()}: {h['content']}" for h in history
        ) + "\n\n"

    fallback = ""
    if confidence == "SYNTHESIZED":
        fallback = "\nNote: No strong document matches found. Provide an independent analysis based on your knowledge and mark paragraphs with [INDEPENDENT ANALYSIS].\n"

    return f"""{hist_str}CONTEXT DOCUMENTS:
{ctx}
{fallback}
QUESTION: {query}"""

# ─── Streaming Generators ────────────────────────────────────────────────────
import httpx

async def stream_groq(messages: list, model: str = None):
    model    = model or GROQ_MODEL
    in_think = False
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST", "https://api.groq.com/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"},
            json={"model": model, "messages": messages, "stream": True}
        ) as resp:
            if resp.status_code != 200:
                err_body = await resp.aread()
                try: err_json = json.loads(err_body)
                except: err_json = {"error": {"message": err_body.decode()}}
                msg = err_json.get("error", {}).get("message", "Unknown Groq error")
                raise Exception(f"Groq API Error ({resp.status_code}): {msg}")

            async for line in resp.aiter_lines():
                if not line.startswith("data: "): continue
                data = line[6:]
                if data.strip() == "[DONE]": break
                try:
                    pdata = json.loads(data)
                    if "error" in pdata:
                        raise Exception(f"Groq Stream Error: {pdata['error'].get('message', 'Unknown')}")
                    token = pdata["choices"][0]["delta"].get("content", "")
                    if not token: continue
                    clean, in_think = strip_think_tags(token, in_think)
                    if clean: yield clean
                except Exception as e:
                    if "Stream Error" in str(e) or "API Error" in str(e): raise e
                    pass

async def stream_openrouter(messages: list, model: str = None):
    model    = model or OPENROUTER_MODEL
    in_think = False
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST", "https://openrouter.ai/api/v1/chat/completions",
            headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"},
            json={"model": model, "messages": messages, "stream": True}
        ) as resp:
            if resp.status_code != 200:
                err_body = await resp.aread()
                try: err_json = json.loads(err_body)
                except: err_json = {"error": {"message": err_body.decode()}}
                msg = err_json.get("error", {}).get("message", "Unknown OpenRouter error")
                raise Exception(f"OpenRouter API Error ({resp.status_code}): {msg}")

            async for line in resp.aiter_lines():
                if not line.startswith("data: "): continue
                data = line[6:]
                if data.strip() == "[DONE]": break
                try:
                    pdata = json.loads(data)
                    if "error" in pdata:
                        raise Exception(f"OpenRouter Stream Error: {pdata['error'].get('message', 'Unknown')}")
                    token = pdata["choices"][0]["delta"].get("content", "")
                    if not token: continue
                    clean, in_think = strip_think_tags(token, in_think)
                    if clean: yield clean
                except Exception as e:
                    if "Stream Error" in str(e) or "API Error" in str(e): raise e
                    pass

async def stream_ollama(messages: list, model: str = None):
    model = model or OLLAMA_MODEL
    async with httpx.AsyncClient(timeout=180.0) as client:
        async with client.stream(
            "POST", OLLAMA_URL,
            json={"model": model, "messages": messages, "stream": True}
        ) as resp:
            if resp.status_code != 200:
                err_body = await resp.aread()
                raise Exception(f"Ollama API Error ({resp.status_code}): {err_body.decode()}")
            async for line in resp.aiter_lines():
                if not line.strip():
                    continue
                try:
                    chunk = json.loads(line)
                    token = chunk.get("message", {}).get("content", "")
                    if token: yield token
                    if chunk.get("done"): break
                except Exception:
                    pass

async def stream_provider(messages: list, provider: str, model: str = None):
    if provider == "groq":
        async for t in stream_groq(messages, model): yield t
    elif provider == "openrouter":
        async for t in stream_openrouter(messages, model): yield t
    elif provider == "ollama":
        async for t in stream_ollama(messages, model): yield t
    else:
        async for t in stream_groq(messages, model): yield t

# ─── Legacy sync query (CLI) ──────────────────────────────────────────────────
def query_rag(question: str, jurisdiction: str = None, source_type: str = None,
              top_k: int = None, provider: str = None, session_id: str = "default") -> dict:
    from api.memory import save_message, get_history
    from api.utils import parse_citations
    provider   = provider or LLM_PROVIDER
    top_k      = top_k or auto_context_depth(question)
    jurisdiction = jurisdiction or detect_jurisdiction(question)
    docs       = search_and_rerank(question, jurisdiction, top_k)
    confidence = tier_sources(docs)
    history    = get_history(session_id, limit=5)
    prompt     = build_prompt(question, docs, history, confidence)
    messages   = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}]
    try:
        import httpx as _h
        if provider == "groq":
            r = _h.Client(timeout=60).post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
                json={"model": GROQ_MODEL, "messages": messages}
            ).json()["choices"][0]["message"]["content"]
        elif provider == "openrouter":
            r = _h.Client(timeout=60).post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}"},
                json={"model": OPENROUTER_MODEL, "messages": messages}
            ).json()["choices"][0]["message"]["content"]
        else:
            r = "Provider not supported in sync mode."
        r = parse_citations(r)
    except Exception as e:
        r = f"Error: {e}"
    save_message(session_id, "user", question)
    save_message(session_id, "assistant", r, sources=docs, provider=provider)
    return {"answer": r, "sources": [{"title": d["doc_title"], "source": d["source"],
            "jurisdiction": d["jurisdiction"], "type": d["source_type"], "url": d["url"],
            "score": round(d.get("rerank_score", 0), 3)} for d in docs],
            "context_used": len(docs), "provider": provider, "session_id": session_id,
            "confidence": confidence, "jurisdiction": jurisdiction}

if __name__ == "__main__":
    r = query_rag("What is the GST rate on online gaming contest entry fees in India?")
    print("ANSWER:", r["answer"][:300])
    print("CONFIDENCE:", r["confidence"])
    print("JURISDICTION:", r["jurisdiction"])