nvtitan/graphRAG

1"""2Graph Builder - constructs knowledge graph from canonical triples3Handles entity canonicalization, node/edge creation, and graph pruning4"""5from typing import List, Dict, Any, Set, Tuple6from loguru import logger7from models import CanonicalTriple, GraphNode, GraphEdge, SupportingChunk, NodeType8from graph_store import GraphStore9from embedding_service import EmbeddingService10from config import settings11import numpy as np12from collections import defaultdict13 14 15class GraphBuilder:16    """17    Builds and refines knowledge graph from canonical triples18    Implements entity canonicalization, deduplication, and pruning19    """20 21    def __init__(self, graph_store: GraphStore, embedding_service: EmbeddingService):22        self.graph_store = graph_store23        self.embedding_service = embedding_service24        self.entity_embeddings: Dict[str, np.ndarray] = {}25 26    async def build_graph(self, triples: List[CanonicalTriple]) -> Tuple[int, int]:27        """28        Build graph from canonical triples29 30        Args:31            triples: List of canonical triples32 33        Returns:34            Tuple of (num_nodes_added, num_edges_added)35        """36        logger.info(f"Building graph from {len(triples)} triples")37 38        # Step 1: Entity canonicalization - merge similar entities39        entity_map = await self._canonicalize_entities(triples)40 41        # Step 2: Create nodes42        nodes_created = 043        logger.info(f"Creating nodes from {len(entity_map)} canonical entities")44 45        for entity_label in entity_map.keys():46            node = await self._create_node(entity_label, entity_map, triples)47            if self.graph_store.add_node(node):48                nodes_created += 149                logger.debug(f"Created node: {node.label} (type: {node.type.value})")50 51        logger.info(f"✓ Successfully created {nodes_created} nodes")52 53        # Step 3: Create edges54        edges_created = 055        for triple in triples:56            # Map to canonical entities57            canonical_subject = entity_map.get(triple.subject_label, triple.subject_label)58            canonical_object = entity_map.get(triple.object_label, triple.object_label)59 60            # Skip self-loops61            if canonical_subject == canonical_object:62                continue63 64            # Get node IDs65            subject_node = self.graph_store.get_node_by_label(canonical_subject)66            object_node = self.graph_store.get_node_by_label(canonical_object)67 68            if not subject_node or not object_node:69                continue70 71            # Create edge72            edge = self._create_edge(subject_node, object_node, triple)73            if self.graph_store.add_edge(edge):74                edges_created += 175 76        logger.info(f"Created {nodes_created} nodes and {edges_created} edges")77 78        # Step 4: Compute importance scores79        self._compute_importance_scores()80 81        # Step 5: Prune low-importance nodes and edges82        pruned_nodes, pruned_edges = self._prune_graph()83 84        logger.info(f"Pruned {pruned_nodes} nodes and {pruned_edges} edges")85        logger.info(f"Final graph: {nodes_created - pruned_nodes} nodes, {edges_created - pruned_edges} edges")86 87        return nodes_created - pruned_nodes, edges_created - pruned_edges88 89    async def _canonicalize_entities(self, triples: List[CanonicalTriple]) -> Dict[str, str]:90        """91        ⚡ OPTIMIZATION: Skip expensive canonicalization (identity mapping)92 93        With 2 nodes per page hard cap and strict technical filtering,94        we have very few duplicates and highly specific entities.95        Embedding computation + O(n²) similarity checks not worth the cost.96 97        Args:98            triples: List of triples99 100        Returns:101            Dict mapping entity_label -> canonical_label (identity map)102        """103        # Collect all unique entities104        entities = set()105        for triple in triples:106            entities.add(triple.subject_label)107            entities.add(triple.object_label)108 109        # DETERMINISTIC: Sort entities for consistent ordering across runs110        entities_list = sorted(list(entities))111        logger.info(f"⚡ FAST MODE: Skipping entity canonicalization for {len(entities_list)} unique entities")112        logger.info(f"Each entity maps to itself (no merging)")113 114        # Return identity mapping - each entity maps to itself115        entity_map = {entity: entity for entity in entities_list}116 117        logger.info(f"✓ Identity mapping created (0 merges, {len(entities_list)} canonical entities)")118 119        return entity_map120 121    def _entity_to_text(self, entity: str) -> str:122        """Convert entity label to text for embedding"""123        # Simple approach: use the label as-is124        return entity125 126    async def _create_node(127        self,128        label: str,129        entity_map: Dict[str, str],130        triples: List[CanonicalTriple]131    ) -> GraphNode:132        """133        Create a graph node for an entity134 135        Args:136            label: Canonical entity label137            entity_map: Entity canonicalization map138            triples: All triples (to find supporting chunks)139 140        Returns:141            GraphNode142        """143        # Find all triples mentioning this entity144        supporting_chunks = []145        aliases = []146 147        for original_label, canonical_label in entity_map.items():148            if canonical_label == label:149                if original_label != label:150                    aliases.append(original_label)151 152        # Collect supporting chunks from triples153        chunk_scores = defaultdict(float)154        for triple in triples:155            canonical_subject = entity_map.get(triple.subject_label, triple.subject_label)156            canonical_object = entity_map.get(triple.object_label, triple.object_label)157 158            if canonical_subject == label or canonical_object == label:159                # This triple supports the node160                chunk_key = (triple.page_number, triple.justification[:100])  # Use justification as proxy161                chunk_scores[chunk_key] += triple.confidence162 163        # Convert to SupportingChunk objects164        for (page_number, snippet), score in chunk_scores.items():165            supporting_chunks.append(SupportingChunk(166                chunk_id=f"page_{page_number}",  # Placeholder167                score=score,168                page_number=page_number,169                snippet=snippet170            ))171 172        # DETERMINISTIC: Sort by score (desc) then page_number (asc) for stable ordering173        supporting_chunks.sort(key=lambda x: (-x.score, x.page_number))174        supporting_chunks = supporting_chunks[:10]175 176        # Infer node type (simple heuristic)177        node_type = self._infer_node_type(label)178 179        node = GraphNode(180            label=label,181            type=node_type,182            aliases=aliases,183            supporting_chunks=supporting_chunks,184            importance_score=0.0  # Will be computed later185        )186 187        return node188 189    def _infer_node_type(self, label: str) -> NodeType:190        """Infer node type from label (simple heuristics)"""191        label_lower = label.lower()192 193        # Check for common patterns194        if any(word in label_lower for word in ["function", "method", "algorithm"]):195            return NodeType.FUNCTION196        elif any(word in label_lower for word in ["class", "type", "struct"]):197            return NodeType.CLASS198        elif label[0].isupper() and " " not in label:  # Capitalized single word199            return NodeType.PERSON200        elif any(word in label_lower for word in ["definition", "term", "concept"]):201            return NodeType.TERM202        else:203            return NodeType.CONCEPT204 205    def _create_edge(206        self,207        from_node: GraphNode,208        to_node: GraphNode,209        triple: CanonicalTriple210    ) -> GraphEdge:211        """Create a graph edge from a triple"""212        supporting_chunk = SupportingChunk(213            chunk_id=f"page_{triple.page_number}",214            score=triple.confidence,215            page_number=triple.page_number,216            snippet=triple.justification217        )218 219        edge = GraphEdge(220            from_node=from_node.node_id,221            to_node=to_node.node_id,222            relation=triple.relation,223            confidence=triple.confidence,224            supporting_chunks=[supporting_chunk]225        )226 227        return edge228 229    def _compute_importance_scores(self):230        """231        ⚡ OPTIMIZATION: Simplified importance scoring (skip expensive PageRank)232 233        Since we're not pruning, we only need basic scores for display purposes.234        """235        logger.info("⚡ FAST MODE: Computing simplified importance scores (no PageRank)")236 237        # Update node importance with simple metric (just degree centrality)238        for node in self.graph_store.get_all_nodes():239            # Simple importance = number of connections (fast to compute)240            num_neighbors = len(self.graph_store.get_neighbors(node.node_id))241 242            # Normalize to 0-1 range (assume max 10 connections)243            importance = min(num_neighbors / 10.0, 1.0)244 245            node.importance_score = importance246 247            # Update in store (for NetworkX)248            if not self.graph_store.use_neo4j:249                self.graph_store.nodes_dict[node.node_id] = node250 251        logger.info(f"✓ Importance scores computed (based on degree centrality only)")252 253    def _prune_graph(self) -> Tuple[int, int]:254        """255        ⚡ OPTIMIZATION: Skip pruning (we already filter at extraction)256 257        Pruning is expensive (PageRank + multiple graph traversals).258        With strict filtering at extraction (technical concepts only, 2 per page),259        we don't need additional pruning.260 261        Returns:262            Tuple of (nodes_removed, edges_removed) - always (0, 0)263        """264        logger.info(f"⚡ FAST MODE: Skipping graph pruning")265        logger.info(f"Nodes already filtered at extraction with strict technical validation")266        logger.info(f"Final graph: {len(self.graph_store.get_all_nodes())} nodes, {len(self.graph_store.get_all_edges())} edges")267 268        return 0, 0269