""" 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] def clean_message_for_payload(message) -> dict: """Helper to clean ChatCompletionMessage so it is accepted by Groq API.""" msg_dict = { "role": message.role, "content": message.content, } if getattr(message, "tool_calls", None): msg_dict["tool_calls"] = [ { "id": tc.id, "type": tc.type, "function": { "name": tc.function.name, "arguments": tc.function.arguments, } } for tc in message.tool_calls ] return msg_dict # ── 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(clean_message_for_payload(choice.message)) 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(clean_message_for_payload(choice.message)) 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), }