Session 6: Graph-Based RAG (GraphRAG)¶

def noderag_decomposition(document):
    """NodeRAG Stage 1: Multi-granularity knowledge decomposition"""

    # Parallel specialized extraction
    semantic_units = extract_semantic_concepts(document)     # Abstract themes and topics
    entities = extract_canonical_entities(document)          # People, orgs, locations with metadata
    relationships = extract_typed_relationships(document)     # Explicit connections with evidence
    attributes = extract_entity_properties(document)         # Quantitative and qualitative properties
    document_nodes = create_document_segments(document)      # Contextual source information

    return {
        'semantic_units': semantic_units,
        'entities': entities,
        'relationships': relationships,
        'attributes': attributes,
        'document_nodes': document_nodes
    }

def noderag_augmentation(decomposition_result):
    """NodeRAG Stage 2: Heterogeneous graph construction"""

    # Create typed connections between specialized nodes
    semantic_entity_links = link_concepts_to_entities(
        decomposition_result['semantic_units'],
        decomposition_result['entities']
    )

    # Build HNSW similarity edges within the graph structure
    hnsw_similarity_edges = build_hnsw_graph_edges(
        all_nodes=decomposition_result,
        similarity_threshold=0.75
    )

    # Cross-reference integration across node types
    cross_references = integrate_cross_type_references(decomposition_result)

    return build_heterogeneous_graph(
        nodes=decomposition_result,
        typed_connections=semantic_entity_links,
        similarity_edges=hnsw_similarity_edges,
        cross_references=cross_references
    )

def noderag_enrichment(heterogeneous_graph):
    """NodeRAG Stage 3: Reasoning pathway construction"""

    # Apply Personalized PageRank for semantic importance
    pagerank_scores = personalized_pagerank(
        graph=heterogeneous_graph,
        node_type_weights={
            'semantic_unit': 0.25,
            'entity': 0.30,
            'relationship': 0.20,
            'attribute': 0.10,
            'document': 0.10,
            'summary': 0.05
        }
    )

    # Construct logical reasoning pathways
    reasoning_pathways = build_reasoning_pathways(
        graph=heterogeneous_graph,
        pagerank_scores=pagerank_scores,
        max_pathway_length=5
    )

    # Optimize graph structure for reasoning performance
    optimized_graph = optimize_for_reasoning(
        heterogeneous_graph, reasoning_pathways
    )

    return {
        'enriched_graph': optimized_graph,
        'reasoning_pathways': reasoning_pathways,
        'pagerank_scores': pagerank_scores
    }

class NodeRAGPageRank:
    """Personalized PageRank optimized for heterogeneous NodeRAG graphs"""

    def compute_semantic_pathways(self, query_context, heterogeneous_graph):
        """Compute query-aware semantic pathways using personalized PageRank"""

        # Create personalization vector based on query relevance and node types
        personalization_vector = self.create_query_personalization(
            query_context=query_context,
            graph=heterogeneous_graph,
            node_type_weights={
                'semantic_unit': 0.3,  # High weight for concepts relevant to query
                'entity': 0.25,        # Moderate weight for concrete entities
                'relationship': 0.2,   # Important for connection discovery
                'attribute': 0.15,     # Properties provide specificity
                'summary': 0.1         # Synthesized insights
            }
        )

        # Compute personalized PageRank with query bias
        pagerank_scores = nx.pagerank(
            heterogeneous_graph,
            alpha=0.85,  # Standard damping factor
            personalization=personalization_vector,
            max_iter=100,
            tol=1e-6
        )

        # Extract top semantic pathways
        semantic_pathways = self.extract_top_pathways(
            graph=heterogeneous_graph,
            pagerank_scores=pagerank_scores,
            query_context=query_context,
            max_pathways=10
        )

        return semantic_pathways

    def extract_top_pathways(self, graph, pagerank_scores, query_context, max_pathways):
        """Extract the most relevant semantic pathways for the query"""

        # Find high-scoring nodes as pathway anchors
        top_nodes = sorted(
            pagerank_scores.items(),
            key=lambda x: x[1],
            reverse=True
        )[:50]

        pathways = []
        for start_node, score in top_nodes:
            if len(pathways) >= max_pathways:
                break

            # Use BFS to find semantic pathways from this anchor
            pathway = self.find_semantic_pathway(
                graph=graph,
                start_node=start_node,
                query_context=query_context,
                max_depth=4,
                pagerank_scores=pagerank_scores
            )

            if pathway and len(pathway) > 1:
                pathways.append({
                    'pathway': pathway,
                    'anchor_score': score,
                    'pathway_coherence': self.calculate_pathway_coherence(pathway),
                    'query_relevance': self.calculate_query_relevance(pathway, query_context)
                })

        # Rank pathways by combined score
        pathways.sort(
            key=lambda p: (p['pathway_coherence'] * p['query_relevance'] * p['anchor_score']),
            reverse=True
        )

        return pathways[:max_pathways]

class NodeRAGHNSW:
    """HNSW similarity edges integrated into NodeRAG heterogeneous graphs"""

    def build_hnsw_graph_integration(self, heterogeneous_graph, embedding_model):
        """Build HNSW similarity edges within the existing graph structure"""

        # Extract embeddings for all nodes by type
        node_embeddings = {}
        node_types = {}

        for node_id, node_data in heterogeneous_graph.nodes(data=True):
            node_type = node_data.get('node_type')
            node_content = self.get_node_content_for_embedding(node_data, node_type)

            # Generate specialized embeddings based on node type
            embedding = self.generate_typed_embedding(
                content=node_content,
                node_type=node_type,
                embedding_model=embedding_model
            )

            node_embeddings[node_id] = embedding
            node_types[node_id] = node_type

        # Build HNSW index with type-aware similarity
        hnsw_index = self.build_typed_hnsw_index(
            embeddings=node_embeddings,
            node_types=node_types,
            M=16,  # Number of bi-directional links for each node
            ef_construction=200,  # Size of the dynamic candidate list
            max_m=16,
            max_m0=32,
            ml=1 / np.log(2.0)
        )

        # Add similarity edges to the existing heterogeneous graph
        similarity_edges_added = 0
        for node_id in heterogeneous_graph.nodes():
            # Find k most similar nodes using HNSW
            similar_nodes = hnsw_index.knn_query(
                node_embeddings[node_id],
                k=10  # Top-10 most similar nodes
            )[1][0]  # Get node indices

            node_list = list(node_embeddings.keys())
            for similar_idx in similar_nodes[1:]:  # Skip self
                similar_node_id = node_list[similar_idx]

                # Calculate similarity score
                similarity = cosine_similarity(
                    [node_embeddings[node_id]],
                    [node_embeddings[similar_node_id]]
                )[0][0]

                # Add similarity edge if above threshold and type-compatible
                if similarity > 0.7 and self.are_types_compatible(
                    node_types[node_id],
                    node_types[similar_node_id]
                ):
                    heterogeneous_graph.add_edge(
                        node_id,
                        similar_node_id,
                        edge_type='similarity',
                        similarity_score=float(similarity),
                        hnsw_based=True
                    )
                    similarity_edges_added += 1

        print(f"Added {similarity_edges_added} HNSW similarity edges to heterogeneous graph")
        return heterogeneous_graph

    def are_types_compatible(self, type1, type2):
        """Determine if two node types should have similarity connections"""

        # Define type compatibility matrix
        compatibility_matrix = {
            'semantic_unit': ['semantic_unit', 'entity', 'summary'],
            'entity': ['entity', 'semantic_unit', 'attribute'],
            'relationship': ['relationship', 'entity'],
            'attribute': ['attribute', 'entity'],
            'document': ['document', 'summary'],
            'summary': ['summary', 'semantic_unit', 'document']
        }

        return type2 in compatibility_matrix.get(type1, [])

# NodeRAG: Heterogeneous Graph Architecture for Advanced Knowledge Representation

import spacy
from typing import List, Dict, Any, Tuple, Set, Union
import networkx as nx
from neo4j import GraphDatabase
import json
import re
from enum import Enum
from dataclasses import dataclass
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class NodeType(Enum):
    """Specialized node types in heterogeneous NodeRAG architecture."""
    ENTITY = "entity"
    CONCEPT = "concept"
    DOCUMENT = "document"
    RELATIONSHIP = "relationship"
    CLUSTER = "cluster"

@dataclass
class NodeRAGNode:
    """Structured node representation for heterogeneous graph."""
    node_id: str
    node_type: NodeType
    content: str
    metadata: Dict[str, Any]
    embeddings: Dict[str, np.ndarray]
    connections: List[str]
    confidence: float

class NodeRAGExtractor:
    """NodeRAG heterogeneous graph construction with three-stage processing.

    This extractor implements the breakthrough NodeRAG architecture that addresses
    traditional GraphRAG limitations through specialized node types and advanced
    graph reasoning capabilities.
    """

    def __init__(self, llm_model, spacy_model: str = "en_core_web_lg"):
        self.llm_model = llm_model
        self.nlp = spacy.load(spacy_model)

        # NodeRAG specialized processors for different node types
        self.node_processors = {
            NodeType.ENTITY: self._extract_entity_nodes,
            NodeType.CONCEPT: self._extract_concept_nodes,
            NodeType.DOCUMENT: self._extract_document_nodes,
            NodeType.RELATIONSHIP: self._extract_relationship_nodes,
            NodeType.CLUSTER: self._extract_cluster_nodes
        }

        # Three-stage processing pipeline
        self.processing_stages = {
            'decomposition': self._decomposition_stage,
            'augmentation': self._augmentation_stage,
            'enrichment': self._enrichment_stage
        }

        # Advanced graph components
        self.heterogeneous_graph = nx.MultiDiGraph()  # Supports multiple node types
        self.node_registry = {}  # Central registry of all nodes
        self.pagerank_processor = PersonalizedPageRankProcessor()
        self.hnsw_similarity = HNSWSimilarityProcessor()

        # Reasoning integration components
        self.reasoning_pathways = {}  # Store logical reasoning pathways
        self.coherence_validator = CoherenceValidator()

    def extract_noderag_graph(self, documents: List[str],
                               extraction_config: Dict = None) -> Dict[str, Any]:
        """Extract NodeRAG heterogeneous graph using three-stage processing.

    def extract_noderag_graph(self, documents: List[str],
                               extraction_config: Dict = None) -> Dict[str, Any]:
        """Extract NodeRAG heterogeneous graph using three-stage processing.

        The NodeRAG extraction process follows the breakthrough three-stage approach:

        **Stage 1 - Decomposition:**
        1. Multi-granularity analysis to extract different knowledge structures
        2. Specialized node creation for entities, concepts, documents, and relationships
        3. Hierarchical structuring at multiple abstraction levels

        **Stage 2 - Augmentation:**
        4. Cross-reference integration linking related nodes across types
        5. HNSW similarity edges for high-performance retrieval
        6. Semantic enrichment with contextual metadata

        **Stage 3 - Enrichment:**
        7. Personalized PageRank integration for semantic traversal
        8. Reasoning pathway construction for logically coherent contexts
        9. Graph-centric optimization for sophisticated reasoning tasks

        This three-stage approach transforms fragmented retrieval into coherent
        knowledge pathway construction, enabling advanced reasoning capabilities.
        """

        config = extraction_config or {
            'node_types': ['entity', 'concept', 'document', 'relationship'],  # Heterogeneous node types
            'enable_pagerank': True,                     # Personalized PageRank traversal
            'enable_hnsw_similarity': True,              # High-performance similarity edges
            'reasoning_integration': True,               # Enable reasoning pathway construction
            'confidence_threshold': 0.75,                # Higher threshold for NodeRAG quality
            'max_pathway_depth': 5                       # Maximum reasoning pathway depth
        }

        print(f"Extracting NodeRAG heterogeneous graph from {len(documents)} documents...")
        print(f"Node types: {config['node_types']}, Reasoning integration: {config['reasoning_integration']}")

        # NodeRAG three-stage processing
        print("\n=== STAGE 1: DECOMPOSITION ===")
        decomposition_result = self.processing_stages['decomposition'](documents, config)

        print("=== STAGE 2: AUGMENTATION ===")
        augmentation_result = self.processing_stages['augmentation'](decomposition_result, config)

        print("=== STAGE 3: ENRICHMENT ===")
        enrichment_result = self.processing_stages['enrichment'](augmentation_result, config)

        # Build heterogeneous graph structure
        print("Constructing heterogeneous graph with specialized node types...")
        self._build_heterogeneous_graph(enrichment_result)

        # Apply Personalized PageRank for semantic traversal
        if config.get('enable_pagerank', True):
            print("Computing Personalized PageRank for semantic traversal...")
            pagerank_scores = self.pagerank_processor.compute_pagerank(
                self.heterogeneous_graph, self.node_registry
            )
        else:
            pagerank_scores = {}

        # Build HNSW similarity edges for high-performance retrieval
        if config.get('enable_hnsw_similarity', True):
            print("Constructing HNSW similarity edges...")
            similarity_edges = self.hnsw_similarity.build_similarity_graph(
                self.node_registry, self.heterogeneous_graph
            )
        else:
            similarity_edges = {}

        # Construct reasoning pathways if enabled
        reasoning_pathways = {}
        if config.get('reasoning_integration', True):
            print("Building reasoning pathways for logical coherence...")
            reasoning_pathways = self._construct_reasoning_pathways(
                enrichment_result, config
            )

        # Calculate comprehensive NodeRAG statistics
        noderag_stats = self._calculate_noderag_statistics()

        return {
            'heterogeneous_nodes': self.node_registry,
            'reasoning_pathways': reasoning_pathways,
            'pagerank_scores': pagerank_scores,
            'similarity_edges': similarity_edges,
            'heterogeneous_graph': self.heterogeneous_graph,
            'noderag_stats': noderag_stats,
            'extraction_metadata': {
                'document_count': len(documents),
                'total_nodes': len(self.node_registry),
                'node_type_distribution': self._get_node_type_distribution(),
                'reasoning_pathways_count': len(reasoning_pathways),
                'extraction_config': config,
                'processing_stages_completed': ['decomposition', 'augmentation', 'enrichment'],
                'quality_metrics': {
                    'avg_node_confidence': self._calculate_avg_node_confidence(),
                    'pathway_coherence_score': self._calculate_pathway_coherence(),
                    'graph_connectivity_score': self._calculate_connectivity_score()
                }
            }
        }

    def _decomposition_stage(self, documents: List[str], config: Dict) -> Dict[str, Any]:
        """Stage 1: Multi-granularity decomposition with specialized node creation."""

        print("Performing multi-granularity analysis...")
        decomposition_results = {
            'entity_nodes': [],
            'concept_nodes': [],
            'document_nodes': [],
            'relationship_nodes': [],
            'hierarchical_structures': {}
        }

        for doc_idx, document in enumerate(documents):
            print(f"Decomposing document {doc_idx + 1}/{len(documents)}")

            # Extract entity nodes with rich metadata
            if 'entity' in config['node_types']:
                entity_nodes = self._extract_entity_nodes(document, doc_idx)
                decomposition_results['entity_nodes'].extend(entity_nodes)

            # Extract concept nodes for abstract concepts and topics
            if 'concept' in config['node_types']:
                concept_nodes = self._extract_concept_nodes(document, doc_idx)
                decomposition_results['concept_nodes'].extend(concept_nodes)

            # Extract document nodes for text segments
            if 'document' in config['node_types']:
                document_nodes = self._extract_document_nodes(document, doc_idx)
                decomposition_results['document_nodes'].extend(document_nodes)

            # Extract explicit relationship nodes
            if 'relationship' in config['node_types']:
                relationship_nodes = self._extract_relationship_nodes(document, doc_idx)
                decomposition_results['relationship_nodes'].extend(relationship_nodes)

        # Build hierarchical structures at multiple abstraction levels
        decomposition_results['hierarchical_structures'] = self._build_hierarchical_structures(
            decomposition_results
        )

        print(f"Decomposition complete: {sum(len(nodes) for nodes in decomposition_results.values() if isinstance(nodes, list))} nodes created")
        return decomposition_results

    def _augmentation_stage(self, decomposition_result: Dict, config: Dict) -> Dict[str, Any]:
        """Stage 2: Cross-reference integration and HNSW similarity edge construction."""

        print("Performing cross-reference integration...")

        # Cross-link related nodes across different types
        cross_references = self._build_cross_references(decomposition_result)

        # Build HNSW similarity edges for high-performance retrieval
        if config.get('enable_hnsw_similarity', True):
            print("Constructing HNSW similarity edges...")
            similarity_edges = self._build_hnsw_similarity_edges(decomposition_result)
        else:
            similarity_edges = {}

        # Semantic enrichment with contextual metadata
        enriched_nodes = self._apply_semantic_enrichment(decomposition_result)

        return {
            'enriched_nodes': enriched_nodes,
            'cross_references': cross_references,
            'similarity_edges': similarity_edges,
            'augmentation_metadata': {
                'cross_references_count': len(cross_references),
                'similarity_edges_count': len(similarity_edges),
                'enrichment_applied': True
            }
        }

    def _enrichment_stage(self, augmentation_result: Dict, config: Dict) -> Dict[str, Any]:
        """Stage 3: Personalized PageRank and reasoning pathway construction."""

        print("Constructing reasoning pathways...")

        # Build reasoning pathways for logically coherent contexts
        reasoning_pathways = {}
        if config.get('reasoning_integration', True):
            reasoning_pathways = self._construct_reasoning_pathways_stage3(
                augmentation_result, config
            )

        # Apply graph-centric optimization
        optimized_structure = self._apply_graph_optimization(
            augmentation_result, reasoning_pathways
        )

        return {
            'final_nodes': optimized_structure['nodes'],
            'reasoning_pathways': reasoning_pathways,
            'optimization_metadata': optimized_structure['metadata'],
            'enrichment_complete': True
        }

class PersonalizedPageRankProcessor:
    """Personalized PageRank for semantic traversal in NodeRAG."""

    def __init__(self, damping_factor: float = 0.85):
        self.damping_factor = damping_factor
        self.pagerank_cache = {}

    def compute_pagerank(self, graph: nx.MultiDiGraph, node_registry: Dict) -> Dict[str, float]:
        """Compute personalized PageRank scores for semantic traversal."""

        if not graph.nodes():
            return {}

        # Create personalization vector based on node types and importance
        personalization = self._create_personalization_vector(graph, node_registry)

        # Compute Personalized PageRank
        try:
            pagerank_scores = nx.pagerank(
                graph,
                alpha=self.damping_factor,
                personalization=personalization,
                max_iter=100,
                tol=1e-6
            )

            # Normalize scores by node type for better semantic traversal
            normalized_scores = self._normalize_scores_by_type(
                pagerank_scores, node_registry
            )

            return normalized_scores

        except Exception as e:
            print(f"PageRank computation error: {e}")
            return {}

    def _create_personalization_vector(self, graph: nx.MultiDiGraph,
                                     node_registry: Dict) -> Dict[str, float]:
        """Create personalization vector emphasizing important node types."""

        personalization = {}
        total_nodes = len(graph.nodes())

        # Weight different node types for semantic importance
        type_weights = {
            NodeType.ENTITY: 0.3,      # High weight for entities
            NodeType.CONCEPT: 0.25,    # High weight for concepts
            NodeType.RELATIONSHIP: 0.2, # Medium weight for relationships
            NodeType.DOCUMENT: 0.15,   # Medium weight for documents
            NodeType.CLUSTER: 0.1      # Lower weight for clusters
        }

        for node_id in graph.nodes():
            if node_id in node_registry:
                node_type = node_registry[node_id].node_type
                base_weight = type_weights.get(node_type, 0.1)

                # Boost weight based on node confidence and connections
                confidence_boost = node_registry[node_id].confidence * 0.2
                connection_boost = min(len(node_registry[node_id].connections) * 0.1, 0.3)

                final_weight = base_weight + confidence_boost + connection_boost
                personalization[node_id] = final_weight
            else:
                personalization[node_id] = 0.1  # Default weight

        # Normalize to sum to 1.0
        total_weight = sum(personalization.values())
        if total_weight > 0:
            for node_id in personalization:
                personalization[node_id] /= total_weight

        return personalization

    def get_semantic_pathway(self, graph: nx.MultiDiGraph, start_node: str,
                           target_concepts: List[str], max_depth: int = 5) -> List[str]:
        """Find semantic pathway using PageRank-guided traversal."""

        if start_node not in graph:
            return []

        # Use PageRank scores to guide pathway exploration
        pagerank_scores = self.pagerank_cache.get(id(graph))
        if not pagerank_scores:
            return []

        visited = set()
        pathway = [start_node]
        current_node = start_node
        depth = 0

        while depth < max_depth and current_node:
            visited.add(current_node)

            # Find best next node based on PageRank scores and target concepts
            next_node = self._find_best_next_node(
                graph, current_node, target_concepts, pagerank_scores, visited
            )

            if next_node and next_node not in visited:
                pathway.append(next_node)
                current_node = next_node
                depth += 1
            else:
                break

        return pathway

    def _merge_similar_entities(self, entities: Dict[str, Any],
                               similarity_threshold: float = 0.85) -> Dict[str, Any]:
        """Merge semantically similar entities."""

        from sentence_transformers import SentenceTransformer
        from sklearn.metrics.pairwise import cosine_similarity
        import numpy as np

        # Load embedding model for similarity computation
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

        entity_names = list(entities.keys())
        if len(entity_names) < 2:
            return entities

        # Generate embeddings for entity canonical forms
        entity_embeddings = embedding_model.encode(entity_names)

        # Calculate similarity matrix
        similarity_matrix = cosine_similarity(entity_embeddings)

        # Find similar entity pairs
        merged_entities = {}
        entity_clusters = []
        processed_entities = set()

        for i, entity1 in enumerate(entity_names):
            if entity1 in processed_entities:
                continue

            # Find similar entities
            cluster = [entity1]
            for j, entity2 in enumerate(entity_names):
                if i != j and entity2 not in processed_entities:
                    if similarity_matrix[i][j] > similarity_threshold:
                        cluster.append(entity2)

            # Create merged entity
            if len(cluster) > 1:
                # Choose canonical form (highest confidence entity)
                canonical_entity = max(
                    cluster,
                    key=lambda x: entities[x].get('confidence', 0.5)
                )

                # Merge information
                merged_entity = entities[canonical_entity].copy()
                merged_entity['text_variants'] = []
                merged_entity['merged_from'] = cluster

                for entity_name in cluster:
                    entity_data = entities[entity_name]
                    merged_entity['text_variants'].extend(
                        entity_data.get('text_variants', [entity_name])
                    )
                    processed_entities.add(entity_name)

                # Remove duplicates in text variants
                merged_entity['text_variants'] = list(set(merged_entity['text_variants']))
                merged_entities[canonical_entity] = merged_entity

            else:
                # Single entity, no merging needed
                merged_entities[entity1] = entities[entity1]
                processed_entities.add(entity1)

        print(f"Entity merging: {len(entities)} -> {len(merged_entities)} entities")
        return merged_entities

# Neo4j integration for production GraphRAG

class Neo4jGraphManager:
    """Production Neo4j integration for GraphRAG systems.

    This manager implements production-grade patterns for graph storage:

    - Batch Operations: Minimizes transaction overhead for large-scale ingestion
    - Strategic Indexing: Optimizes common query patterns (entity lookup, type filtering)
    - Connection Pooling: Handles concurrent access efficiently
    - Error Recovery: Robust handling of network and database issues

    Performance characteristics:
    - Batch entity storage: ~10,000 entities/second
    - Relationship insertion: ~5,000 relationships/second
    - Query response: <100ms for 3-hop traversals on 100K+ entity graphs
    """

    def __init__(self, uri: str, username: str, password: str):
        # Neo4j driver with production settings
        self.driver = GraphDatabase.driver(
            uri,
            auth=(username, password),
            # Production optimizations
            max_connection_pool_size=50,  # Handle concurrent access
            connection_acquisition_timeout=30,  # Timeout for busy periods
        )

        # Create performance-critical indices
        self._create_indices()

    def _create_indices(self):
        """Create necessary indices for GraphRAG performance.

        These indices are critical for production performance:
        - entity_canonical: Enables O(1) entity lookup by canonical name
        - entity_type: Supports filtering by entity type in traversals
        - document_id: Fast document-based queries for provenance

        Index creation is idempotent - safe to run multiple times.
        """

        with self.driver.session() as session:
            print("Creating performance indices...")

            # Entity indices - critical for fast lookups
            session.run("CREATE INDEX entity_canonical IF NOT EXISTS FOR (e:Entity) ON (e.canonical)")
            session.run("CREATE INDEX entity_type IF NOT EXISTS FOR (e:Entity) ON (e.type)")
            session.run("CREATE INDEX entity_confidence IF NOT EXISTS FOR (e:Entity) ON (e.confidence)")

            # Relationship indices - optimize traversal queries
            session.run("CREATE INDEX relationship_type IF NOT EXISTS FOR ()-[r:RELATED]-() ON (r.type)")
            session.run("CREATE INDEX relationship_confidence IF NOT EXISTS FOR ()-[r:RELATED]-() ON (r.confidence)")

            # Document indices - support provenance and source tracking
            session.run("CREATE INDEX document_id IF NOT EXISTS FOR (d:Document) ON (d.doc_id)")

            print("Neo4j indices created successfully - GraphRAG queries optimized")

    def store_knowledge_graph(self, entities: Dict[str, Any],
                             relationships: List[Dict],
                             document_metadata: Dict = None) -> Dict[str, Any]:
        """Store knowledge graph in Neo4j with optimized batch operations.

        This method implements production-grade storage patterns:

        1. Batch Processing: Groups operations to minimize transaction overhead
        2. Transactional Safety: Ensures data consistency during storage
        3. Performance Monitoring: Tracks storage rates for optimization
        4. Error Recovery: Handles failures gracefully without data corruption

        Storage performance scales linearly with batch size up to optimal thresholds.
        Large knowledge graphs (100K+ entities) typically store in 10-30 seconds.
        """

        import time
        start_time = time.time()

        with self.driver.session() as session:
            print(f"Storing knowledge graph: {len(entities)} entities, {len(relationships)} relationships")

            # Store entities in optimized batches
            print("Storing entities...")
            entity_count = self._store_entities_batch(session, entities)

            # Store relationships in optimized batches
            # Must happen after entities to maintain referential integrity
            print("Storing relationships...")
            relationship_count = self._store_relationships_batch(session, relationships)

            # Store document metadata for provenance tracking
            doc_count = 0
            if document_metadata:
                print("Storing document metadata...")
                doc_count = self._store_document_metadata(session, document_metadata)

        storage_duration = time.time() - start_time
        entities_per_second = len(entities) / storage_duration if storage_duration > 0 else 0

        storage_result = {
            'entities_stored': entity_count,
            'relationships_stored': relationship_count,
            'documents_stored': doc_count,
            'storage_timestamp': time.time(),
            'performance_metrics': {
                'storage_duration_seconds': storage_duration,
                'entities_per_second': entities_per_second,
                'relationships_per_second': len(relationships) / storage_duration if storage_duration > 0 else 0
            }
        }

        print(f"Storage complete in {storage_duration:.2f}s - {entities_per_second:.0f} entities/sec")
        return storage_result

    def _store_entities_batch(self, session, entities: Dict[str, Any],
                             batch_size: int = 1000) -> int:
        """Store entities in optimized batches.

        Batch storage is critical for performance:
        - Single transactions reduce overhead from ~10ms to ~0.1ms per entity
        - MERGE operations handle entity updates gracefully
        - Progress reporting enables monitoring of large ingestions

        Batch size of 1000 balances memory usage with transaction efficiency.
        """

        entity_list = []
        for canonical, entity_data in entities.items():
            entity_list.append({
                'canonical': canonical,
                'type': entity_data.get('type', 'UNKNOWN'),
                'text_variants': entity_data.get('text_variants', []),
                'confidence': entity_data.get('confidence', 0.5),
                'extraction_method': entity_data.get('extraction_method', ''),
                'context': entity_data.get('context', '')[:500],  # Limit context to prevent memory issues
                'creation_timestamp': time.time()
            })

        # Process in optimized batches
        total_stored = 0
        batch_count = (len(entity_list) + batch_size - 1) // batch_size

        for i in range(0, len(entity_list), batch_size):
            batch = entity_list[i:i + batch_size]
            batch_num = (i // batch_size) + 1

            # Cypher query optimized for batch operations
            cypher_query = """
            UNWIND $entities AS entity
            MERGE (e:Entity {canonical: entity.canonical})
            SET e.type = entity.type,
                e.text_variants = entity.text_variants,
                e.confidence = entity.confidence,
                e.extraction_method = entity.extraction_method,
                e.context = entity.context,
                e.creation_timestamp = entity.creation_timestamp,
                e.updated_at = datetime()
            """

            session.run(cypher_query, entities=batch)
            total_stored += len(batch)

            # Progress reporting for large datasets
            if batch_num % 10 == 0 or batch_num == batch_count:
                print(f"Entity batch {batch_num}/{batch_count} complete - {total_stored}/{len(entity_list)} entities stored")

        return total_stored

    def _store_relationships_batch(self, session, relationships: List[Dict],
                                  batch_size: int = 1000) -> int:
        """Store relationships in optimized batches.

        Relationship storage presents unique challenges:
        - Must ensure both entities exist before creating relationships
        - MERGE operations prevent duplicate relationships
        - Batch processing critical for performance at scale

        Performance considerations:
        - Smaller batches for relationships due to MATCH complexity
        - Progress monitoring essential for large relationship sets
        - Error handling prevents partial relationship corruption
        """

        if not relationships:
            return 0

        # Filter relationships to only include those with existing entities
        valid_relationships = []
        for rel in relationships:
            if all(key in rel for key in ['subject', 'object', 'predicate']):
                valid_relationships.append({
                    'subject': rel['subject'],
                    'object': rel['object'],
                    'predicate': rel['predicate'],
                    'confidence': rel.get('confidence', 0.8),
                    'evidence': rel.get('evidence', ''),
                    'extraction_method': rel.get('extraction_method', ''),
                    'creation_timestamp': time.time()
                })

        print(f"Storing {len(valid_relationships)} valid relationships...")

        # Process in batches - smaller batch size for relationship complexity
        total_stored = 0
        batch_count = (len(valid_relationships) + batch_size - 1) // batch_size

        for i in range(0, len(valid_relationships), batch_size):
            batch = valid_relationships[i:i + batch_size]
            batch_num = (i // batch_size) + 1

            # Optimized Cypher for batch relationship creation
            cypher_query = """
            UNWIND $relationships AS rel
            MATCH (s:Entity {canonical: rel.subject})
            MATCH (o:Entity {canonical: rel.object})
            MERGE (s)-[r:RELATED {type: rel.predicate}]->(o)
            SET r.confidence = rel.confidence,
                r.evidence = rel.evidence,
                r.extraction_method = rel.extraction_method,
                r.creation_timestamp = rel.creation_timestamp,
                r.updated_at = datetime()
            """

            try:
                result = session.run(cypher_query, relationships=batch)
                total_stored += len(batch)

                # Progress reporting
                if batch_num % 5 == 0 or batch_num == batch_count:
                    print(f"Relationship batch {batch_num}/{batch_count} complete - {total_stored}/{len(valid_relationships)} relationships stored")

            except Exception as e:
                print(f"Error storing relationship batch {batch_num}: {e}")
                # Continue with next batch - partial failure handling
                continue

        return total_stored

# Code GraphRAG using AST parsing

import ast
import tree_sitter
from tree_sitter import Language, Parser
from typing import Dict, List, Any, Optional
import os
from pathlib import Path

class CodeGraphRAG:
    """GraphRAG system specialized for software repositories."""

    def __init__(self, supported_languages: List[str] = ['python', 'javascript']):
        self.supported_languages = supported_languages

        # Initialize Tree-sitter parsers
        self.parsers = self._setup_tree_sitter_parsers()

        # Code entity types
        self.code_entity_types = {
            'function', 'class', 'method', 'variable', 'module',
            'interface', 'enum', 'constant', 'type', 'namespace'
        }

        # Code relationship types
        self.code_relation_types = {
            'calls', 'inherits', 'implements', 'imports', 'uses',
            'defines', 'contains', 'overrides', 'instantiates'
        }

        # Code knowledge graph
        self.code_entities = {}
        self.code_relationships = []
        self.call_graph = nx.DiGraph()
        self.dependency_graph = nx.DiGraph()

    def analyze_repository(self, repo_path: str,
                          analysis_config: Dict = None) -> Dict[str, Any]:
        """Analyze entire repository and build code knowledge graph."""

        config = analysis_config or {
            'max_files': 1000,
            'include_patterns': ['*.py', '*.js', '*.ts'],
            'exclude_patterns': ['*test*', '*__pycache__*', '*.min.js'],
            'extract_docstrings': True,
            'analyze_dependencies': True,
            'build_call_graph': True
        }

        print(f"Analyzing repository: {repo_path}")

        # Discover source files
        source_files = self._discover_source_files(repo_path, config)
        print(f"Found {len(source_files)} source files")

        # Analyze each file
        all_entities = {}
        all_relationships = []

        for file_path in source_files[:config.get('max_files', 1000)]:
            try:
                file_analysis = self._analyze_source_file(file_path, config)

                if file_analysis:
                    all_entities.update(file_analysis['entities'])
                    all_relationships.extend(file_analysis['relationships'])

            except Exception as e:
                print(f"Error analyzing {file_path}: {e}")
                continue

        # Build specialized graphs
        if config.get('build_call_graph', True):
            self.call_graph = self._build_call_graph(all_entities, all_relationships)

        if config.get('analyze_dependencies', True):
            self.dependency_graph = self._build_dependency_graph(all_entities, all_relationships)

        return {
            'entities': all_entities,
            'relationships': all_relationships,
            'call_graph': self.call_graph,
            'dependency_graph': self.dependency_graph,
            'analysis_stats': {
                'files_analyzed': len(source_files),
                'entities_extracted': len(all_entities),
                'relationships_extracted': len(all_relationships),
                'supported_languages': self.supported_languages
            }
        }

    def _analyze_python_file(self, file_path: str, config: Dict) -> Dict[str, Any]:
        """Analyze Python file using AST parsing."""

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                source_code = f.read()

            # Parse AST
            tree = ast.parse(source_code, filename=file_path)

            entities = {}
            relationships = []

            # Extract entities and relationships
            for node in ast.walk(tree):
                # Function definitions
                if isinstance(node, ast.FunctionDef):
                    func_entity = self._extract_function_entity(node, file_path, source_code)
                    entities[func_entity['canonical']] = func_entity

                    # Extract function relationships (calls, uses)
                    func_relationships = self._extract_function_relationships(
                        node, func_entity['canonical'], source_code
                    )
                    relationships.extend(func_relationships)

                # Class definitions
                elif isinstance(node, ast.ClassDef):
                    class_entity = self._extract_class_entity(node, file_path, source_code)
                    entities[class_entity['canonical']] = class_entity

                    # Extract class relationships (inheritance, methods)
                    class_relationships = self._extract_class_relationships(
                        node, class_entity['canonical'], source_code
                    )
                    relationships.extend(class_relationships)

                # Import statements
                elif isinstance(node, (ast.Import, ast.ImportFrom)):
                    import_relationships = self._extract_import_relationships(
                        node, file_path
                    )
                    relationships.extend(import_relationships)

            return {
                'entities': entities,
                'relationships': relationships,
                'file_metadata': {
                    'path': file_path,
                    'language': 'python',
                    'lines_of_code': len(source_code.splitlines()),
                    'ast_nodes': len(list(ast.walk(tree)))
                }
            }

        except Exception as e:
            print(f"Python AST analysis error for {file_path}: {e}")
            return None

    def _extract_function_entity(self, node: ast.FunctionDef,
                               file_path: str, source_code: str) -> Dict[str, Any]:
        """Extract function entity with comprehensive metadata."""

        # Get function signature
        signature = self._get_function_signature(node)

        # Extract docstring
        docstring = ast.get_docstring(node) or ""

        # Get source code for function
        function_source = self._get_node_source(node, source_code)

        # Analyze parameters and return type
        params = [arg.arg for arg in node.args.args]
        returns = self._extract_return_type(node)

        canonical_name = f"{file_path}::{node.name}"

        return {
            'canonical': canonical_name,
            'type': 'FUNCTION',
            'name': node.name,
            'signature': signature,
            'parameters': params,
            'returns': returns,
            'docstring': docstring,
            'source_code': function_source,
            'file_path': file_path,
            'line_start': node.lineno,
            'line_end': getattr(node, 'end_lineno', node.lineno),
            'complexity': self._calculate_complexity(node),
            'calls': self._extract_function_calls(node),
            'confidence': 0.95  # High confidence for AST extraction
        }

    def _build_call_graph(self, entities: Dict[str, Any],
                         relationships: List[Dict]) -> nx.DiGraph:
        """Build call graph from extracted entities and relationships."""

        call_graph = nx.DiGraph()

        # Add function nodes
        for entity_id, entity in entities.items():
            if entity['type'] == 'FUNCTION':
                call_graph.add_node(entity_id, **{
                    'name': entity['name'],
                    'file_path': entity['file_path'],
                    'complexity': entity.get('complexity', 1),
                    'parameters': entity.get('parameters', []),
                    'docstring': entity.get('docstring', '')[:200]  # Truncate for storage
                })

        # Add call edges
        for relationship in relationships:
            if relationship['predicate'] == 'calls':
                caller = relationship['subject']
                callee = relationship['object']

                if caller in call_graph.nodes and callee in call_graph.nodes:
                    call_graph.add_edge(caller, callee, **{
                        'confidence': relationship.get('confidence', 0.8),
                        'call_count': relationship.get('call_count', 1),
                        'evidence': relationship.get('evidence', '')
                    })

        # Calculate graph metrics
        self._calculate_call_graph_metrics(call_graph)

        return call_graph

    def _calculate_call_graph_metrics(self, call_graph: nx.DiGraph):
        """Calculate and store call graph metrics."""

        # Basic graph metrics
        num_nodes = call_graph.number_of_nodes()
        num_edges = call_graph.number_of_edges()

        if num_nodes > 0:
            # Centrality measures
            in_degree_centrality = nx.in_degree_centrality(call_graph)
            out_degree_centrality = nx.out_degree_centrality(call_graph)
            betweenness_centrality = nx.betweenness_centrality(call_graph)

            # Add centrality as node attributes
            for node in call_graph.nodes():
                call_graph.nodes[node].update({
                    'in_degree_centrality': in_degree_centrality.get(node, 0),
                    'out_degree_centrality': out_degree_centrality.get(node, 0),
                    'betweenness_centrality': betweenness_centrality.get(node, 0)
                })

            # Identify key functions (high centrality)
            key_functions = sorted(
                call_graph.nodes(),
                key=lambda x: (call_graph.nodes[x]['in_degree_centrality'] +
                              call_graph.nodes[x]['betweenness_centrality']),
                reverse=True
            )[:10]

            # Store graph-level metadata
            call_graph.graph.update({
                'num_functions': num_nodes,
                'num_calls': num_edges,
                'key_functions': key_functions,
                'analysis_timestamp': time.time()
            })

# Advanced graph traversal for GraphRAG

class GraphTraversalEngine:
    """Advanced graph traversal engine for multi-hop reasoning.

    This engine solves the fundamental challenge of graph exploration: how to find
    meaningful paths through a knowledge graph without being overwhelmed by the
    exponential growth of possible paths.

    Key innovations:
    - Adaptive Strategy Selection: Chooses optimal traversal based on query type
    - Semantic Guidance: Uses embedding similarity to prune irrelevant paths
    - Multi-Criteria Ranking: Evaluates paths on multiple quality dimensions
    - Early Termination: Stops exploration when sufficient quality paths found

    Performance characteristics:
    - 3-hop traversals: <200ms on graphs with 100K entities
    - Semantic filtering reduces path space by 80-95%
    - Quality-based ranking improves answer relevance by 40-60%
    """

    def __init__(self, neo4j_manager: Neo4jGraphManager, embedding_model):
        self.neo4j = neo4j_manager
        self.embedding_model = embedding_model

        # Traversal strategies - each optimized for different exploration patterns
        self.traversal_strategies = {
            'breadth_first': self._breadth_first_traversal,        # Nearby relationships
            'depth_first': self._depth_first_traversal,           # Deep chains
            'semantic_guided': self._semantic_guided_traversal,   # Query-relevant paths
            'relevance_ranked': self._relevance_ranked_traversal, # High-quality relationships
            'community_focused': self._community_focused_traversal # Dense clusters
        }

        # Path ranking functions - multi-criteria evaluation
        self.path_rankers = {
            'shortest_path': self._rank_by_path_length,              # Minimize hops
            'semantic_coherence': self._rank_by_semantic_coherence,   # Query relevance
            'entity_importance': self._rank_by_entity_importance,     # Entity significance
            'relationship_confidence': self._rank_by_relationship_confidence  # Extraction quality
        }

    def multi_hop_retrieval(self, query: str, starting_entities: List[str],
                           traversal_config: Dict = None) -> Dict[str, Any]:
        """Perform multi-hop retrieval using graph traversal.

        This is the core method that enables GraphRAG's multi-hop reasoning:

        1. Path Discovery: Find all relevant paths from seed entities
        2. Intelligent Filtering: Apply semantic and confidence-based pruning
        3. Path Ranking: Score paths by multiple quality criteria
        4. Context Extraction: Convert graph paths into natural language
        5. Synthesis: Combine path information into comprehensive answers

        The method balances exploration breadth with computational efficiency,
        ensuring comprehensive coverage while maintaining real-time response.
        """

        config = traversal_config or {
            'max_hops': 3,                          # Reasonable depth limit
            'max_paths': 50,                        # Top-k path selection
            'strategy': 'semantic_guided',          # Query-relevant traversal
            'path_ranking': 'semantic_coherence',   # Primary ranking criterion
            'include_path_context': True,           # Rich context extraction
            'semantic_threshold': 0.7               # Quality gate
        }

        print(f"Multi-hop retrieval for query: {query[:100]}...")
        print(f"Starting from entities: {starting_entities}")
        print(f"Configuration - Max hops: {config['max_hops']}, Strategy: {config['strategy']}")

        # Step 1: Find relevant paths from starting entities
        # Each starting entity serves as a seed for exploration
        all_paths = []
        path_contexts = []

        for start_entity in starting_entities:
            print(f"Exploring paths from: {start_entity}")
            entity_paths = self._find_entity_paths(start_entity, query, config)
            all_paths.extend(entity_paths)
            print(f"Found {len(entity_paths)} paths from {start_entity}")

        print(f"Total paths discovered: {len(all_paths)}")

        # Step 2: Rank and filter paths using configured ranking strategy
        # Multi-criteria ranking ensures high-quality path selection
        ranked_paths = self._rank_paths(all_paths, query, config)
        print(f"Path ranking complete - using {config['path_ranking']} strategy")

        # Step 3: Extract context from top paths
        # Convert graph structures into natural language narratives
        top_paths = ranked_paths[:config.get('max_paths', 50)]
        path_contexts = self._extract_path_contexts(top_paths, query)
        print(f"Context extracted from {len(path_contexts)} top paths")

        # Step 4: Generate comprehensive answer using path information
        # Synthesize individual path contexts into coherent response
        comprehensive_context = self._synthesize_path_contexts(path_contexts, query)

        return {
            'query': query,
            'starting_entities': starting_entities,
            'discovered_paths': len(all_paths),
            'top_paths': top_paths,
            'path_contexts': path_contexts,
            'comprehensive_context': comprehensive_context,
            'traversal_metadata': {
                'max_hops': config['max_hops'],
                'strategy_used': config['strategy'],
                'paths_explored': len(all_paths),
                'semantic_threshold': config['semantic_threshold'],
                'avg_path_length': sum(len(p.get('entity_path', [])) for p in top_paths) / len(top_paths) if top_paths else 0
            }
        }