GraphRAG: Graph-Based Retrieval-Augmented Generation

"""
Traditional RAG limitations illustrated.
"""

# Problem 1: Flat knowledge representation
flat_chunks = [
    "Alice works at Company X.",
    "Company X was founded in 2020.",
    "Company X has 500 employees.",
    "Bob is the CEO of Company X."
]

# Query: "Who founded Company X?"
# Traditional RAG retrieves: ["Company X was founded in 2020."]
# Missing: WHO founded it (founder information is in different chunk)

# Problem 2: Loss of relationship context
# Without explicit relationships, we lose:
# - "founded_by" relationships
# - "CEO_of" relationships  
# - "works_at" relationships

"""
GraphRAG representation.
"""

# Knowledge Graph
knowledge_graph = {
    'nodes': [
        {'id': 'Alice', 'type': 'Person', 'role': 'Employee'},
        {'id': 'Bob', 'type': 'Person', 'role': 'CEO'},
        {'id': 'Company X', 'type': 'Organization'},
        {'id': '2020', 'type': 'Year'}
    ],
    'edges': [
        {'from': 'Alice', 'to': 'Company X', 'relation': 'works_at'},
        {'from': 'Bob', 'to': 'Company X', 'relation': 'CEO_of'},
        {'from': 'Company X', 'to': '2020', 'relation': 'founded_in'},
        {'from': '?', 'to': 'Company X', 'relation': 'founded_by'}
    ]
}

# Now query "Who founded Company X?"
# GraphRAG traverses: founded_in -> Company X <- founded_by
# Can answer: "It was founded in 2020 (founder info needed)"

import re
from typing import List, Dict, Tuple

class KnowledgeGraphExtractor:
    """
    Extract entities and relationships from text to build a knowledge graph.
    """
    
    def __init__(self, llm=None, entity_types=None):
        self.llm = llm
        self.entity_types = entity_types or ['Person', 'Organization', 'Location', 'Date', 'Event']
        
    def extract_from_document(self, document: str) -> Dict:
        """
        Extract knowledge graph from a document.
        
        Returns:
            Dictionary with 'entities' and 'relations'
        """
        if self.llm:
            return self._llm_extract(document)
        else:
            return self._rule_based_extract(document)
    
    def _llm_extract(self, document: str) -> Dict:
        """
        Use LLM for extraction (more accurate).
        """
        prompt = f"""Extract entities and relationships from the following text.

Text: {document}

Extract:
1. Entities (with types): Person, Organization, Location, Date, Event
2. Relationships between entities

Return as JSON:
{{
    "entities": [{{"name": "...", "type": "...", "description": "..."}}],
    "relations": [{{"from": "...", "to": "...", "type": "..."}}]
}}
"""
        # Use LLM to extract (simplified)
        result = self.llm.generate(prompt)
        return self._parse_json_result(result)
    
    def _rule_based_extract(self, document: str) -> Dict:
        """
        Rule-based extraction (no LLM required).
        """
        entities = []
        relations = []
        
        # Simple NER patterns
        patterns = {
            'Person': r'\b([A-Z][a-z]+ [A-Z][a-z]+)\b',
            'Organization': r'\b([A-Z][a-zA-Z]+ (Inc|Corp|LLC|Company))\b',
            'Date': r'\b(\d{4})\b'
        }
        
        for entity_type, pattern in patterns.items():
            matches = re.findall(pattern, document)
            for match in matches:
                entities.append({
                    'name': match,
                    'type': entity_type
                })
        
        # Simple relation extraction
        relation_patterns = [
            (r'(\w+) works at (\w+)', 'works_at'),
            (r'(\w+) is the CEO of (\w+)', 'CEO_of'),
            (r'(\w+) founded (\w+)', 'founded_by'),
            (r'(\w+) is located in (\w+)', 'located_in'),
        ]
        
        for pattern, rel_type in relation_patterns:
            matches = re.findall(pattern, document)
            for match in matches:
                relations.append({
                    'from': match[0],
                    'to': match[1],
                    'type': rel_type
                })
        
        return {'entities': entities, 'relations': relations}
    
    def _parse_json_result(self, result: str) -> Dict:
        """Parse LLM JSON output."""
        import json
        # Simplified - real implementation would handle errors
        try:
            return json.loads(result)
        except:
            return {'entities': [], 'relations': []}


class GraphBuilder:
    """
    Build and maintain the knowledge graph.
    """
    
    def __init__(self):
        self.entities = {}  # {entity_id: entity_data}
        self.relations = []  # List of (from, to, relation_type)
        
    def add_entity(self, entity: Dict):
        """Add or update an entity."""
        entity_id = entity['name'].lower().replace(' ', '_')
        if entity_id not in self.entities:
            self.entities[entity_id] = entity
            
    def add_relation(self, relation: Dict):
        """Add a relationship."""
        from_id = relation['from'].lower().replace(' ', '_')
        to_id = relation['to'].lower().replace(' ', '_')
        
        # Ensure entities exist
        self.add_entity({'name': relation['from'], 'type': 'Entity'})
        self.add_entity({'name': relation['to'], 'type': 'Entity'})
        
        self.relations.append({
            'from': from_id,
            'to': to_id,
            'type': relation['type']
        })
    
    def build_from_extractions(self, extractions: List[Dict]):
        """Build graph from extraction results."""
        for extraction in extractions:
            for entity in extraction.get('entities', []):
                self.add_entity(entity)
            for relation in extraction.get('relations', []):
                self.add_relation(relation)
    
    def get_entity(self, entity_name: str) -> Dict:
        """Get entity by name."""
        entity_id = entity_name.lower().replace(' ', '_')
        return self.entities.get(entity_id)
    
    def get_neighbors(self, entity_name: str, relation_type: str = None) -> List[Dict]:
        """Get neighboring entities."""
        entity_id = entity_name.lower().replace(' ', '_')
        neighbors = []
        
        for rel in self.relations:
            if rel['from'] == entity_id:
                if relation_type is None or rel['type'] == relation_type:
                    neighbors.append({
                        'entity': self.entities.get(rel['to']),
                        'relation': rel['type']
                    })
            elif rel['to'] == entity_id:
                if relation_type is None or rel['type'] == relation_type:
                    neighbors.append({
                        'entity': self.entities.get(rel['from']),
                        'relation': f"reverse_{rel['type']}"
                    })
                    
        return neighbors
    
    def traverse(self, start: str, path: List[str], max_depth: int = 3) -> List:
        """
        Traverse graph following a path.
        
        Args:
            start: Starting entity
            path: List of relation types to follow
            max_depth: Maximum traversal depth
        
        Returns:
            All paths found
        """
        results = []
        
        def dfs(current, path_remaining, current_path):
            if not path_remaining or len(current_path) > max_depth:
                results.append(current_path + [current])
                return
                
            neighbors = self.get_neighbors(current, path_remaining[0])
            for neighbor in neighbors:
                dfs(
                    neighbor['entity']['name'] if neighbor['entity'] else None,
                    path_remaining[1:],
                    current_path + [(current, neighbor['relation'], neighbor['entity']['name'] if neighbor['entity'] else None)]
                )
                
        dfs(start, path, [])
        return results

import torch
import torch.nn as nn

class GraphEncoder(nn.Module):
    """
    Encode graph structure into embeddings.
    """
    
    def __init__(self, node_dim=768, hidden_dim=256, num_relations=10):
        super().__init__()
        
        # Node embedding
        self.node_encoder = nn.Sequential(
            nn.Linear(node_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        
        # Relation embedding
        self.relation_encoder = nn.Embedding(num_relations, hidden_dim)
        
        # Graph aggregation (Graph Neural Network style)
        self.conv1 = GraphConv(hidden_dim, hidden_dim)
        self.conv2 = GraphConv(hidden_dim, hidden_dim)
        
    def forward(self, node_features, edge_index, edge_type):
        """
        Forward pass through graph.
        
        Args:
            node_features: [num_nodes, node_dim]
            edge_index: [2, num_edges]
            edge_type: [num_edges]
        
        Returns:
            node_embeddings: [num_nodes, hidden_dim]
        """
        # Initial encoding
        x = self.node_encoder(node_features)
        
        # Graph convolutions
        x = self.conv1(x, edge_index, edge_type)
        x = torch.relu(x)
        x = self.conv2(x, edge_index, edge_type)
        
        return x


class GraphConv(nn.Module):
    """
    Simple Graph Convolution layer.
    """
    
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)
        
    def forward(self, x, edge_index, edge_type=None):
        """
        Message passing on graph.
        """
        row, col = edge_index
        
        # Simple aggregation (mean)
        out = torch.zeros_like(x)
        
        # For each node, aggregate neighbor features
        for i in range(len(row)):
            out[row[i]] += x[col[i]]
            
        # Normalize
        deg = torch.bincount(row, minlength=x.size(0)).float()
        deg = deg.clamp(min=1)
        out = out / deg.unsqueeze(1)
        
        return self.linear(out)


class GraphVectorStore:
    """
    Vector store for graph elements with hybrid search.
    """
    
    def __init__(self, embedding_dim=768):
        self.embedding_dim = embedding_dim
        self.node_store = {}  # {node_id: (embedding, metadata)}
        self.relation_store = {}
        
    def add_node(self, node_id: str, embedding: torch.Tensor, metadata: Dict):
        """Add a node to the store."""
        self.node_store[node_id] = (embedding, metadata)
    
    def search(self, query_embedding: torch.Tensor, 
               top_k: int = 5,
               filter_type: str = None) -> List[Dict]:
        """
        Search nodes by embedding similarity.
        """
        results = []
        
        for node_id, (embedding, metadata) in self.node_store.items():
            if filter_type and metadata.get('type') != filter_type:
                continue
                
            # Cosine similarity
            sim = torch.nn.functional.cosine_similarity(
                query_embedding.unsqueeze(0), 
                embedding.unsqueeze(0)
            )
            
            results.append({
                'node_id': node_id,
                'score': sim.item(),
                'metadata': metadata
            })
            
        # Sort by score
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]
    
    def hybrid_search(self, query_embedding: torch.Tensor,
                     query_text: str,
                     vector_weight: float = 0.5,
                     keyword_weight: float = 0.5) -> List[Dict]:
        """
        Hybrid search combining vector and keyword matching.
        """
        # Vector search
        vector_results = self.search(query_embedding, top_k=10)
        
        # Keyword search (simple BM25 or keyword match)
        keyword_results = self._keyword_search(query_text, top_k=10)
        
        # Combine scores
        combined = {}
        
        for result in vector_results:
            combined[result['node_id']] = {
                'score': vector_weight * result['score'],
                'metadata': result['metadata']
            }
            
        for result in keyword_results:
            if result['node_id'] in combined:
                combined[result['node_id']]['score'] += keyword_weight * result['score']
            else:
                combined[result['node_id']] = {
                    'score': keyword_weight * result['score'],
                    'metadata': result['metadata']
                }
        
        # Sort and return
        results = sorted(combined.values(), 
                       key=lambda x: x['score'], 
                       reverse=True)
        return results[:10]
    
    def _keyword_search(self, query: str, top_k: int = 10) -> List[Dict]:
        """Simple keyword-based search."""
        query_terms = set(query.lower().split())
        results = []
        
        for node_id, (embedding, metadata) in self.node_store.items():
            text = metadata.get('text', '').lower()
            text_terms = set(text.split())
            
            # Jaccard similarity
            if len(query_terms) > 0:
                overlap = len(query_terms & text_terms)
                score = overlap / len(query_terms)
            else:
                score = 0
                
            if score > 0:
                results.append({
                    'node_id': node_id,
                    'score': score,
                    'metadata': metadata
                })
                
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]

class MultiHopRetriever:
    """
    Retrieve information through multiple hops in the knowledge graph.
    """
    
    def __init__(self, graph: GraphBuilder, vector_store: GraphVectorStore):
        self.graph = graph
        self.vector_store = vector_store
        
    def retrieve(self, query: str, query_embedding: torch.Tensor,
                num_hops: int = 2) -> Dict:
        """
        Multi-hop retrieval.
        
        Args:
            query: Query string
            query_embedding: Embedded query
            num_hops: Number of reasoning hops
        
        Returns:
            Retrieved context and sources
        """
        # First hop: Find relevant entities
        initial_results = self.vector_store.search(
            query_embedding, top_k=10
        )
        
        # Collect context from multiple hops
        all_context = []
        all_sources = []
        
        for result in initial_results:
            entity = result['node_id']
            
            # Get neighbors (1 hop)
            neighbors = self.graph.get_neighbors(entity)
            all_context.extend(neighbors)
            all_sources.append(entity)
            
            if num_hops >= 2:
                # Get neighbors of neighbors (2 hops)
                for neighbor in neighbors[:3]:  # Limit for efficiency
                    if neighbor.get('entity'):
                        neighbor_name = neighbor['entity'].get('name')
                        if neighbor_name:
                            second_hops = self.graph.get_neighbors(neighbor_name)
                            all_context.extend(second_hops)
                            all_sources.append(neighbor_name)
        
        return {
            'context': all_context,
            'sources': list(set(all_sources)),
            'initial_entities': [r['node_id'] for r in initial_results]
        }
    
    def retrieve_by_path(self, query: str, 
                        query_embedding: torch.Tensor,
                        path_patterns: List[List[str]]) -> List[Dict]:
        """
        Retrieve by specific path patterns.
        
        Example paths:
        - ["CEO_of", "works_at"] -> Find CEO of company where person works
        - ["founded_by", "located_in"] -> Find where founder is located
        """
        # Find starting entities
        start_entities = self.vector_store.search(
            query_embedding, top_k=5
        )
        
        results = []
        
        for entity_result in start_entities:
            start = entity_result['node_id']
            
            for pattern in path_patterns:
                # Traverse the specified path
                paths = self.graph.traverse(start, pattern)
                
                for path in paths:
                    results.append({
                        'start': start,
                        'path': path,
                        'endpoint': path[-1] if path else None
                    })
                    
        return results


class GraphRAGPipeline:
    """
    Complete GraphRAG pipeline.
    """
    
    def __init__(self, config: Dict):
        self.config = config
        
        # Components
        self.extractor = KnowledgeGraphExtractor(llm=config.get('llm'))
        self.graph = GraphBuilder()
        self.vector_store = GraphVectorStore(
            embedding_dim=config.get('embedding_dim', 768)
        )
        self.retriever = MultiHopRetriever(self.graph, self.vector_store)
        self.llm = config.get('llm')
        
    def index_documents(self, documents: List[str]):
        """Index documents into the knowledge graph."""
        for i, doc in enumerate(documents):
            # Extract entities and relations
            extraction = self.extractor.extract_from_document(doc)
            
            # Build graph
            self.graph.build_from_extractions([extraction])
            
            # Create embeddings and add to vector store
            embedding = self._create_embedding(doc)
            self.vector_store.add_node(
                f"doc_{i}",
                embedding,
                {'text': doc, 'type': 'document'}
            )
            
            # Add entity embeddings
            for entity in extraction.get('entities', []):
                entity_embedding = self._create_embedding(
                    entity.get('description', entity['name'])
                )
                self.vector_store.add_node(
                    entity['name'].lower().replace(' ', '_'),
                    entity_embedding,
                    entity
                )
                
    def query(self, query: str) -> Dict:
        """Answer a query using GraphRAG."""
        # Embed query
        query_embedding = self._create_embedding(query)
        
        # Retrieve from graph
        retrieval_result = self.retriever.retrieve(
            query, query_embedding, num_hops=2
        )
        
        # Build context for generation
        context = self._build_context(retrieval_result)
        
        # Generate answer
        if self.llm:
            answer = self.llm.generate(
                self._create_prompt(query, context)
            )
        else:
            answer = self._rule_based_generate(query, context)
            
        return {
            'answer': answer,
            'sources': retrieval_result['sources'],
            'reasoning': self._explain_reasoning(retrieval_result)
        }
    
    def _create_embedding(self, text: str) -> torch.Tensor:
        """Create embedding for text (simplified)."""
        # In practice, use sentence transformers or similar
        # Placeholder random embedding
        return torch.randn(self.config.get('embedding_dim', 768))
    
    def _build_context(self, retrieval_result: Dict) -> str:
        """Build context string from retrieval results."""
        context_parts = []
        
        for item in retrieval_result['context']:
            if item.get('entity'):
                entity = item['entity']
                if entity:
                    text = f"{entity.get('name', 'Unknown')} - {item.get('relation', '')}"
                    context_parts.append(text)
                    
        return "\n".join(context_parts[:10])  # Limit context length
    
    def _create_prompt(self, query: str, context: str) -> str:
        """Create prompt for LLM."""
        return f"""Based on the following knowledge graph context, answer the question.

Context:
{context}

Question: {query}

Answer:"""
    
    def _explain_reasoning(self, retrieval_result: Dict) -> str:
        """Explain how the answer was derived."""
        initial = retrieval_result['initial_entities']
        sources = retrieval_result['sources']
        
        return f"""Reasoning:
1. Found initial relevant entities: {initial}
2. Expanded to neighbors through graph traversal
3. Retrieved {len(sources)} source entities
4. Generated answer from combined context"""

class GraphSummarizer:
    """
    Summarize knowledge graph communities for better retrieval.
    """
    
    def __init__(self, llm=None):
        self.llm = llm
        
    def community_detection(self, graph: GraphBuilder) -> List[List[str]]:
        """
        Detect communities in the graph.
        
        Returns:
            List of communities (each community is a list of entity IDs)
        """
        # Simple community detection using connected components
        # In practice, use Louvain, Label Propagation, etc.
        
        # Build adjacency for each entity
        adj = {}
        for rel in graph.relations:
            if rel['from'] not in adj:
                adj[rel['from']] = set()
            if rel['to'] not in adj:
                adj[rel['to']] = set()
            adj[rel['from']].add(rel['to'])
            adj[rel['to']].add(rel['from'])
        
        # Find connected components
        visited = set()
        communities = []
        
        def dfs(node, component):
            visited.add(node)
            component.append(node)
            for neighbor in adj.get(node, []):
                if neighbor not in visited:
                    dfs(neighbor, component)
        
        for node in adj:
            if node not in visited:
                component = []
                dfs(node, component)
                if len(component) > 1:  # Only non-trivial communities
                    communities.append(component)
                    
        return communities
    
    def summarize_community(self, graph: GraphBuilder, 
                          community: List[str]) -> str:
        """
        Summarize a community of entities.
        """
        # Collect all info about community
        info = []
        for entity_id in community:
            entity = graph.entities.get(entity_id, {})
            info.append(f"Entity: {entity.get('name', entity_id)}")
            
            # Get relations
            neighbors = graph.get_neighbors(entity.get('name', entity_id))
            for n in neighbors[:5]:
                if n.get('entity'):
                    info.append(f"  - {n['relation']}: {n['entity'].get('name', '')}")
        
        context = "\n".join(info)
        
        if self.llm:
            prompt = f"""Summarize this knowledge graph community:

{context}

Provide a brief summary (2-3 sentences):"""
            return self.llm.generate(prompt)
        
        return context[:500]  # Simple truncation fallback

class DynamicGraphUpdater:
    """
    Handle dynamic updates to the knowledge graph.
    """
    
    def __init__(self, graph: GraphBuilder):
        self.graph = graph
        self.version = 0
        
    def add_document(self, document: str, extractor: KnowledgeGraphExtractor):
        """Add new document to graph."""
        extraction = extractor.extract_from_document(document)
        
        # Add new entities and relations
        self.graph.build_from_extractions([extraction])
        
        # Update version
        self.version += 1
        
    def update_entity(self, entity_name: str, new_data: Dict):
        """Update entity information."""
        entity_id = entity_name.lower().replace(' ', '_')
        
        if entity_id in self.graph.entities:
            # Update existing
            self.graph.entities[entity_id].update(new_data)
        else:
            # Add new
            self.graph.entities[entity_id] = new_data
            
        self.version += 1
        
    def remove_entity(self, entity_name: str):
        """Remove entity and its relations."""
        entity_id = entity_name.lower().replace(' ', '_')
        
        # Remove entity
        if entity_id in self.graph.entities:
            del self.graph.entities[entity_id]
            
        # Remove related edges
        self.graph.relations = [
            r for r in self.graph.relations
            if r['from'] != entity_id and r['to'] != entity_id
        ]
        
        self.version += 1
        
    def get_changes_since(self, version: int) -> Dict:
        """Get graph changes since a version."""
        # Simplified - track incremental changes
        return {
            'from_version': version,
            'to_version': self.version,
            'changes': f"Updated to version {self.version}"
        }

class GraphAugmentedGenerator:
    """
    Generate responses with graph-augmented context.
    """
    
    def __init__(self, llm):
        self.llm = llm
        
    def generate(self, query: str, 
                retrieval_context: Dict,
                use_graph_reasoning: bool = True) -> str:
        """
        Generate with graph-augmented context.
        """
        # Build prompt with graph context
        context_parts = []
        
        # Text context
        if retrieval_context.get('text_context'):
            context_parts.append("Text Context:")
            context_parts.append(retrieval_context['text_context'])
            
        # Graph context
        if use_graph_reasoning and retrieval_context.get('graph_context'):
            context_parts.append("\nGraph Relationships:")
            
            for rel in retrieval_context['graph_context'][:10]:
                if rel.get('entity'):
                    entity = rel['entity']
                    context_parts.append(
                        f"- {rel.get('relation', 'related')}: "
                        f"{entity.get('name', '')}"
                    )
        
        context = "\n".join(context_parts)
        
        # Create prompt
        prompt = f"""You are a helpful AI assistant. Use the provided context to answer the question accurately.

{context}

Question: {query}

Instructions:
1. Use the context to provide a factual answer
2. If the context doesn't contain enough information, say so
3. Cite specific relationships from the graph when relevant

Answer:"""
        
        return self.llm.generate(prompt)

"""
Microsoft GraphRAG Pipeline (conceptual implementation).
"""

class MicrosoftGraphRAG:
    """
    Implementation following Microsoft's GraphRAG approach.
    
    Key innovations:
    1. LLM-based entity extraction
    2. Community summarization
    3. Local and global search
    """
    
    def __init__(self, config: Dict):
        self.llm = config['llm']
        self.embedding_model = config['embedding_model']
        
        # Storage
        self.entity_graph = nx.Graph()
        self.entity_embeddings = {}
        self.community_summaries = {}
        
    def index_documents(self, documents: List[str]):
        """Index documents using GraphRAG pipeline."""
        
        # Step 1: Extract entities and relationships
        entities = []
        relationships = []
        
        for doc in documents:
            extraction = self._llm_extract(doc)
            entities.extend(extraction['entities'])
            relationships.extend(extraction['relations'])
            
        # Step 2: Build graph
        for entity in entities:
            self.entity_graph.add_node(
                entity['name'],
                **entity
            )
            
        for rel in relationships:
            self.entity_graph.add_edge(
                rel['from'], 
                rel['to'],
                relation=rel['type']
            )
            
        # Step 3: Detect communities
        communities = self._detect_communities()
        
        # Step 4: Generate community summaries
        for i, community in enumerate(communities):
            subgraph = self.entity_graph.subgraph(community)
            summary = self._summarize_community(subgraph)
            self.community_summaries[i] = summary
            
    def local_search(self, query: str, top_k: int = 10) -> str:
        """
        Local search: Retrieve specific entities and relationships.
        """
        # Embed query
        query_emb = self.embedding_model.embed(query)
        
        # Find relevant entities
        relevant_entities = self._find_relevant_entities(
            query_emb, top_k
        )
        
        # Collect local context
        context = []
        for entity_name in relevant_entities:
            # Get entity info
            entity_data = self.entity_graph.nodes[entity_name]
            context.append(f"Entity: {entity_data}")
            
            # Get relationships
            neighbors = list(self.entity_graph.neighbors(entity_name))
            for neighbor in neighbors[:5]:
                edge_data = self.entity_graph.edges[entity_name, neighbor]
                context.append(
                    f"  - {edge_data.get('relation', 'related')}: {neighbor}"
                )
                
        return "\n".join(context)
    
    def global_search(self, query: str) -> str:
        """
        Global search: Use community summaries.
        """
        # Find relevant communities
        community_context = []
        
        for comm_id, summary in self.community_summaries.items():
            # Check relevance
            if self._is_relevant(query, summary):
                community_context.append(summary)
                
        return "\n\n".join(community_context[:5])
    
    def _llm_extract(self, text: str) -> Dict:
        """Extract entities using LLM."""
        # Implementation uses LLM with prompting
        prompt = f"""Extract entities and relationships from:

{text}

Return JSON with 'entities' (name, type, description) and 
'relationships' (from, to, type)."""
        
        return self.llm.extract_json(prompt)
    
    def _detect_communities(self) -> List[List[str]]:
        """Detect communities using Louvain."""
        import networkx as nx
        
        # Use Louvain community detection
        try:
            import community
            partition = community.best_partition(self.entity_graph)
            
            # Group by community
            communities = {}
            for node, comm_id in partition.items():
                if comm_id not in communities:
                    communities[comm_id] = []
                communities[comm_id].append(node)
                
            return list(communities.values())
        except:
            # Fallback
            return [list(self.entity_graph.nodes())]
    
    def _summarize_community(self, subgraph) -> str:
        """Summarize a community using LLM."""
        # Collect subgraph info
        nodes = list(subgraph.nodes())
        edges = list(subgraph.edges(data=True))
        
        info = f"Entities: {', '.join(nodes)}\n\n"
        info += "Relationships:\n"
        for e in edges:
            info += f"- {e[0]} {e[2].get('relation', '')} {e[1]}\n"
            
        prompt = f"""Summarize this knowledge graph community:

{info}

Provide a concise summary:"""
        
        return self.llm.generate(prompt)

class EnterpriseGraphRAG:
    """
    GraphRAG for enterprise knowledge bases.
    """
    
    def __init__(self, config):
        self.pipeline = MicrosoftGraphRAG(config)
        
    def index_company_documents(self, documents: List[str]):
        """Index company documents."""
        # Extract and index
        self.pipeline.index_documents(documents)
        
    def query_knowledge_base(self, question: str) -> Dict:
        """Query the knowledge base."""
        # Try local search first
        local_result = self.pipeline.local_search(question)
        
        # If not enough, try global
        if len(local_result) < 200:
            global_result = self.pipeline.global_search(question)
            context = local_result + "\n\n" + global_result
        else:
            context = local_result
            
        # Generate answer
        answer = self.pipeline.llm.generate(
            f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
        )
        
        return {
            'answer': answer,
            'sources': self.pipeline.local_search(question, top_k=5)
        }

class ResearchGraphRAG:
    """
    Analyze research papers with GraphRAG.
    """
    
    def __init__(self, config):
        self.pipeline = MicrosoftGraphRAG(config)
        
    def index_papers(self, papers: List[Dict]):
        """
        Index research papers.
        
        papers: List of {title, abstract, content}
        """
        texts = [f"{p['title']}\n{p['abstract']}\n{p.get('content', '')}" 
                for p in papers]
        self.pipeline.index_documents(texts)
        
    def find_related_work(self, paper_title: str, query: str) -> str:
        """Find related work based on citations and topics."""
        # Search for relevant papers
        result = self.pipeline.local_search(query)
        
        return f"Related to '{paper_title}' based on knowledge graph:\n\n{result}"

class SupportGraphRAG:
    """
    GraphRAG for customer support.
    """
    
    def __init__(self, config):
        self.pipeline = MicrosoftGraphRAG(config)
        
    def index_support_docs(self, docs: List[str]):
        """Index support documentation."""
        self.pipeline.index_documents(docs)
        
    def answer_support_question(self, question: str) -> Dict:
        """Answer customer question."""
        # Try both local and global
        local = self.pipeline.local_search(question)
        global_s = self.pipeline.global_search(question)
        
        # Combine for comprehensive answer
        context = f"{local}\n\n{global_s}"
        
        answer = self.pipeline.llm.generate(
            f"Customer Question: {question}\n\n"
            f"Support Context:\n{context}\n\n"
            f"Provide helpful answer:"
        )
        
        return {
            'answer': answer,
            'related_topics': self.pipeline.local_search(question, top_k=3)
        }

class EntityExtractionOptimizer:
    """
    Optimize entity extraction quality.
    """
    
    @staticmethod
    def improve_extraction(document: str, llm) -> Dict:
        """
        Use multiple strategies for better extraction.
        """
        # Strategy 1: Few-shot prompting
        prompt_with_examples = f"""
Extract entities from the text.

Examples:
Text: "Apple Inc. was founded by Steve Jobs in Cupertino."
Entities: {{"name": "Apple Inc.", "type": "Organization"}, {"name": "Steve Jobs", "type": "Person"}}
Relations: {{"from": "Steve Jobs", "to": "Apple Inc.", "type": "founded"}}

Text: {document}
Entities:"""
        
        # Strategy 2: Entity resolution
        extracted = llm.extract(prompt_with_examples)
        resolved = EntityResolver().resolve(extracted)
        
        return resolved

class GraphMaintenance:
    """
    Best practices for graph maintenance.
    """
    
    @staticmethod
    def periodic_rebuild(graph: GraphBuilder, 
                        documents: List[str],
                        threshold: float = 0.3):
        """
        Periodically rebuild graph when drift is high.
        """
        # Track changes
        old_entities = set(graph.entities.keys())
        
        # Rebuild
        new_graph = GraphBuilder()
        # ... rebuild ...
        
        # Check drift
        new_entities = set(new_graph.entities.keys())
        drift = 1 - len(old_entities & new_entities) / len(old_entities | new_entities)
        
        if drift > threshold:
            return new_graph  # Significant changes
        return graph  # Keep existing