Optimizing GraphRAG: Five Essential Techniques for Production Performance – 2 of 4

# Traditional RAG approach
def traditional_rag_query(query, vector_db):
    """Find semantically similar documents."""
    # Convert query to embedding
    query_embedding = embed_text(query)
    
    # Find similar documents
    similar_docs = vector_db.similarity_search(query_embedding, k=5)
    
    # Return content for LLM context
    return [doc.content for doc in similar_docs]

# GraphRAG approach
def graphrag_query(query, vector_db, graph_db):
    """Find semantically similar content AND related information."""
    # Step 1: Find semantically similar content
    query_embedding = embed_text(query)
    entry_points = vector_db.similarity_search(query_embedding, k=3)
    
    # Step 2: Explore relationships from those entry points
    enriched_context = []
    for doc in entry_points:
        # Get entities from this document
        entities = doc.metadata.get('entities', [])
        
        # Traverse graph to find related information
        for entity in entities:
            related = graph_db.query("""
                MATCH (e:Entity {id: $entity_id})-[r]-(related)
                RETURN related, r
                LIMIT 10
            """, entity_id=entity)
            enriched_context.extend(related)
    
    # Step 3: Combine semantic and relational context
    return merge_contexts(entry_points, enriched_context)

import re
from typing import List, Tuple
from dataclasses import dataclass

@dataclass
class Chunk:
    content: str
    chunk_type: str  # 'prose', 'code', 'table', 'header'
    metadata: dict

class SemanticChunker:
    """
    Semantic-aware document chunker that preserves document structure
    and creates coherent chunks for optimal GraphRAG processing.
    """
    
    def __init__(self, 
                 min_chunk_size: int = 100,
                 max_chunk_size: int = 1500,
                 overlap_size: int = 50):
        self.min_chunk_size = min_chunk_size
        self.max_chunk_size = max_chunk_size
        self.overlap_size = overlap_size
        
    def chunk_document(self, text: str) -> List[Chunk]:
        """
        Process a document into semantically coherent chunks.
        """
        chunks = []
        
        # Step 1: Identify document structure
        sections = self._split_by_headers(text)
        
        for section_header, section_content in sections:
            # Step 2: Process each section based on content type
            section_chunks = self._process_section(section_header, section_content)
            chunks.extend(section_chunks)
            
        # Step 3: Add overlap for context continuity
        chunks = self._add_overlap(chunks)
        
        return chunks
    
    def _split_by_headers(self, text: str) -> List[Tuple[str, str]]:
        """Split document by markdown headers while preserving hierarchy."""
        # Pattern matches h1-h6 headers
        header_pattern = r'^(#{1,6})\s+(.+)$'
        
        sections = []
        current_header = "Document Start"
        current_content = []
        
        for line in text.split('\n'):
            header_match = re.match(header_pattern, line)
            
            if header_match:
                # Save previous section
                if current_content:
                    sections.append((current_header, '\n'.join(current_content)))
                
                # Start new section
                current_header = line
                current_content = []
            else:
                current_content.append(line)
        
        # Don't forget the last section
        if current_content:
            sections.append((current_header, '\n'.join(current_content)))
            
        return sections
    
    def _process_section(self, header: str, content: str) -> List[Chunk]:
        """Process a section based on its content type."""
        chunks = []
        
        # Extract code blocks first (they should remain intact)
        code_blocks = self._extract_code_blocks(content)
        remaining_content = content
        
        for code_block in code_blocks:
            # Replace code block with placeholder
            placeholder = f"<CODE_BLOCK_{len(chunks)}>"
            remaining_content = remaining_content.replace(code_block, placeholder, 1)
            
            # Create code chunk
            chunks.append(Chunk(
                content=code_block,
                chunk_type='code',
                metadata={'header': header}
            ))
        
        # Process remaining prose content
        prose_chunks = self._chunk_prose(remaining_content, header)
        
        # Merge chunks back in order
        final_chunks = []
        prose_idx = 0
        
        for part in remaining_content.split('<CODE_BLOCK_'):
            if prose_idx < len(prose_chunks):
                final_chunks.append(prose_chunks[prose_idx])
                prose_idx += 1
            
            # Check if this part references a code block
            if part.startswith(str(len(final_chunks) - 1) + '>'):
                code_idx = int(part.split('>')[0])
                if code_idx < len(chunks):
                    final_chunks.append(chunks[code_idx])
        
        return final_chunks
    
    def _chunk_prose(self, text: str, header: str) -> List[Chunk]:
        """Chunk prose content at natural boundaries."""
        chunks = []
        
        # First, try to split by paragraphs
        paragraphs = text.split('\n\n')
        current_chunk = []
        current_size = 0
        
        for paragraph in paragraphs:
            paragraph_size = len(paragraph)
            
            # If adding this paragraph exceeds max size, finalize current chunk
            if current_size + paragraph_size > self.max_chunk_size and current_chunk:
                chunks.append(Chunk(
                    content='\n\n'.join(current_chunk),
                    chunk_type='prose',
                    metadata={'header': header}
                ))
                current_chunk = []
                current_size = 0
            
            # If a single paragraph is too large, split by sentences
            if paragraph_size > self.max_chunk_size:
                sentence_chunks = self._split_by_sentences(paragraph, header)
                chunks.extend(sentence_chunks)
            else:
                current_chunk.append(paragraph)
                current_size += paragraph_size + 2  # +2 for \n\n
        
        # Don't forget the last chunk
        if current_chunk:
            chunks.append(Chunk(
                content='\n\n'.join(current_chunk),
                chunk_type='prose',
                metadata={'header': header}
            ))
            
        return chunks

# Benchmarking semantic vs. fixed-size chunking
def benchmark_chunking_approaches(documents):
    """Compare performance and quality metrics between chunking approaches."""
    
    results = {
        'fixed_size': {'chunks': 0, 'extraction_accuracy': 0, 'processing_time': 0},
        'semantic': {'chunks': 0, 'extraction_accuracy': 0, 'processing_time': 0}
    }
    
    for doc in documents:
        # Fixed-size chunking
        start_time = time.time()
        fixed_chunks = simple_chunk(doc, chunk_size=1000)
        fixed_entities = extract_entities(fixed_chunks)
        results['fixed_size']['processing_time'] += time.time() - start_time
        results['fixed_size']['chunks'] += len(fixed_chunks)
        
        # Semantic chunking
        start_time = time.time()
        semantic_chunks = semantic_chunker.chunk_document(doc)
        semantic_entities = extract_entities(semantic_chunks)
        results['semantic']['processing_time'] += time.time() - start_time
        results['semantic']['chunks'] += len(semantic_chunks)
    
    # Calculate improvements
    chunk_reduction = 1 - (results['semantic']['chunks'] / results['fixed_size']['chunks'])
    time_reduction = 1 - (results['semantic']['processing_time'] / results['fixed_size']['processing_time'])
    
    print(f"Chunk Reduction: {chunk_reduction:.1%}")
    print(f"Processing Time Reduction: {time_reduction:.1%}")
    
    return results

# The naive approach - what NOT to do
def create_entities_naive(entities, neo4j_driver):
    """WARNING: This approach will destroy your performance at scale."""
    created_count = 0
    
    with neo4j_driver.session() as session:
        for entity in entities:
            # Each iteration creates a new transaction!
            result = session.run("""
                CREATE (n:Entity {id: $id, name: $name, type: $type})
                RETURN n
            """, id=entity.id, name=entity.name, type=entity.type)
            
            created_count += 1
            
    return created_count

# What's really happening behind the scenes:
# 1. Open network connection (if pooled: ~1ms, if not: ~10-50ms)
# 2. Begin transaction (~5-10ms)
# 3. Acquire locks (~1-5ms)
# 4. Execute query (~1-2ms) <-- The actual work!
# 5. Commit transaction (~10-20ms)
# 6. Release locks (~1-2ms)
# 7. Close connection/return to pool (~1ms)
# Total: ~20-90ms for 1-2ms of actual work!

from typing import List, Dict, Any
import logging
from neo4j import GraphDatabase
from neo4j.exceptions import TransientError, SessionExpired

class OptimizedNeo4jBatchProcessor:
    """
    High-performance batch processor for Neo4j with adaptive sizing
    and comprehensive error handling.
    """
    
    def __init__(self, 
                 driver,
                 initial_node_batch_size: int = 500,
                 initial_rel_batch_size: int = 1000,
                 max_retries: int = 3):
        self.driver = driver
        self.node_batch_size = initial_node_batch_size
        self.rel_batch_size = initial_rel_batch_size
        self.max_retries = max_retries
        self.logger = logging.getLogger(__name__)
        
        # Adaptive sizing parameters
        self.batch_size_history = []
        self.performance_threshold = 0.8  # Target 80% success rate
        
    def batch_create_nodes(self, 
                          nodes: List[Dict[str, Any]], 
                          label: str = "Entity") -> int:
        """
        Create nodes in optimized batches with automatic size adjustment.
        """
        total_created = 0
        failed_nodes = []
        
        # Process nodes in batches
        for i in range(0, len(nodes), self.node_batch_size):
            batch = nodes[i:i + self.node_batch_size]
            
            for attempt in range(self.max_retries):
                try:
                    created = self._execute_node_batch(batch, label)
                    total_created += created
                    
                    # Record success for adaptive sizing
                    self._record_batch_performance(True, len(batch))
                    break
                    
                except TransientError as e:
                    self.logger.warning(f"Transient error on attempt {attempt + 1}: {e}")
                    if attempt == self.max_retries - 1:
                        failed_nodes.extend(batch)
                        self._record_batch_performance(False, len(batch))
                    else:
                        # Exponential backoff
                        time.sleep(2 ** attempt)
                
                except Exception as e:
                    self.logger.error(f"Unexpected error in batch creation: {e}")
                    failed_nodes.extend(batch)
                    self._record_batch_performance(False, len(batch))
                    break
            
            # Adjust batch size based on performance
            self._adjust_batch_size('node')
        
        # Handle failed nodes with smaller batches
        if failed_nodes:
            self.logger.info(f"Retrying {len(failed_nodes)} failed nodes with smaller batches")
            original_size = self.node_batch_size
            self.node_batch_size = max(10, self.node_batch_size // 10)
            
            retry_created = self.batch_create_nodes(failed_nodes, label)
            total_created += retry_created
            
            self.node_batch_size = original_size
        
        return total_created
    
    def _execute_node_batch(self, batch: List[Dict], label: str) -> int:
        """Execute a single batch of node creations."""
        with self.driver.session() as session:
            # Use UNWIND for efficient batch processing
            result = session.run(f"""
                UNWIND $batch AS node
                MERGE (n:{label} {{id: node.id}})
                ON CREATE SET n += node.properties
                ON MATCH SET n += node.properties
                RETURN count(n) as created
            """, batch=[{
                'id': node['id'],
                'properties': {k: v for k, v in node.items() if k != 'id'}
            } for node in batch])
            
            return result.single()['created']
    
    def batch_create_relationships(self,
                                 relationships: List[Dict],
                                 rel_type: str = "RELATES_TO") -> int:
        """
        Create relationships in batches with intelligent grouping.
        """
        # Group relationships by type for better performance
        grouped_rels = self._group_relationships_by_type(relationships)
        total_created = 0
        
        for rel_type, rels in grouped_rels.items():
            # Further batch by size
            for i in range(0, len(rels), self.rel_batch_size):
                batch = rels[i:i + self.rel_batch_size]
                
                try:
                    created = self._execute_relationship_batch(batch, rel_type)
                    total_created += created
                    self._record_batch_performance(True, len(batch))
                    
                except Exception as e:
                    self.logger.error(f"Error creating relationship batch: {e}")
                    self._record_batch_performance(False, len(batch))
                    
                    # Try smaller batches for failed relationships
                    if len(batch) > 10:
                        smaller_batches = [batch[j:j+10] for j in range(0, len(batch), 10)]
                        for small_batch in smaller_batches:
                            try:
                                created = self._execute_relationship_batch(small_batch, rel_type)
                                total_created += created
                            except Exception as e2:
                                self.logger.error(f"Failed even with small batch: {e2}")
                
                # Adjust batch size based on performance
                self._adjust_batch_size('relationship')
        
        return total_created
    
    def _execute_relationship_batch(self, batch: List[Dict], rel_type: str) -> int:
        """Execute a single batch of relationship creations."""
        with self.driver.session() as session:
            # Prepare batch data with proper structure
            batch_data = [{
                'source_id': rel['source_id'],
                'target_id': rel['target_id'],
                'properties': rel.get('properties', {})
            } for rel in batch]
            
            # Use parameterized query for safety and performance
            query = f"""
                UNWIND $batch AS rel
                MATCH (source:Entity {{id: rel.source_id}})
                MATCH (target:Entity {{id: rel.target_id}})
                MERGE (source)-[r:{rel_type}]->(target)
                ON CREATE SET r = rel.properties
                ON MATCH SET r += rel.properties
                RETURN count(r) as created
            """
            
            result = session.run(query, batch=batch_data)
            return result.single()['created']
    
    def _adjust_batch_size(self, operation_type: str):
        """
        Dynamically adjust batch size based on recent performance.
        """
        if len(self.batch_size_history) < 10:
            return  # Not enough data yet
            
        recent_success_rate = sum(1 for success, _ in self.batch_size_history[-10:] if success) / 10
        
        if operation_type == 'node':
            if recent_success_rate < self.performance_threshold:
                # Reduce batch size
                self.node_batch_size = max(50, int(self.node_batch_size * 0.8))
                self.logger.info(f"Reduced node batch size to {self.node_batch_size}")
            elif recent_success_rate > 0.95:
                # Increase batch size
                self.node_batch_size = min(2000, int(self.node_batch_size * 1.2))
                self.logger.info(f"Increased node batch size to {self.node_batch_size}")
        
        elif operation_type == 'relationship':
            if recent_success_rate < self.performance_threshold:
                self.rel_batch_size = max(100, int(self.rel_batch_size * 0.8))
                self.logger.info(f"Reduced relationship batch size to {self.rel_batch_size}")
            elif recent_success_rate > 0.95:
                self.rel_batch_size = min(5000, int(self.rel_batch_size * 1.2))
                self.logger.info(f"Increased relationship batch size to {self.rel_batch_size}")

# Thread 1 is creating: Alice -> knows -> Bob
# Thread 2 is creating: Bob -> knows -> Alice

# What happens:
# Thread 1: Lock Alice (success) → Try to lock Bob (waiting for Thread 2)
# Thread 2: Lock Bob (success) → Try to lock Alice (waiting for Thread 1)
# Result: DEADLOCK! Both threads waiting forever

import networkx as nx
from collections import defaultdict
from typing import List, Dict, Set, Tuple

class RelationshipGrouper:
    """
    Groups relationships to eliminate deadlocks using graph coloring.
    This ensures relationships in the same group never conflict.
    """
    
    def __init__(self, conflict_threshold: int = 1000):
        self.conflict_threshold = conflict_threshold
        self.logger = logging.getLogger(__name__)
        
    def group_relationships(self, 
                          relationships: List[Tuple[str, str, str, Dict]]) -> Dict[int, List]:
        """
        Group relationships to prevent deadlocks during parallel creation.
        
        Args:
            relationships: List of (source_id, target_id, rel_type, properties) tuples
            
        Returns:
            Dictionary mapping group IDs to lists of non-conflicting relationships
        """
        # Step 1: Build conflict graph
        conflict_graph = self._build_conflict_graph(relationships)
        
        # Step 2: Apply graph coloring
        coloring = self._color_graph(conflict_graph)
        
        # Step 3: Group relationships by color
        groups = self._organize_by_color(relationships, coloring)
        
        self.logger.info(f"Grouped {len(relationships)} relationships into {len(groups)} non-conflicting groups")
        
        return groups
    
    def _build_conflict_graph(self, relationships: List[Tuple]) -> nx.Graph:
        """
        Build a graph where nodes are relationships and edges connect
        relationships that share nodes (and thus could conflict).
        """
        conflict_graph = nx.Graph()
        
        # Track which relationships involve each entity
        entity_to_rels = defaultdict(set)
        
        # Add all relationships as nodes
        for idx, rel in enumerate(relationships):
            source_id, target_id, rel_type, _ = rel
            
            # Add relationship to conflict graph
            conflict_graph.add_node(idx, relationship=rel)
            
            # Track entity involvement
            entity_to_rels[source_id].add(idx)
            entity_to_rels[target_id].add(idx)
        
        # Add edges between conflicting relationships
        for entity_id, rel_indices in entity_to_rels.items():
            # All relationships involving this entity potentially conflict
            rel_list = list(rel_indices)
            for i in range(len(rel_list)):
                for j in range(i + 1, len(rel_list)):
                    conflict_graph.add_edge(rel_list[i], rel_list[j])
        
        self.logger.info(f"Conflict graph has {conflict_graph.number_of_nodes()} nodes "
                        f"and {conflict_graph.number_of_edges()} edges")
        
        return conflict_graph
    
    def _color_graph(self, graph: nx.Graph) -> Dict[int, int]:
        """
        Apply graph coloring to find non-conflicting groups.
        Uses various strategies based on graph characteristics.
        """
        # For sparse graphs, use a simple greedy algorithm
        if graph.number_of_edges() < graph.number_of_nodes() * 2:
            return nx.greedy_color(graph, strategy='largest_first')
        
        # For dense graphs, use a more sophisticated approach
        # Welsh-Powell algorithm tends to use fewer colors
        return self._welsh_powell_coloring(graph)
    
    def _welsh_powell_coloring(self, graph: nx.Graph) -> Dict[int, int]:
        """
        Implement Welsh-Powell algorithm for better coloring of dense graphs.
        """
        # Sort nodes by degree (descending)
        nodes_by_degree = sorted(graph.nodes(), 
                               key=lambda n: graph.degree(n), 
                               reverse=True)
        
        coloring = {}
        color = 0
        
        while nodes_by_degree:
            # Start new color
            current_color_nodes = []
            remaining_nodes = []
            
            for node in nodes_by_degree:
                # Check if this node conflicts with any node of current color
                conflicts = False
                for colored_node in current_color_nodes:
                    if graph.has_edge(node, colored_node):
                        conflicts = True
                        break
                
                if not conflicts:
                    coloring[node] = color
                    current_color_nodes.append(node)
                else:
                    remaining_nodes.append(node)
            
            nodes_by_degree = remaining_nodes
            color += 1
        
        return coloring
    
    def _organize_by_color(self, 
                          relationships: List[Tuple], 
                          coloring: Dict[int, int]) -> Dict[int, List]:
        """Organize relationships by their assigned color."""
        groups = defaultdict(list)
        
        for idx, color in coloring.items():
            groups[color].append(relationships[idx])
        
        return dict(groups)
    
    def optimize_for_supernodes(self, 
                               relationships: List[Tuple],
                               supernode_threshold: int = 100) -> Dict[int, List]:
        """
        Special handling for graphs with supernodes (highly connected nodes).
        """
        # Identify supernodes
        node_degrees = defaultdict(int)
        for source_id, target_id, _, _ in relationships:
            node_degrees[source_id] += 1
            node_degrees[target_id] += 1
        
        supernodes = {node for node, degree in node_degrees.items() 
                     if degree > supernode_threshold}
        
        if not supernodes:
            # No supernodes, use standard grouping
            return self.group_relationships(relationships)
        
        self.logger.info(f"Identified {len(supernodes)} supernodes")
        
        # Separate supernode relationships
        supernode_rels = []
        regular_rels = []
        
        for rel in relationships:
            source_id, target_id, _, _ = rel
            if source_id in supernodes or target_id in supernodes:
                supernode_rels.append(rel)
            else:
                regular_rels.append(rel)
        
        # Group regular relationships normally
        regular_groups = self.group_relationships(regular_rels)
        
        # Handle supernode relationships with finer granularity
        supernode_groups = self._group_supernode_relationships(supernode_rels, supernodes)
        
        # Merge groups
        all_groups = {}
        group_id = 0
        
        for group in regular_groups.values():
            all_groups[group_id] = group
            group_id += 1
            
        for group in supernode_groups.values():
            all_groups[group_id] = group
            group_id += 1
        
        return all_groups

def process_relationships_with_grouping(relationships, neo4j_driver):
    """
    Process relationships using grouping to prevent deadlocks.
    """
    # Initialize components
    grouper = RelationshipGrouper()
    batch_processor = OptimizedNeo4jBatchProcessor(neo4j_driver)
    
    # Group relationships
    groups = grouper.group_relationships(relationships)
    
    total_created = 0
    processing_times = []
    
    # Process each group
    for group_id, group_relationships in groups.items():
        start_time = time.time()
        
        # Convert to format expected by batch processor
        formatted_rels = [
            {
                'source_id': rel[0],
                'target_id': rel[1],
                'properties': rel[3]
            }
            for rel in group_relationships
        ]
        
        # Process this group (no conflicts within group)
        created = batch_processor.batch_create_relationships(
            formatted_rels, 
            rel_type=group_relationships[0][2]  # Assuming same type in group
        )
        
        total_created += created
        processing_times.append(time.time() - start_time)
        
        print(f"Group {group_id}: Created {created} relationships "
              f"in {processing_times[-1]:.2f} seconds")
    
    # Summary statistics
    avg_time = sum(processing_times) / len(processing_times)
    print(f"\nTotal relationships created: {total_created}")
    print(f"Average time per group: {avg_time:.2f} seconds")
    print(f"Total processing time: {sum(processing_times):.2f} seconds")
    
    return total_created

# The inefficient approach we want to avoid
def extract_entities_one_by_one(chunks, llm_client):
    """WARNING: This will burn through your API budget and patience."""
    all_entities = []
    
    for chunk in chunks:  # If you have 10,000 chunks...
        # Each call has:
        # - Network latency: ~50-200ms
        # - LLM processing: ~500-2000ms  
        # - API rate limiting delays
        # - Context window underutilization
        
        response = llm_client.complete(
            prompt=f"Extract entities from: {chunk.content}"
        )
        entities = parse_response(response)
        all_entities.extend(entities)
    
    return all_entities

from typing import List, Dict, Tuple, Optional
import json
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor
import backoff

@dataclass
class ExtractionResult:
    entities: Dict[str, Dict]
    relationships: List[Dict]
    metadata: Dict

class IntelligentLLMExtractor:
    """
    Optimized LLM-based extraction with adaptive batching,
    parallel processing, and quality optimization.
    """
    
    def __init__(self, 
                 llm_client,
                 initial_batch_size: int = 5,
                 max_batch_size: int = 10,
                 min_batch_size: int = 1,
                 parallel_workers: int = 3):
        self.llm_client = llm_client
        self.batch_size = initial_batch_size
        self.max_batch_size = max_batch_size
        self.min_batch_size = min_batch_size
        self.parallel_workers = parallel_workers
        self.logger = logging.getLogger(__name__)
        
        # Performance tracking
        self.batch_performance = []
        
    def extract_from_chunks(self, 
                           chunks: List[Chunk],
                           domain: str = "general") -> ExtractionResult:
        """
        Extract entities and relationships from chunks using optimized batching.
        """
        # Prepare batches
        batches = self._create_intelligent_batches(chunks)
        
        # Process batches in parallel
        all_entities = {}
        all_relationships = []
        
        with ThreadPoolExecutor(max_workers=self.parallel_workers) as executor:
            # Submit all batches for processing
            future_to_batch = {
                executor.submit(self._process_batch, batch, domain): batch
                for batch in batches
            }
            
            # Collect results as they complete
            for future in future_to_batch:
                try:
                    result = future.result(timeout=60)
                    
                    # Merge results
                    all_entities.update(result['entities'])
                    all_relationships.extend(result['relationships'])
                    
                    # Track performance
                    self._update_batch_size(success=True, 
                                          batch_size=len(future_to_batch[future]))
                    
                except Exception as e:
                    self.logger.error(f"Batch processing failed: {e}")
                    self._update_batch_size(success=False, 
                                          batch_size=len(future_to_batch[future]))
                    
                    # Reprocess failed batch with smaller size
                    failed_batch = future_to_batch[future]
                    if len(failed_batch) > 1:
                        self._process_failed_batch(failed_batch, all_entities, all_relationships, domain)
        
        # Post-process to resolve duplicates and enhance quality
        refined_entities, refined_relationships = self._post_process_extraction(
            all_entities, all_relationships
        )
        
        return ExtractionResult(
            entities=refined_entities,
            relationships=refined_relationships,
            metadata={'total_chunks': len(chunks), 'batches_processed': len(batches)}
        )
    
    def _create_intelligent_batches(self, chunks: List[Chunk]) -> List[List[Chunk]]:
        """
        Create batches that optimize for LLM context window usage
        and semantic coherence.
        """
        batches = []
        current_batch = []
        current_tokens = 0
        
        # Estimate token limits (leaving room for prompt and response)
        max_tokens_per_batch = 2000  # Adjust based on your LLM
        
        for chunk in chunks:
            # Estimate tokens (rough: 1 token ≈ 4 characters)
            chunk_tokens = len(chunk.content) // 4
            
            # Check if adding this chunk would exceed limits
            if (len(current_batch) >= self.batch_size or 
                current_tokens + chunk_tokens > max_tokens_per_batch):
                
                if current_batch:
                    batches.append(current_batch)
                current_batch = [chunk]
                current_tokens = chunk_tokens
            else:
                current_batch.append(chunk)
                current_tokens += chunk_tokens
        
        # Don't forget the last batch
        if current_batch:
            batches.append(current_batch)
        
        self.logger.info(f"Created {len(batches)} batches from {len(chunks)} chunks")
        return batches
    
    @backoff.on_exception(backoff.expo, Exception, max_tries=3)
    def _process_batch(self, batch: List[Chunk], domain: str) -> Dict:
        """
        Process a batch of chunks through the LLM with sophisticated prompting.
        """
        # Create structured prompt for batch processing
        batch_prompt = self._create_batch_prompt(batch, domain)
        
        # Call LLM with structured output format
        response = self.llm_client.complete(
            prompt=batch_prompt,
            temperature=0.1,  # Low temperature for consistent extraction
            max_tokens=2000
        )
        
        # Parse structured response
        return self._parse_batch_response(response, batch)
    
    def _create_batch_prompt(self, batch: List[Chunk], domain: str) -> str:
        """
        Create an optimized prompt for batch extraction with examples.
        """
        # Get domain-specific examples
        examples = self._get_domain_examples(domain)
        
        prompt = f"""You are an expert at extracting entities and relationships from text.
        
{examples}

Now extract entities and relationships from the following {len(batch)} text chunks.
For each chunk, identify key entities and their relationships.

"""
        
        # Add chunks with clear separation
        for i, chunk in enumerate(batch):
            prompt += f"\n--- Chunk {i+1} ---\n{chunk.content}\n"
        
        prompt += """
--- Instructions ---
Return a JSON object with this exact structure:
{
    "chunks": [
        {
            "chunk_id": 0,
            "entities": [
                {
                    "id": "unique_id",
                    "name": "Entity Name",
                    "type": "Person|Organization|Location|Concept|Other",
                    "confidence": 0.9
                }
            ],
            "relationships": [
                {
                    "source": "entity_id_1",
                    "target": "entity_id_2",
                    "type": "RELATES_TO|WORKS_FOR|LOCATED_IN|etc",
                    "confidence": 0.8
                }
            ]
        }
    ]
}

Be comprehensive but precise. Only extract clearly stated information."""
        
        return prompt
    
    def _get_domain_examples(self, domain: str) -> str:
        """
        Provide domain-specific examples to improve extraction quality.
        """
        examples = {
            "technical": """Example for technical documentation:
Text: "The React framework, developed by Facebook, uses a virtual DOM for efficient rendering."
Entities: React (Technology), Facebook (Organization), virtual DOM (Concept)
Relationships: React -[DEVELOPED_BY]-> Facebook, React -[USES]-> virtual DOM""",
            
            "academic": """Example for academic text:
Text: "Dr. Smith from MIT published groundbreaking research on quantum computing in Nature."
Entities: Dr. Smith (Person), MIT (Organization), quantum computing (Concept), Nature (Publication)
Relationships: Dr. Smith -[AFFILIATED_WITH]-> MIT, Dr. Smith -[RESEARCHES]-> quantum computing""",
            
            "business": """Example for business content:
Text: "Apple Inc. acquired Beats Electronics for $3 billion in 2014."
Entities: Apple Inc. (Organization), Beats Electronics (Organization), 2014 (Date)
Relationships: Apple Inc. -[ACQUIRED]-> Beats Electronics"""
        }
        
        return examples.get(domain, examples["technical"])
    
    def _parse_batch_response(self, response: str, batch: List[Chunk]) -> Dict:
        """
        Parse the LLM response and map back to original chunks.
        """
        try:
            # Extract JSON from response
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            json_str = response[json_start:json_end]
            
            parsed = json.loads(json_str)
            
            # Process each chunk's results
            entities = {}
            relationships = []
            
            for chunk_result in parsed.get('chunks', []):
                chunk_id = chunk_result.get('chunk_id', 0)
                
                # Process entities
                for entity in chunk_result.get('entities', []):
                    entity_id = f"{batch[chunk_id].metadata.get('doc_id', 'unknown')}_{entity['id']}"
                    entities[entity_id] = {
                        'name': entity['name'],
                        'type': entity['type'],
                        'confidence': entity.get('confidence', 0.8),
                        'source_chunk': chunk_id
                    }
                
                # Process relationships
                for rel in chunk_result.get('relationships', []):
                    relationships.append({
                        'source': f"{batch[chunk_id].metadata.get('doc_id', 'unknown')}_{rel['source']}",
                        'target': f"{batch[chunk_id].metadata.get('doc_id', 'unknown')}_{rel['target']}",
                        'type': rel['type'],
                        'confidence': rel.get('confidence', 0.7),
                        'source_chunk': chunk_id
                    })
            
            return {'entities': entities, 'relationships': relationships}
            
        except Exception as e:
            self.logger.error(f"Failed to parse LLM response: {e}")
            return {'entities': {}, 'relationships': []}
    
    def _update_batch_size(self, success: bool, batch_size: int):
        """
        Dynamically adjust batch size based on success rates.
        """
        self.batch_performance.append((success, batch_size))
        
        # Only adjust after sufficient data
        if len(self.batch_performance) < 10:
            return
        
        # Calculate recent success rate
        recent = self.batch_performance[-10:]
        success_rate = sum(1 for s, _ in recent if s) / 10
        
        if success_rate < 0.7 and self.batch_size > self.min_batch_size:
            # Reduce batch size
            self.batch_size = max(self.min_batch_size, self.batch_size - 1)
            self.logger.info(f"Reduced batch size to {self.batch_size}")
        elif success_rate > 0.95 and self.batch_size < self.max_batch_size:
            # Increase batch size
            self.batch_size = min(self.max_batch_size, self.batch_size + 1)
            self.logger.info(f"Increased batch size to {self.batch_size}")

import math
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Set, Tuple
import hashlib

class MixAndBatchLoader:
    """
    Implements the Mix and Batch technique for massively parallel
    relationship loading without deadlocks.
    """
    
    def __init__(self,
                 neo4j_driver,
                 num_partitions: int = 10,
                 parallel_workers: int = 4):
        self.driver = neo4j_driver
        self.num_partitions = num_partitions
        self.parallel_workers = parallel_workers
        self.logger = logging.getLogger(__name__)
        
    def load_relationships_parallel(self, 
                                  relationships: List[Tuple[str, str, str, Dict]]) -> int:
        """
        Load relationships using Mix and Batch technique for parallel processing.
        
        Args:
            relationships: List of (source_id, target_id, rel_type, properties)
            
        Returns:
            Number of relationships created
        """
        start_time = time.time()
        
        # Step 1: Partition nodes
        self.logger.info("Step 1: Partitioning nodes...")
        node_partitions = self._partition_nodes(relationships)
        
        # Step 2: Assign partition codes
        self.logger.info("Step 2: Creating partition codes...")
        partition_codes = self._create_partition_codes(relationships, node_partitions)
        
        # Step 3: Organize into batches
        self.logger.info("Step 3: Organizing batches...")
        batches = self._organize_batches(partition_codes, relationships)
        
        # Step 4: Process batches
        self.logger.info(f"Step 4: Processing {len(batches)} batches...")
        total_created = self._process_batches_parallel(batches, relationships)
        
        elapsed = time.time() - start_time
        rate = total_created / elapsed if elapsed > 0 else 0
        
        self.logger.info(f"Created {total_created} relationships in {elapsed:.2f} seconds "
                        f"({rate:.0f} relationships/second)")
        
        return total_created
    
    def _partition_nodes(self, relationships: List[Tuple]) -> Dict[str, int]:
        """
        Assign each node to a partition using consistent hashing.
        """
        node_partitions = {}
        
        # Extract all unique nodes
        nodes = set()
        for source, target, _, _ in relationships:
            nodes.add(source)
            nodes.add(target)
        
        # Assign partitions
        for node in nodes:
            if isinstance(node, (int, float)):
                # For numeric IDs, use modulo
                partition = int(node) % self.num_partitions
            else:
                # For string IDs, use consistent hashing
                hash_val = int(hashlib.md5(str(node).encode()).hexdigest(), 16)
                partition = hash_val % self.num_partitions
            
            node_partitions[node] = partition
        
        self.logger.info(f"Partitioned {len(nodes)} nodes into {self.num_partitions} partitions")
        
        # Log partition distribution for monitoring
        partition_counts = {}
        for partition in node_partitions.values():
            partition_counts[partition] = partition_counts.get(partition, 0) + 1
        
        self.logger.debug(f"Partition distribution: {partition_counts}")
        
        return node_partitions
    
    def _create_partition_codes(self, 
                               relationships: List[Tuple],
                               node_partitions: Dict[str, int]) -> Dict[int, str]:
        """
        Create partition codes for each relationship.
        """
        partition_codes = {}
        
        for idx, (source, target, _, _) in enumerate(relationships):
            source_partition = node_partitions[source]
            target_partition = node_partitions[target]
            
            # Create deterministic partition code
            partition_code = f"{source_partition}-{target_partition}"
            partition_codes[idx] = partition_code
        
        # Log partition code distribution
        code_counts = {}
        for code in partition_codes.values():
            code_counts[code] = code_counts.get(code, 0) + 1
        
        self.logger.debug(f"Created {len(code_counts)} unique partition codes")
        
        return partition_codes
    
    def _organize_batches(self, 
                         partition_codes: Dict[int, str],
                         relationships: List[Tuple]) -> List[List[int]]:
        """
        Organize relationships into non-conflicting batches using
        the Mix and Batch algorithm.
        """
        # Group relationships by partition code
        code_to_indices = {}
        for idx, code in partition_codes.items():
            if code not in code_to_indices:
                code_to_indices[code] = []
            code_to_indices[code].append(idx)
        
        batches = []
        
        # For bipartite graphs (source and target from different sets)
        if self._is_bipartite(relationships):
            batches = self._organize_bipartite_batches(code_to_indices)
        else:
            # For general graphs
            batches = self._organize_monopartite_batches(code_to_indices)
        
        self.logger.info(f"Organized {len(relationships)} relationships into {len(batches)} batches")
        
        return batches
    
    def _organize_bipartite_batches(self, code_to_indices: Dict[str, List[int]]) -> List[List[int]]:
        """
        Organize batches for bipartite graphs using diagonal pattern.
        """
        batches = []
        
        for offset in range(self.num_partitions):
            batch = []
            
            for i in range(self.num_partitions):
                j = (i + offset) % self.num_partitions
                code = f"{i}-{j}"
                
                if code in code_to_indices:
                    batch.extend(code_to_indices[code])
            
            if batch:
                batches.append(batch)
        
        return batches
    
    def _organize_monopartite_batches(self, code_to_indices: Dict[str, List[int]]) -> List[List[int]]:
        """
        Organize batches for monopartite graphs with more complex patterns.
        """
        batches = []
        processed_codes = set()
        
        # Process diagonal patterns
        for offset in range(self.num_partitions):
            batch = []
            
            for i in range(self.num_partitions):
                j = (i + offset) % self.num_partitions
                
                # Handle both directions for undirected relationships
                codes = [f"{i}-{j}", f"{j}-{i}"]
                if i == j:
                    codes = [f"{i}-{j}"]  # Self-loops only need one code
                
                for code in codes:
                    if code in code_to_indices and code not in processed_codes:
                        batch.extend(code_to_indices[code])
                        processed_codes.add(code)
            
            if batch:
                batches.append(batch)
        
        # Handle any remaining relationships
        remaining = []
        for code, indices in code_to_indices.items():
            if code not in processed_codes:
                remaining.extend(indices)
        
        if remaining:
            # Add remaining as final batch
            batches.append(remaining)
            self.logger.warning(f"Had {len(remaining)} relationships in overflow batch")
        
        return batches
    
    def _process_batches_parallel(self, 
                                 batches: List[List[int]], 
                                 relationships: List[Tuple]) -> int:
        """
        Process batches sequentially, but within each batch use parallel processing.
        """
        total_created = 0
        
        for batch_idx, batch in enumerate(batches):
            self.logger.info(f"Processing batch {batch_idx + 1}/{len(batches)} "
                           f"with {len(batch)} relationships")
            
            # Split batch into chunks for parallel workers
            chunk_size = max(1, len(batch) // self.parallel_workers)
            chunks = [batch[i:i + chunk_size] for i in range(0, len(batch), chunk_size)]
            
            # Process chunks in parallel
            with ThreadPoolExecutor(max_workers=self.parallel_workers) as executor:
                futures = []
                
                for chunk in chunks:
                    # Extract relationships for this chunk
                    chunk_relationships = [relationships[idx] for idx in chunk]
                    
                    future = executor.submit(
                        self._process_relationship_chunk,
                        chunk_relationships
                    )
                    futures.append(future)
                
                # Collect results
                for future in as_completed(futures):
                    try:
                        created = future.result()
                        total_created += created
                    except Exception as e:
                        self.logger.error(f"Error processing chunk: {e}")
        
        return total_created
    
    def _process_relationship_chunk(self, chunk_relationships: List[Tuple]) -> int:
        """
        Process a chunk of relationships in a single transaction.
        """
        with self.driver.session() as session:
            try:
                # Prepare batch data
                batch_data = []
                for source, target, rel_type, properties in chunk_relationships:
                    batch_data.append({
                        'source': source,
                        'target': target,
                        'type': rel_type,
                        'props': properties or {}
                    })
                
                # Execute batch creation
                result = session.run("""
                    UNWIND $batch AS rel
                    MATCH (source {id: rel.source})
                    MATCH (target {id: rel.target})
                    CALL apoc.create.relationship(source, rel.type, rel.props, target) 
                    YIELD rel as created
                    RETURN count(created) as count
                """, batch=batch_data)
                
                return result.single()['count']
                
            except Exception as e:
                self.logger.error(f"Failed to create relationships: {e}")
                return 0