graph_rag.py

"""
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),
        }