| 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 = None |
| def get_reranker(): |
| global _reranker |
| if _reranker is None: |
| print("Initializing Reranker (Lazy)...") |
| _reranker = CrossEncoder('BAAI/bge-reranker-base') |
| return _reranker |
|
|
| |
| 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 = """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.""" |
|
|
| |
| 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] |
|
|
| |
| 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}""" |
|
|
| |
| 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 |
|
|
| |
| 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"]) |
