E2E GraphRAG engine; multi-hop relational paths + Louvain community detection

e2e_graphrag.py

import networkx as nx
import numpy as np
import math
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Any

# --- core data structures ---

@dataclass
class Entity:
    name: str
    type: str
    description: str
    embedding: np.ndarray = field(default_factory=lambda: np.zeros(64))

@dataclass
class Relationship:
    source: str
    target: str
    description: str
    weight: float = 1.0

@dataclass
class CommunityReport:
    community_id: int
    title: str
    summary: str
    findings: List[str]


# --- mock infrastructure providers ---
# replace these methods with actual api calls (e.g., openai or ollama) in production

class MockLLMClient:
    """simulates structured llm extraction and generation."""
    
    def extract_triplets(self, text: str) -> Tuple[List[Entity], List[Relationship]]:
        # simulates extracting structured entity/relationship primitives from raw data
        # in practice, you would use instructor, outlines, or json-mode prompting here
        if "quantum" in text.lower():
            entities = [
                Entity("Quantum Computing", "Technology", "a computing paradigm using quantum-mechanical phenomena."),
                Entity("Qubit", "Component", "the basic unit of quantum information."),
                Entity("Superposition", "Principle", "the ability of a quantum system to be in multiple states simultaneously.")
            ]
            relationships = [
                Relationship("Quantum Computing", "Qubit", "utilizes qubits as core hardware infrastructure.", 5.0),
                Relationship("Qubit", "Superposition", "depends on superposition to maintain quantum states.", 4.0),
                Relationship("Quantum Computing", "Superposition", "leverages superposition for parallel computation operations.", 4.0)
            ]
            return entities, relationships
        else:
            entities = [
                Entity("Artificial Intelligence", "Field", "the simulation of human intelligence processes by machines."),
                Entity("Neural Network", "Architecture", "computational models inspired by biological brain networks."),
                Entity("Deep Learning", "Subfield", "a subset of machine learning based on artificial neural networks.")
            ]
            relationships = [
                Relationship("Artificial Intelligence", "Neural Network", "uses neural networks to process complex data configurations.", 5.0),
                Relationship("Deep Learning", "Neural Network", "is built primarily on multi-layered neural network stacks.", 5.0),
                Relationship("Artificial Intelligence", "Deep Learning", "drives advanced capabilities via deep learning breakthroughs.", 4.0)
            ]
            return entities, relationships

    def generate(self, prompt: str) -> str:
        # simulates context-grounded final answer synthesis
        return f"[llm response grounded via graph context]\nprocessed input request based on provided structure."


class MockEmbeddingClient:
    """simulates generation of dense semantic vectors."""
    def get_embedding(self, text: str) -> np.ndarray:
        # returns a deterministic dummy vector for execution stability
        np.random.seed(sum(ord(c) for c in text) % 1000)
        vec = np.random.randn(64)
        return vec / np.linalg.norm(vec)


# --- the graphrag system engine ---

class GraphRAGIndexer:
    """handles ingestion, entity extraction, graph compilation, and community building."""
    def __init__(self, llm: MockLLMClient, embedder: MockEmbeddingClient):
        self.llm = llm
        self.embedder = embedder
        self.graph = nx.Graph()
        self.entities: Dict[str, Entity] = {}
        self.community_reports: Dict[int, CommunityReport] = {}

    def ingest_document(self, text: str):
        # step 1: extract entity and relationship primitives via llm
        extracted_entities, extracted_relationships = self.llm.extract_triplets(text)
        
        # step 2: upsert entities into global index and update semantic embeddings
        for entity in extracted_entities:
            if entity.name not in self.entities:
                entity.embedding = self.embedder.get_embedding(f"{entity.name}: {entity.description}")
                self.entities[entity.name] = entity
                self.graph.add_node(entity.name, type=entity.type, description=entity.description)
            else:
                # merge descriptions if entity already exists
                self.entities[entity.name].description += f" | {entity.description}"
                
        # step 3: upsert edges with topological weights
        for rel in extracted_relationships:
            if self.graph.has_edge(rel.source, rel.target):
                self.graph[rel.source][rel.target]['weight'] += rel.weight
                self.graph[rel.source][rel.target]['description'] += f" | {rel.description}"
            else:
                self.graph.add_edge(
                    rel.source, 
                    rel.target, 
                    weight=rel.weight, 
                    description=rel.description
                )

    def build_communities(self):
        """uncovers network community clusters and runs partition-level summarization."""
        # run louvain modularity optimization to find highly-connected entity clusters
        communities = nx.community.louvain_communities(self.graph, weight='weight')
        
        for idx, community in enumerate(communities):
            # compile local text descriptors for this specific community subgraph
            subgraph_details = []
            for node in community:
                entity = self.entities[node]
                subgraph_details.append(f"entity: {entity.name} ({entity.type}) - {entity.description}")
                
            # fetch internal edges within this specific community
            subgraph = self.graph.subgraph(community)
            for u, v, data in subgraph.edges(data=True):
                subgraph_details.append(f"relation: {u} -> {v}: {data['description']}")
                
            # format a map-reduce style summary prompt for the partition
            community_text = "\n".join(subgraph_details)
            summary_prompt = f"summarize the high-level themes of this semantic community block:\n{community_text}"
            
            # store compiled structural reports
            self.community_reports[idx] = CommunityReport(
                community_id=idx,
                title=f"semantic cluster group {idx}",
                summary=self.llm.generate(summary_prompt),
                findings=[f"key structural insight from cluster partition {idx}"]
            )


class GraphRAGQueryEngine:
    """contains both the localized vector-neighborhood search and macro map-reduce global engines."""
    def __init__(self, indexer: GraphRAGIndexer):
        self.indexer = indexer

    def local_search(self, query: str, top_k_entities: int = 2) -> str:
        """resolves queries by finding closest entities and expanding to their immediate neighbors."""
        query_emb = self.indexer.embedder.get_embedding(query)
        
        # calculate cosine similarity across all stored entity vectors
        scores = []
        for name, entity in self.indexer.entities.items():
            sim = np.dot(query_emb, entity.embedding)
            scores.append((name, sim))
            
        scores.sort(key=lambda x: x[1], reverse=True)
        target_entities = [name for name, _ in scores[:top_k_entities]]
        
        # traverse graph to harvest the surrounding multi-hop neighborhood context
        context_parts = []
        visited_edges = set()
        
        for entity_name in target_entities:
            entity = self.indexer.entities[entity_name]
            context_parts.append(f"[entity context] {entity.name}: {entity.description}")
            
            # grab 1-hop connected neighbors
            for neighbor in self.indexer.graph.neighbors(entity_name):
                edge_key = tuple(sorted((entity_name, neighbor)))
                if edge_key not in visited_edges:
                    visited_edges.add(edge_key)
                    edge_data = self.indexer.graph[entity_name][neighbor]
                    context_parts.append(
                        f"[relational link] {entity_name} is linked to {neighbor}. connection detail: {edge_data['description']}"
                    )
                    
        # synthesize final localized answer using collected structural context
        final_prompt = f"answer the query based strictly on the extracted sub-graph neighborhood:\n\ncontext:\n" + "\n".join(context_parts) + f"\n\nquery: {query}"
        return self.indexer.llm.generate(final_prompt)

    def global_search(self, query: str) -> str:
        """resolves high-level thematic queries via map-reduce across precompiled community reports."""
        # map phase: extract insights targeting the query from every single distinct community report
        intermediate_answers = []
        for comm_id, report in self.indexer.community_reports.items():
            map_prompt = (
                f"review this community summary report and pull out details relevant to the query.\n"
                f"report title: {report.title}\nsummary: {report.summary}\nquery: {query}"
            )
            intermediate_answers.append(self.indexer.llm.generate(map_prompt))
            
        # reduce phase: combine scattered macro insights into one consolidated response
        reduce_prompt = (
            f"synthesize the following intermediate community reports into a unified global analysis:\n"
            + "\n".join(intermediate_answers) +
            f"\nfinal query objective: {query}"
        )
        return self.indexer.llm.generate(reduce_prompt)


# --- execution verification harness ---
if __name__ == "__main__":
    # initialize dependencies
    llm = MockLLMClient()
    embedder = MockEmbeddingClient()
    
    # initialize indexer
    indexer = GraphRAGIndexer(llm, embedder)
    
    # simulate raw text unstructured document corpus ingestion
    doc_1 = "quantum computing relies on building stable hardware architectures. at the center of this are qubits, which exploit the mechanics of superposition to evaluate multiple problem dimensions simultaneously."
    doc_2 = "modern artificial intelligence systems scale out capability using neural networks. deep learning represents a major subfield, chaining dense mathematical layers together to form predictive transformers."
    
    print("=== processing unstructured ingestion pipeline ===")
    indexer.ingest_document(doc_1)
    indexer.ingest_document(doc_2)
    print(f"graph compiled successfully. nodes: {list(indexer.graph.nodes)}")
    print(f"edges: {list(indexer.graph.edges(data=False))}\n")
    
    # compile community partitions
    print("=== partition analysis & community detection ===")
    indexer.build_communities()
    print(f"generated {len(indexer.community_reports)} distinct macro community reports.\n")
    
    # instantiate runtime query routers
    query_engine = GraphRAGQueryEngine(indexer)
    
    # execute local search query execution (specific, entity-targeted facts)
    print("=== executing query: local neighborhood search ===")
    local_res = query_engine.local_search("what core principles govern qubit operations?")
    print(local_res, "\n")
    
    # execute global search query execution (thematic, global multi-document abstractions)
    print("=== executing query: global community map-reduce ===")
    global_res = query_engine.global_search("compare the technical focus areas across the whole data archive.")
    print(global_res)

note: 2e3cb[3]