Graph RAG Reference Engine (Python)

A standalone, runnable Python implementation showing AI engineers how to build a Hybrid Graph RAG Engine combining Pydantic entity extraction, NetworkX topology traversals, and Reciprocal Rank Fusion (RRF).

import math
from typing import List, Dict, Tuple
from pydantic import BaseModel, Field
import networkx as nx

# 1. Pydantic Schemas for Knowledge Graph Nodes & Edges
class Node(BaseModel):
    id: str = Field(description="Unique node identifier (e.g., 'PR_412')")
    type: str = Field(description="Category: Service, Developer, PullRequest, Issue")
    properties: Dict[str, str] = Field(default_factory=dict)

class Edge(BaseModel):
    source: str
    target: str
    relation: str
    weight: float = 1.0

# 2. Hybrid Graph RAG Retrieval Engine
class GraphRAGEngine:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.vector_store: Dict[str, str] = {}  # Mock vector index (node_id -> text)

    def add_knowledge_triple(self, source_node: Node, target_node: Node, edge: Edge):
        self.graph.add_node(source_node.id, type=source_node.type, **source_node.properties)
        self.graph.add_node(target_node.id, type=target_node.type, **target_node.properties)
        self.graph.add_edge(edge.source, edge.target, relation=edge.relation, weight=edge.weight)

    def reciprocal_rank_fusion(self, vector_ranks: List[str], graph_ranks: List[str], k: int = 60) -> List[Tuple[str, float]]:
        """Reciprocal Rank Fusion (RRF) scoring algorithm"""
        scores: Dict[str, float] = {}
        for rank, node_id in enumerate(vector_ranks, start=1):
            scores[node_id] = scores.get(node_id, 0.0) + (1.0 / (k + rank))
        for rank, node_id in enumerate(graph_ranks, start=1):
            scores[node_id] = scores.get(node_id, 0.0) + (1.0 / (k + rank))
        return sorted(scores.items(), key=lambda x: x[1], reverse=True)

    def query_graph_rag(self, seed_node: str, depth: int = 2) -> List[str]:
        """Multi-hop localized graph traversal"""
        if seed_node not in self.graph:
            return []
        subgraph_nodes = nx.single_source_shortest_path_length(self.graph, seed_node, cutoff=depth)
        return list(subgraph_nodes.keys())

# Example Execution
if __name__ == "__main__":
    engine = GraphRAGEngine()
    # Populate knowledge triples
    engine.add_knowledge_triple(
        Node(id="AUTH_API", type="Service", properties={"name": "Auth API"}),
        Node(id="PR_412", type="PullRequest", properties={"title": "Custom LRU Cache"}),
        Edge(source="PR_412", target="AUTH_API", relation="MODIFIES", weight=0.95)
    )
    
    # Vector & Graph Candidate Retrieval
    vector_candidates = ["PR_412", "JIRA_SEC_402"]
    graph_candidates = engine.query_graph_rag("AUTH_API", depth=2)
    
    # RRF Fusion Ranking
    fused_results = engine.reciprocal_rank_fusion(vector_candidates, graph_candidates)
    print("Fused Rank Context Queue:", fused_results)