Upload folder using huggingface_hub

api.py

@@ -1,4 +1,11 @@
 import os
+# Configure thread pool limits BEFORE PyTorch/numpy are imported to conserve memory in Codespace
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+
 import gc
 import hashlib
 import re
@@ -21,6 +28,7 @@ from slowapi.errors import RateLimitExceeded
 from bhashini import translate_text, LANGUAGE_CODES
 from eligibility.engine import check_eligibility
 from neuro_symbolic import run_neuro_symbolic_pipeline
+from graph_rag import run_graph_rag
 from whatsapp.webhook import router as whatsapp_router
 from config import settings
 
@@ -39,9 +47,17 @@ embedding_model = None
 cross_encoder_model = None
 
 def get_embedding_model():
+    global embedding_model
+    if embedding_model is None:
+        logger.info("⏳ Lazily loading Nomic Embedding Model (768-dim)...")
+        embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
     return embedding_model
 
 def get_reranker_model():
+    global cross_encoder_model
+    if cross_encoder_model is None:
+        logger.info("⏳ Lazily loading Cross-Encoder Reranker Model...")
+        cross_encoder_model = CrossEncoder("cross-encoder/ettin-reranker-68m-v1", max_length=512, device='cpu', trust_remote_code=True)
     return cross_encoder_model
 
 # --- SECURE KEYS ---
@@ -55,12 +71,7 @@ groq_client = Groq(api_key=GROQ_API_KEY)
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global embedding_model, cross_encoder_model
-    logger.info("⏳ Pre-loading Nomic Embedding Model (768-dim)...")
-    embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
-    logger.info("⏳ Pre-loading Cross-Encoder Reranker Model...")
-    cross_encoder_model = CrossEncoder("cross-encoder/ettin-reranker-68m-v1", max_length=512, device='cpu', trust_remote_code=True)
-    logger.info("✅ Models pre-loaded successfully.")
+    logger.info("🚀 Lifespan started (lazy model loading enabled).")
     yield
 
 app = FastAPI(lifespan=lifespan)
@@ -484,26 +495,53 @@ async def rag_query(request: Request, query: SearchQuery):
                 if c.get('scheme_title')
             ]))
 
-        # --- STEP 6: GENERATE RESPONSE IN ENGLISH ---
-        # Note: We await the full block of text rather than stream it piece-by-piece, 
-        # because translation APIs require complete sentences to maintain grammar rules.
-        response = groq_client.chat.completions.create(
-            model="llama-3.3-70b-versatile",
-            messages=[
-                {"role": "system", "content": 
-                 "You are GovBridge AI. Answer using ONLY the provided context. "
-                 "Respond exclusively in English. "
-                 "Keep your answer under 200 words. Be precise, not exhaustive. "
-                 "Cite the scheme name. Format benefits as bullet points where applicable."
-                },
-                {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {english_query}"}
-            ],
-            temperature=0.1,
-            max_tokens=512, 
-            stream=False 
+        # --- STEP 6: GENERATE RESPONSE WITH GRAPH-RAG (Sprint 35) ---
+        # The Graph-RAG engine wraps LLM generation with optional autonomous
+        # graph traversal. If the LLM determines that legal relationship context
+        # would help, it invokes traverse_legal_graph via tool calling.
+        # Max 2 traversal iterations. Falls back to standard generation if no
+        # tool calls are made.
+        graph_rag_result = run_graph_rag(
+            query=english_query,
+            context_chunks=chunks if chunks else [],
+            groq_client=groq_client,
+            supabase_client=supabase,
+            embedding_model=get_embedding_model(),
+            user_language=user_language,
         )
         
-        english_response = response.choices[0].message.content or ""
+        english_response = graph_rag_result.get("response", "")
+        
+        # Log graph traversal metrics
+        if graph_rag_result.get("graph_traversals", 0) > 0:
+            logger.info(
+                f"🔗 Graph-RAG: {graph_rag_result['graph_traversals']} traversal(s), "
+                f"TOON context: {len(graph_rag_result.get('graph_context', ''))} chars"
+            )
+        
+        # If Graph-RAG returned additional sources from graph traversal, merge them
+        graph_sources = graph_rag_result.get("sources", [])
+        if graph_sources:
+            source_titles = list(set(source_titles + graph_sources))
+        
+        # Fallback: if Graph-RAG returned empty (e.g., LLM error), use direct generation
+        if not english_response:
+            response = groq_client.chat.completions.create(
+                model="llama-3.3-70b-versatile",
+                messages=[
+                    {"role": "system", "content": 
+                     "You are GovBridge AI. Answer using ONLY the provided context. "
+                     "Respond exclusively in English. "
+                     "Keep your answer under 200 words. Be precise, not exhaustive. "
+                     "Cite the scheme name. Format benefits as bullet points where applicable."
+                    },
+                    {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {english_query}"}
+                ],
+                temperature=0.1,
+                max_tokens=512, 
+                stream=False 
+            )
+            english_response = response.choices[0].message.content or ""
         
         # Write to cache after successful generation
         try:

graph_rag.py

@@ -0,0 +1,348 @@
+"""
+GovBridge India — Graph-RAG Autonomous Traversal Engine (Sprint 35)
+
+PROJECT INDRA Phase 4.0: Autonomous Graph-RAG
+
+This module implements a ReAct-style tool-calling loop that enables
+the LLM (Groq/Llama-3.3) to autonomously traverse the tensor_edges
+knowledge graph to enrich its RAG context.
+
+Pipeline:
+1. First LLM call includes a tool definition for `traverse_legal_graph`
+2. If the LLM invokes the tool, we intercept the call and execute
+   the Supabase `get_graph_neighborhood` RPC
+3. We apply CatRAG semantic pruning to filter irrelevant edges
+4. We serialize the relevant subgraph as TOON (Token-Oriented Object Notation)
+5. We feed the TOON context back to the LLM for final synthesis
+6. Max 2 traversal iterations. visited_nodes set prevents re-exploration.
+
+ARCHITECTURAL LAW:
+  - No LangChain, no LlamaIndex, no framework bloat.
+  - Pure recursive Python. Zero dependencies beyond groq + supabase.
+  - LLM NEVER does arithmetic. It only traverses and synthesizes prose.
+"""
+
+import json
+import logging
+from typing import Any, Optional
+from groq import Groq
+from supabase import Client
+from sentence_transformers import SentenceTransformer
+import numpy as np
+
+logger = logging.getLogger("govbridge.graph_rag")
+
+# ── Tool Definition (Groq/OpenAI Tool Calling Format) ────────────
+TRAVERSE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "traverse_legal_graph",
+        "description": (
+            "Traverse the legal knowledge graph to find related gazette documents, "
+            "amendments, superseding acts, and cross-references for a given document node. "
+            "Use this when the user's question involves legal relationships, document "
+            "history, amendments, or connections between government notifications."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "anchor_id": {
+                    "type": "string",
+                    "description": (
+                        "UUID of the gazette document node to explore. "
+                        "This should be extracted from the context provided."
+                    )
+                }
+            },
+            "required": ["anchor_id"],
+        },
+    },
+}
+
+# ── TOON Serializer (Token-Oriented Object Notation) ──────────────
+def graph_to_toon(edges: list[dict[str, Any]]) -> str:
+    """
+    Flatten graph edges into a compact Markdown table.
+    
+    This saves ~60% of context tokens vs raw JSON serialization.
+    The LLM reads tabular data efficiently.
+    
+    Returns:
+        Markdown table string or empty string if no edges.
+    """
+    if not edges:
+        return ""
+    
+    lines = ["| Source | Relation | Target | Hop |", "|--------|----------|--------|-----|"]
+    
+    for edge in edges:
+        source = (edge.get("source_title") or "Unknown")[:50]
+        target = (edge.get("target_title") or "Unresolved")[:50]
+        relation = (edge.get("edge_type") or "cross_references").upper()
+        hop = edge.get("hop_depth", 1)
+        lines.append(f"| {source} | {relation} | {target} | {hop} |")
+    
+    return "\n".join(lines)
+
+
+# ── CatRAG Semantic Pruning ───────────────────────────────────────
+def prune_edges_by_relevance(
+    query: str,
+    edges: list[dict[str, Any]],
+    embedding_model: SentenceTransformer,
+    top_k: int = 15,
+) -> list[dict[str, Any]]:
+    """
+    Semantic pruning: Compare user query to edge context and return
+    ONLY the top_k most relevant edges.
+    
+    Uses cosine similarity between the query embedding and a
+    composite text of each edge (source_title + relation + target_title).
+    
+    Args:
+        query: The user's original question.
+        edges: Raw edges from get_graph_neighborhood RPC.
+        embedding_model: The Nomic embedding model instance.
+        top_k: Maximum number of edges to return (default 15).
+    
+    Returns:
+        Pruned list of the top_k most relevant edges.
+    """
+    if len(edges) <= top_k:
+        return edges
+    
+    # Build composite text for each edge
+    edge_texts = []
+    for edge in edges:
+        text = (
+            f"{edge.get('source_title', '')} "
+            f"{edge.get('edge_type', '')} "
+            f"{edge.get('target_title', '')}"
+        )
+        edge_texts.append(text)
+    
+    # Encode query and edge texts
+    query_emb = embedding_model.encode(
+        f"search_query: {query}",
+        normalize_embeddings=True,
+    )
+    edge_embs = embedding_model.encode(
+        [f"search_document: {t}" for t in edge_texts],
+        normalize_embeddings=True,
+        batch_size=32,
+        show_progress_bar=False,
+    )
+    
+    # Compute cosine similarities (vectors are normalized, so dot product = cosine)
+    similarities = np.dot(edge_embs, query_emb)
+    
+    # Get top_k indices by descending similarity
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    return [edges[int(i)] for i in top_indices]
+
+
+# ── Main ReAct Loop ───────────────────────────────────────────────
+def run_graph_rag(
+    query: str,
+    context_chunks: list[dict[str, Any]],
+    groq_client: Groq,
+    supabase_client: Client,
+    embedding_model: SentenceTransformer,
+    user_language: str = "english",
+) -> dict[str, Any]:
+    """
+    Execute the Graph-RAG ReAct loop.
+    
+    This function wraps the standard RAG response generation with
+    an optional graph traversal step. If the LLM determines that
+    graph context would help answer the query, it invokes the
+    traverse_legal_graph tool, and we feed the results back.
+    
+    Args:
+        query: The user's question (in English).
+        context_chunks: Pre-fetched RAG chunks from hybrid search.
+        groq_client: Initialized Groq client.
+        supabase_client: Initialized Supabase client.
+        embedding_model: The Nomic embedding model.
+        user_language: Target language for the response.
+    
+    Returns:
+        Dict with:
+        - "response": str — The final answer text.
+        - "sources": list[str] — Source document titles.
+        - "graph_traversals": int — Number of graph hops performed.
+        - "graph_context": str — TOON-formatted graph context (if any).
+    """
+    MAX_ITERATIONS = 2
+    visited_nodes: set[str] = set()
+    graph_context_toon = ""
+    graph_traversals = 0
+    
+    # Build initial context from RAG chunks
+    rag_context = "\n\n---\n\n".join([
+        f"[Document: {str(c.get('scheme_title', 'Unknown'))}] "
+        f"(ID: {c.get('id', 'N/A')})\n{str(c.get('chunk_text', ''))}"
+        for c in context_chunks
+    ])
+    
+    source_titles = list(set([
+        str(c.get('scheme_title', 'Unknown'))
+        for c in context_chunks
+        if c.get('scheme_title')
+    ]))
+    
+    # System prompt with graph-awareness
+    system_prompt = (
+        "You are GovBridge AI, India's premier government scheme assistant. "
+        "Answer using ONLY the provided context. "
+        "Respond exclusively in English. "
+        "Keep your answer under 200 words. Be precise, not exhaustive. "
+        "Cite the scheme or gazette name. Format benefits as bullet points.\n\n"
+        "You have access to a legal knowledge graph. If the user's question involves "
+        "relationships between gazette documents, amendments, superseding acts, or "
+        "cross-references, you may call the `traverse_legal_graph` tool with the "
+        "document UUID from the context to explore related documents. "
+        "Only call the tool if the graph context would genuinely help answer the question."
+    )
+    
+    # Build messages
+    messages: list[dict[str, Any]] = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"Context:\n{rag_context}\n\nQuestion: {query}"},
+    ]
+    
+    # ── ReAct Loop ────────────────────────────────────────────
+    for iteration in range(MAX_ITERATIONS + 1):  # +1 for final synthesis
+        try:
+            response = groq_client.chat.completions.create(
+                model="llama-3.3-70b-versatile",
+                messages=messages,
+                tools=[TRAVERSE_TOOL] if iteration < MAX_ITERATIONS else None,
+                tool_choice="auto" if iteration < MAX_ITERATIONS else None,
+                temperature=0.1,
+                max_tokens=512,
+                stream=False,
+            )
+        except Exception as e:
+            logger.error(f"Graph-RAG LLM call failed: {e}")
+            return {
+                "response": "",
+                "sources": source_titles,
+                "graph_traversals": graph_traversals,
+                "graph_context": graph_context_toon,
+                "error": str(e),
+            }
+        
+        choice = response.choices[0]
+        
+        # ── Case 1: LLM returns a direct answer (no tool call) ──
+        if choice.finish_reason != "tool_calls" or not choice.message.tool_calls:
+            return {
+                "response": choice.message.content or "",
+                "sources": source_titles,
+                "graph_traversals": graph_traversals,
+                "graph_context": graph_context_toon,
+            }
+        
+        # ── Case 2: LLM invoked traverse_legal_graph ────────────
+        tool_call = choice.message.tool_calls[0]
+        
+        if tool_call.function.name != "traverse_legal_graph":
+            # Unknown tool — force synthesis
+            logger.warning(f"Unknown tool called: {tool_call.function.name}")
+            break
+        
+        try:
+            args = json.loads(tool_call.function.arguments)
+            anchor_id = args.get("anchor_id", "")
+        except (json.JSONDecodeError, AttributeError):
+            logger.warning("Failed to parse tool arguments")
+            break
+        
+        # ── Loop Breaker: Check visited set ─────────────────────
+        if anchor_id in visited_nodes:
+            # Append the assistant's tool call message
+            messages.append(choice.message.model_dump())
+            messages.append({
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "content": "SYSTEM OVERRIDE: Node already explored. Synthesize your answer from existing context.",
+            })
+            continue
+        
+        visited_nodes.add(anchor_id)
+        graph_traversals += 1
+        logger.info(f"🔗 Graph-RAG traversal #{graph_traversals}: anchor={anchor_id}")
+        
+        # ── Execute graph traversal via Supabase RPC ────────────
+        try:
+            result = supabase_client.rpc("get_graph_neighborhood", {
+                "anchor_id": anchor_id,
+                "max_depth": 1,
+                "edge_limit": 50,
+            }).execute()
+            raw_edges = result.data or []
+        except Exception as e:
+            logger.error(f"Graph RPC failed: {e}")
+            raw_edges = []
+        
+        # ── CatRAG Semantic Pruning ─────────────────────────────
+        if raw_edges:
+            pruned_edges = prune_edges_by_relevance(
+                query, raw_edges, embedding_model, top_k=15
+            )
+        else:
+            pruned_edges = []
+        
+        # ── TOON Serialization ──────────────────────────────────
+        toon_result = graph_to_toon(pruned_edges)
+        graph_context_toon = toon_result  # Store for response metadata
+        
+        # ── Feed graph context back to LLM ──────────────────────
+        if toon_result:
+            tool_response = (
+                f"Graph neighborhood for document {anchor_id}:\n\n"
+                f"{toon_result}\n\n"
+                f"({len(pruned_edges)} relevant relationships found out of "
+                f"{len(raw_edges)} total edges.)"
+            )
+        else:
+            tool_response = (
+                f"No graph relationships found for document {anchor_id}. "
+                "Synthesize your answer from the existing context."
+            )
+        
+        # Append tool call and response to messages
+        messages.append(choice.message.model_dump())
+        messages.append({
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "content": tool_response,
+        })
+    
+    # ── Fallback: if loop exhausted without a final answer ────
+    # Make one last call WITHOUT tools to force synthesis
+    try:
+        final_response = groq_client.chat.completions.create(
+            model="llama-3.3-70b-versatile",
+            messages=messages,
+            temperature=0.1,
+            max_tokens=512,
+            stream=False,
+        )
+        return {
+            "response": final_response.choices[0].message.content or "",
+            "sources": source_titles,
+            "graph_traversals": graph_traversals,
+            "graph_context": graph_context_toon,
+        }
+    except Exception as e:
+        logger.error(f"Graph-RAG final synthesis failed: {e}")
+        return {
+            "response": "",
+            "sources": source_titles,
+            "graph_traversals": graph_traversals,
+            "graph_context": graph_context_toon,
+            "error": str(e),
+        }