2b_process_graph_phase1.py

"""
GraphRAG Phase 1: LLM-based Entity and Relationship Extraction
Builds initial knowledge graph from markdown files using LLMs (Cerebras or Gemini)
"""

import json
import logging
import os
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List
from datetime import datetime
import orjson
from json_repair import repair_json
import google.generativeai as genai
import openai
from my_config import MY_CONFIG

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class GraphBuilder:
    
    def __init__(self, llm_provider="cerebras"):
        self.llm_provider = llm_provider.lower()
        
        # Global entity registry for deduplication across files
        self.global_entity_registry = {}
        
        # Initialize graph data structure
        self.graph_data = {"nodes": [], "relationships": []}
        self.processed_files = 0
        
        # Initialize LLM API based on provider
        if self.llm_provider == "cerebras":
            if not MY_CONFIG.CEREBRAS_API_KEY:
                raise ValueError("CEREBRAS_API_KEY environment variable not set. Get free key at: https://cloud.cerebras.ai/")
            
            # Configure Cerebras client
            self.cerebras_client = openai.OpenAI(
                api_key=MY_CONFIG.CEREBRAS_API_KEY,
                base_url="https://api.cerebras.ai/v1"
            )
            self.model_name = "llama-4-scout-17b-16e-instruct"  
            logger.info("🚀 Using Cerebras API")
            
        elif self.llm_provider == "gemini":
            if not MY_CONFIG.GEMINI_API_KEY:
                raise ValueError("GEMINI_API_KEY environment variable not set. Get free key at: https://aistudio.google.com/")
            
            # Configure Gemini with FREE tier
            genai.configure(api_key=MY_CONFIG.GEMINI_API_KEY)
            self.model_name = "gemini-1.5-flash" 
            self.gemini_model = genai.GenerativeModel(self.model_name)
            logger.info("🆓 Using Google Gemini API,)")
            
        else:
            valid_providers = ["cerebras", "gemini"]
            raise ValueError(f"Invalid provider '{llm_provider}'. Choose from: {valid_providers}")

        # Configure extraction parameters
        self.min_entities = int(os.getenv("GRAPH_MIN_ENTITIES", "5"))
        self.max_entities = int(os.getenv("GRAPH_MAX_ENTITIES", "15"))
        self.min_relationships = int(os.getenv("GRAPH_MIN_RELATIONSHIPS", "3"))
        self.max_relationships = int(os.getenv("GRAPH_MAX_RELATIONSHIPS", "8"))
        self.min_confidence = float(os.getenv("GRAPH_MIN_CONFIDENCE", "0.8"))
        self.max_content_chars = int(os.getenv("GRAPH_MAX_CONTENT_CHARS", "12000"))
        self.sentence_boundary_ratio = float(os.getenv("GRAPH_SENTENCE_BOUNDARY_RATIO", "0.7"))
        
        logger.info(f"✅ Initialized {self.llm_provider.upper()} provider with model: {self.model_name}")
        logger.info(f"Extraction config: {self.min_entities}-{self.max_entities} entities, {self.min_relationships}-{self.max_relationships} relationships, min confidence: {self.min_confidence}")
        logger.info(f"Content processing: {self.max_content_chars} chars per chunk with overlap for FULL analysis")

    # STEP 0: Clean Graph Data Folder
    def clean_graph_folder(self, graph_dir: str = None):
        if graph_dir is None:
            graph_dir = "workspace/graph_data"
        try:
            graph_path = Path(graph_dir)
            if graph_path.exists():
                # Remove all files in the directory
                for file_path in graph_path.glob("*"):
                    if file_path.is_file():
                        file_path.unlink()
                        logger.debug(f"Removed: {file_path.name}")
                logger.info(f"Cleaned graph folder: {graph_dir}")
            else:
                # Create directory if it doesn't exist
                graph_path.mkdir(parents=True, exist_ok=True)
                logger.info(f"Created graph folder: {graph_dir}")
        except Exception as e:
            logger.warning(f"Failed to clean graph folder: {e}")

    # STEP 1: Content Preprocessing and Chunking
    def _preprocess_content(self, text: str, max_chars: int = None) -> str:
        # Remove excessive whitespace but keep full content
        text = ' '.join(text.split())
        return text.strip()
    
    def _chunk_content(self, text: str, chunk_size: int = None, overlap: int = 200) -> List[str]:
        if chunk_size is None:
            chunk_size = self.max_content_chars
            
        # If content fits in one chunk, return as-is
        if len(text) <= chunk_size:
            return [text]
        
        chunks = []
        start = 0
        
        while start < len(text):
            # Calculate end position
            end = start + chunk_size
            
            if end >= len(text):
                # Last chunk
                chunks.append(text[start:])
                break
            
            # Try to find good break point (sentence boundary)
            chunk_text = text[start:end]
            last_period = chunk_text.rfind('.')
            last_newline = chunk_text.rfind('\n')
            
            # Use best break point
            break_point = max(last_period, last_newline)
            if break_point > chunk_size * 0.7:  # Good break point
                actual_end = start + break_point + 1
                chunks.append(text[start:actual_end])
                start = actual_end - overlap  # Overlap for context
            else:
                # No good break point, use hard split
                chunks.append(text[start:end])
                start = end - overlap
                
        return chunks

    # STEP 2: LLM Prompt Generation
    def get_entity_extraction_prompt(self) -> str:
        return f"""You are a specialized knowledge graph extraction assistant. Your task is to analyze content and extract entities and relationships to build comprehensive knowledge graphs.

DYNAMIC EXTRACTION REQUIREMENTS:
- Extract {self.min_entities}-{self.max_entities} most important entities from the content
- Create {self.min_relationships}-{self.max_relationships} meaningful relationships between entities
- Confidence threshold: {self.min_confidence} (only include high-confidence extractions)
- Focus on extracting diverse entity types relevant to the content domain

CONSTITUTIONAL AI PRINCIPLES:
1. Content-Adaptive: Determine entity types based on content analysis, not predefined categories
2. Relationship-Rich: Focus on meaningful semantic relationships between entities
3. Context-Aware: Consider document context and domain when extracting entities
4. Quality-First: Prioritize extraction quality over quantity

ENTITY EXTRACTION GUIDELINES:
- Identify the most important concepts, terms, people, places, organizations, technologies, events
- Extract entities that would be valuable for knowledge graph queries
- Include both explicit entities (directly mentioned) and implicit entities (strongly implied)
- Assign appropriate types based on semantic analysis of the entity's role in the content

RELATIONSHIP EXTRACTION GUIDELINES:
- Create relationships that capture semantic meaning, not just co-occurrence
- Use descriptive relationship types that express the nature of the connection
- Include hierarchical, associative, and causal relationships where appropriate
- Ensure relationships are bidirectionally meaningful and contextually accurate

OUTPUT FORMAT (strict JSON):
{{
    "entities": [
        {{
            "text": "Entity Name",
            "type": "DynamicType",
            "content": "Comprehensive description of the entity",
            "confidence": 0.95
        }}
    ],
    "relationships": [
        {{
            "startNode": "Entity Name 1",
            "endNode": "Entity Name 2",
            "type": "DESCRIPTIVE_RELATIONSHIP_TYPE",
            "description": "Clear description of the relationship",
            "evidence": "Direct evidence from text supporting this relationship",
            "confidence": 0.90
        }}
    ]
}}

IMPORTANT: Respond with ONLY the JSON object. No explanations, no markdown formatting, no code blocks."""

    # STEP 3: LLM Inference Methods  
    def _cerebras_inference(self, system_prompt: str, user_prompt: str) -> str:
        try:
            # Cerebras uses OpenAI-compatible chat format
            response = self.cerebras_client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.1,
                max_tokens=2000
            )
            
            # Check for empty response 
            if not response or not response.choices or not response.choices[0].message.content:
                raise ValueError("Empty response from Cerebras")
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            # Check for quota/rate limit exceeded errors
            error_str = str(e).lower()
            if "429" in str(e) and "quota" in error_str:
                logger.error(f"🚫 QUOTA EXCEEDED: Cerebras API rate/quota limit reached - {e}")
                raise Exception("QUOTA_EXCEEDED") from e
            else:
                logger.error(f"Error with Cerebras inference: {e}")
                raise e
    
    def _gemini_inference(self, system_prompt: str, user_prompt: str) -> str:
        try:
            combined_prompt = f"{system_prompt}\n\n{user_prompt}"
            response = self.gemini_model.generate_content(combined_prompt)    
            if not response or not response.text:
                raise ValueError("Empty response from Gemini")
                
            return response.text.strip()
            
        except Exception as e:
            # Check for quota exceeded error
            if "429" in str(e) and "quota" in str(e).lower():
                logger.error(f"🚫 QUOTA EXCEEDED: Gemini API daily limit reached - {e}")
                raise Exception("QUOTA_EXCEEDED") from e
            else:
                logger.error(f"Error with Gemini inference: {e}")
                raise e
    
    # STEP 4: JSON Parsing Pipeline
    def _smart_json_parse(self, json_text: str) -> Dict[str, Any]:
        
        cleaned_text = json_text.strip()
        
        # Step 1: orjson
        try:
            result = orjson.loads(cleaned_text.encode('utf-8'))
            logger.debug("✅ Step 1: orjson succeeded")
            return result
        except Exception as e:
            logger.debug(f"❌ Step 1: orjson failed - {e}")
        
        # Step 2: json-repair
        try:
            repaired = repair_json(cleaned_text)
            result = orjson.loads(repaired.encode('utf-8'))
            logger.debug("✅ Step 2: json-repair + orjson succeeded")
            return result
        except Exception as e:
            logger.debug(f"❌ Step 2: json-repair failed - {e}")
        
        # Step 3: standard json
        try:
            result = json.loads(cleaned_text)
            logger.debug("✅ Step 3: standard json succeeded")
            return result
        except Exception as e:
            logger.debug(f"❌ Step 3: standard json failed - {e}")
        
        # Step 4: json-repair + standard json
        try:
            repaired = repair_json(cleaned_text)
            result = json.loads(repaired)
            logger.debug("✅ Step 4: json-repair + standard json succeeded")
            return result
        except Exception as e:
            logger.debug(f"❌ Step 4: json-repair + standard json failed - {e}")
        
        # Step 5: All failed - this will trigger save failed txt files
        raise ValueError("All 4 JSON parsing steps failed")

    # STEP 5: Response Parsing and Validation
    def _parse_llm_extraction_response(self, llm_response: str, file_name: str) -> Dict[str, Any]:
        
        # Clean up response first
        cleaned_response = llm_response.strip()
        
        # Remove markdown formatting
        if "```json" in cleaned_response:
            parts = cleaned_response.split("```json")
            if len(parts) > 1:
                json_part = parts[1].split("```")[0].strip()
                cleaned_response = json_part
        elif "```" in cleaned_response:
            parts = cleaned_response.split("```")
            if len(parts) >= 3:
                cleaned_response = parts[1].strip()
        
        # Use the 5-step JSON parsing pipeline
        try:
            extraction_data = self._smart_json_parse(cleaned_response)
            
            # Validate complete format
            if self._validate_complete_format(extraction_data):
                return extraction_data
            else:
                self._save_failed_response(cleaned_response, file_name, "Format validation failed", "Missing required fields or empty values")
                return None
        except Exception as e:
            logger.error(f"❌ All JSON parsing steps failed for file {file_name}: {str(e)}")
            self._save_failed_response(cleaned_response, file_name, "All parsing steps failed", str(e))
            return None
    
    # STEP 6: Format Validation
    def _validate_complete_format(self, extraction_data: Dict[str, Any]) -> bool:

        if not isinstance(extraction_data, dict):
            return False

        if "entities" not in extraction_data or "relationships" not in extraction_data:
            return False

        entities = extraction_data.get("entities", [])
        relationships = extraction_data.get("relationships", [])
        if not isinstance(entities, list) or len(entities) == 0:
            return False
        for entity in entities:
            if not isinstance(entity, dict):
                return False
            
            required_fields = ["text", "type", "content", "confidence"]
            for field in required_fields:
                if field not in entity:
                    return False
                value = entity[field]
                if value is None or value == "" or (isinstance(value, str) and not value.strip()):
                    return False
                
            if not isinstance(entity["confidence"], (int, float)) or entity["confidence"] <= 0:
                return False
        
        if isinstance(relationships, list):
            for rel in relationships:
                if not isinstance(rel, dict):
                    return False
                
                required_fields = ["startNode", "endNode", "type", "description", "evidence", "confidence"]
                for field in required_fields:
                    if field not in rel:
                        return False
                    value = rel[field]
                    if value is None or value == "" or (isinstance(value, str) and not value.strip()):
                        return False
                
                if not isinstance(rel["confidence"], (int, float)) or rel["confidence"] <= 0:
                    return False
        
        return True
    
    # STEP 7: Error Handling and Failed Response Logging
    def _save_failed_response(self, llm_response: str, file_name: str, _json_error: str, _repair_error: str):
        try:
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            output_dir = Path("workspace/graph_data")
            output_dir.mkdir(parents=True, exist_ok=True)
    
            with open(output_dir / "failed_responses.txt", 'a', encoding='utf-8') as f:
                f.write(f"# Failed response from file: {file_name} at {timestamp}\n")
                f.write(llm_response)
                f.write("\n---\n") 
                f.flush()  
                
        except Exception as save_error:
            logger.error(f"Failed to save failed response from {file_name}: {save_error}")
            
    # STEP 8: Main Entity Extraction
    def extract_entities_with_llm(self, content: str, file_name: str) -> Dict[str, Any]:
        # Preprocess content
        processed_content = self._preprocess_content(content)
        
        # Split into chunks 
        chunks = self._chunk_content(processed_content)
        
        logger.info(f"📄 Processing {file_name}: {len(processed_content)} chars in {len(chunks)} chunk(s)")
        
        # Collect all entities and relationships from chunks
        all_entities = []
        all_relationships = []
        
        for chunk_idx, chunk in enumerate(chunks):
            logger.info(f"🔄 Processing chunk {chunk_idx + 1}/{len(chunks)} for {file_name}")
            
            # Simple retry mechanism for empty content - just send to LLM again
            max_retries = 3
            for attempt in range(max_retries):
                # Get the optimized prompt for entity extraction based on provider
                system_prompt = self.get_entity_extraction_prompt()
                
                # Create user prompt with chunk content
                chunk_info = f" (chunk {chunk_idx + 1}/{len(chunks)})" if len(chunks) > 1 else ""
                user_prompt = f"""
                Analyze the following content from file "{file_name}"{chunk_info}:
                
                ```
                {chunk}
                ```
                
                Extract all relevant entities, concepts, and their relationships from this content.
                """
                
                # Call appropriate LLM API
                try:
                    if self.llm_provider == "gemini":
                        llm_response = self._gemini_inference(system_prompt, user_prompt)
                    elif self.llm_provider == "cerebras":
                        llm_response = self._cerebras_inference(system_prompt, user_prompt)
                    else:
                        raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
                except Exception as e:
                    if "QUOTA_EXCEEDED" in str(e):
                        logger.error(f"🚫 QUOTA EXCEEDED on file {file_name}, chunk {chunk_idx + 1} - stopping processing")
                        # Return partial results if we have any
                        return {
                            "entities": all_entities,
                            "relationships": all_relationships,
                            "file": file_name,
                            "structure": {"section": "partial_quota_exceeded"},
                            "chunks_processed": chunk_idx,
                            "total_content_length": len(processed_content),
                            "quota_exceeded": True
                        }
                    else:
                        raise e
                    
                # Parse the JSON response
                result = self._parse_llm_extraction_response(llm_response, f"{file_name}_chunk_{chunk_idx}")
                if result is not None or attempt == max_retries - 1:
                    if result is None:
                        logger.warning(f"❌ Chunk {chunk_idx + 1} of {file_name} failed all validation attempts, skipping")
                        break

                    # Chunk results to collections
                    chunk_entities = result.get("entities", [])
                    chunk_relationships = result.get("relationships", [])
                    
                    # Add chunk identifier to entities for deduplication
                    for entity in chunk_entities:
                        entity["chunk_id"] = chunk_idx
                        entity["source_chunk"] = f"chunk_{chunk_idx}"
                    
                    # Add chunk identifier to relationships
                    for rel in chunk_relationships:
                        rel["chunk_id"] = chunk_idx
                        rel["source_chunk"] = f"chunk_{chunk_idx}"
                    
                    all_entities.extend(chunk_entities)
                    all_relationships.extend(chunk_relationships)
                    
                    logger.info(f"✅ Chunk {chunk_idx + 1}: {len(chunk_entities)} entities, {len(chunk_relationships)} relationships")
                    break
                else:
                    logger.info(f"Chunk {chunk_idx + 1} attempt {attempt + 1}/{max_retries}: Validation failed, retrying")
        
        # Deduplicate entities across chunks (same entity name = same entity)
        unique_entities = {}
        for entity in all_entities:
            entity_key = entity.get("text", "").lower().strip()
            if entity_key and entity_key not in unique_entities:
                unique_entities[entity_key] = entity
            elif entity_key:
                # Merge information from duplicate entities
                existing = unique_entities[entity_key]
                existing["confidence"] = max(existing.get("confidence", 0), entity.get("confidence", 0))
                # Combine descriptions
                existing_desc = existing.get("content", "")
                new_desc = entity.get("content", "")
                if new_desc and new_desc not in existing_desc:
                    existing["content"] = f"{existing_desc}; {new_desc}".strip("; ")
        
        # Deduplicate relationships (same startNode+endNode+type = same relationship)
        unique_relationships = {}
        for rel in all_relationships:
            rel_key = f"{rel.get('startNode', '').lower()}||{rel.get('endNode', '').lower()}||{rel.get('type', '').lower()}"
            if rel_key and rel_key not in unique_relationships:
                unique_relationships[rel_key] = rel
            elif rel_key:
                # Keep highest confidence relationship
                existing = unique_relationships[rel_key]
                if rel.get("confidence", 0) > existing.get("confidence", 0):
                    unique_relationships[rel_key] = rel
        
        final_entities = list(unique_entities.values())
        final_relationships = list(unique_relationships.values())
        
        logger.info(f"Final results for {file_name}: {len(final_entities)} unique entities, {len(final_relationships)} unique relationships")
        
        return {
            "entities": final_entities,
            "relationships": final_relationships,
            "file": file_name,
            "structure": {"section": "full_analysis"},
            "chunks_processed": len(chunks),
            "total_content_length": len(processed_content)
        }
    

    
    # STEP 9: Single File Processing
    def process_md_file(self, md_file_path: str) -> Dict[str, Any]:
        logger.info(f"Processing: {md_file_path}")
        
        try:
            # Read file content
            with open(md_file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            file_name = os.path.basename(md_file_path)
            
            # Extract entities and relationships using LLM-only approach
            llm_data = self.extract_entities_with_llm(content, file_name)
            
            # Use LLM data - create nodes and relationships from validated data
            entities_added = 0
            relationships_added = 0
            
            # Check if quota was exceeded during extraction
            quota_exceeded = llm_data.get("quota_exceeded", False)
            if quota_exceeded:
                return {
                    "file": file_name,
                    "status": "quota_exceeded",
                    "entities_extracted": len(llm_data.get("entities", [])),
                    "unique_entities_added": 0,
                    "relationships_generated": 0,
                    "processed_at": datetime.now().isoformat(),
                    "error": "API quota exceeded during processing"
                }
            
            # Process entities from LLM
            for entity in llm_data.get("entities", []):
                entity_text = entity["text"]
                semantic_key = entity_text.lower().strip()
                
                # Add to global registry if new
                if semantic_key not in self.global_entity_registry:
                    # Use LLM data directly
                    entity["id"] = str(uuid.uuid4())
                    entity["source_file"] = file_name
                    
                    self.global_entity_registry[semantic_key] = entity
                    self.graph_data["nodes"].append(entity)
                    entities_added += 1
            
            # Process relationships from LLM 
            for rel in llm_data.get("relationships", []):
                # Apply confidence threshold filtering
                rel_confidence = rel.get("confidence", 0.0)
                if rel_confidence < self.min_confidence:
                    continue  # Skip low-confidence relationships
                
                start_text = rel["startNode"].lower().strip()
                end_text = rel["endNode"].lower().strip()
                
                # Only create if both entities exist
                if start_text in self.global_entity_registry and end_text in self.global_entity_registry:
                    # Use original relationship type without sanitization
                    original_type = rel["type"]
                    
                    # Create clean relationship with only Neo4j fields
                    clean_rel = {
                        "id": str(uuid.uuid4()),
                        "startNode": self.global_entity_registry[start_text]["id"],
                        "endNode": self.global_entity_registry[end_text]["id"],
                        "type": original_type,  # Use original type preserving semantic meaning
                        "description": rel.get("description", ""),
                        "evidence": rel.get("evidence", ""),
                        "confidence": rel_confidence,
                        "chunk_id": rel.get("chunk_id", 0),
                        "source_chunk": rel.get("source_chunk", ""),
                        "source_file": file_name
                    }
                    
                    self.graph_data["relationships"].append(clean_rel)
                    relationships_added += 1
            
            result = {
                "file": file_name,
                "status": "success",
                "entities_extracted": len(llm_data.get("entities", [])),
                "unique_entities_added": entities_added,
                "relationships_generated": relationships_added,
                "processed_at": datetime.now().isoformat()
            }
            
            self.processed_files += 1
            logger.info(f"✅ Processed {file_name}: {entities_added} new entities, {relationships_added} relationships")
            return result
            
        except Exception as e:
            logger.error(f"❌ Error processing {md_file_path}: {e}")
            return {
                "file": os.path.basename(md_file_path),
                "status": "error",
                "error": str(e),
                "processed_at": datetime.now().isoformat()
            }
    
    # STEP 10: Batch File Processing
    def process_all_md_files(self, input_dir: str = None, output_path: str = None) -> Dict[str, Any]:
        if input_dir is None:
            input_dir = "workspace/processed"
        if output_path is None:
            output_path = os.path.join("workspace/graph_data", "graph-data-initial.json")
        
        # Clean the graph folder before starting fresh processing
        graph_dir = os.path.dirname(output_path)
        self.clean_graph_folder(graph_dir)
        
        input_path = Path(input_dir)
        md_files = list(input_path.glob("**/*.md"))  # Include subdirectories
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        if not md_files:
            logger.warning(f"No markdown files found in {input_dir}")
            return {"status": "no_files", "message": "No markdown files found"}
        
        logger.info(f"Found {len(md_files)} markdown files to process")
        
        # Reset data structures for a clean batch processing
        self.graph_data = {"nodes": [], "relationships": []}
        self.global_entity_registry = {}  # Reset global registry
        self.processed_files = 0
        
        logger.info(f"🚀 Starting document processing with Neo4j format output ({self.llm_provider.upper()})...")
        
        # Process files with progress tracking
        results = []
        processed_successfully = []
        failed_files = []
        quota_exceeded_files = []
        start_time = time.time()
        
        for i, md_file in enumerate(md_files, 1):
            file_start_time = time.time()
            logger.info(f"Processing file {i}/{len(md_files)}: {md_file.name}")
            
            # Track registry size before processing
            initial_registry_size = len(self.global_entity_registry)
            initial_relationship_count = len(self.graph_data["relationships"])
            
            # Process the file
            result = self.process_md_file(str(md_file))
            results.append(result)
            
            # Track file status for detailed logging
            file_status = result.get("status", "unknown")
            if file_status == "success":
                processed_successfully.append(md_file.name)
            elif file_status == "quota_exceeded":
                quota_exceeded_files.append(md_file.name)
                logger.warning(f"🚫 QUOTA EXCEEDED - Stopping batch processing at file {i}/{len(md_files)}")
                break  # Stop processing when quota exceeded
            else:
                failed_files.append((md_file.name, result.get("error", "Unknown error")))
            
            # Calculate processing metrics
            file_time = time.time() - file_start_time
            new_entities = len(self.global_entity_registry) - initial_registry_size
            new_relationships = len(self.graph_data["relationships"]) - initial_relationship_count
            
            # Show detailed progress information
            logger.info(f" File processed in {file_time:.2f}s: {new_entities} new entities, {new_relationships} relationships")
            
            # Show batch progress at regular intervals
            if i % 5 == 0 or i == len(md_files):
                successful_so_far = sum(1 for r in results if r.get("status") == "success")
                elapsed = time.time() - start_time
                avg_time = elapsed / i
                remaining = avg_time * (len(md_files) - i)
                
                logger.info(f"Progress: {i}/{len(md_files)} files ({successful_so_far} successful)")
                logger.info(f" Current stats: {len(self.global_entity_registry)} unique entities, {len(self.graph_data['relationships'])} relationships")
                logger.info(f"Time elapsed: {elapsed:.1f}s (avg {avg_time:.1f}s per file, ~{remaining:.1f}s remaining)")
        
        # Generate comprehensive summary with detailed tracking
        elapsed = time.time() - start_time
        successful = len(processed_successfully)
        quota_exceeded = len(quota_exceeded_files)
        failed = len(failed_files)
        unique_entities = len(self.global_entity_registry)
        
        # Save detailed processing lists
        self._save_processing_logs(processed_successfully, quota_exceeded_files, failed_files, output_path)
        
        # Count entity types
        entity_types = {}
        for entity_info in self.global_entity_registry.values():
            entity_type = entity_info["type"]
            entity_types[entity_type] = entity_types.get(entity_type, 0) + 1
        
        # Count relationship types
        relationship_types = {}
        for rel in self.graph_data["relationships"]:
            rel_type = rel["type"]
            relationship_types[rel_type] = relationship_types.get(rel_type, 0) + 1
        
        summary = {
            "status": "completed",
            "total_files": len(md_files),
            "successful": successful,
            "quota_exceeded": quota_exceeded,
            "failed": failed,
            "unique_entities": unique_entities,
            "total_relationships": len(self.graph_data["relationships"]),
            "entity_types": entity_types,
            "relationship_types": relationship_types,
            "processing_time_seconds": elapsed,
            "average_time_per_file": elapsed / len(md_files) if md_files else 0,
            "model": self.model_name,
            "llm_provider": self.llm_provider,
            "processed_at": datetime.now().isoformat()
        }
        
        logger.info(f"✅ Processing complete in {elapsed:.1f}s: {successful}/{len(md_files)} files successful")
        if quota_exceeded > 0:
            logger.warning(f"🚫 {quota_exceeded} files hit quota limit")
        if failed > 0:
            logger.error(f"❌ {failed} files failed with errors")
        logger.info(f"Final stats: {unique_entities} unique entities, {len(self.graph_data['relationships'])} relationships")
        
        # Log entity and relationship type breakdown
        logger.info("Entity types:")
        for entity_type, count in sorted(entity_types.items(), key=lambda x: x[1], reverse=True)[:10]:
            logger.info(f"  - {entity_type}: {count}")
        
        logger.info("Relationship types:")
        for rel_type, count in sorted(relationship_types.items(), key=lambda x: x[1], reverse=True)[:10]:
            logger.info(f"  - {rel_type}: {count}")
            
            return summary
    
    # STEP 10.5: Processing Logs Tracking
    def _save_processing_logs(self, successful_files: List[str], quota_exceeded_files: List[str], failed_files: List[tuple], output_path: str):
        try:
            output_dir = Path(output_path).parent
            
            # Save successfully processed files
            with open(output_dir / "processed_successfully.txt", 'w', encoding='utf-8') as f:
                f.write(f"# Successfully Processed Files ({len(successful_files)} total)\n")
                f.write(f"# Generated: {datetime.now().isoformat()}\n\n")
                for file_name in successful_files:
                    f.write(f"{file_name}\n")
            
            # Save quota exceeded files
            if quota_exceeded_files:
                with open(output_dir / "quota_exceeded_files.txt", 'w', encoding='utf-8') as f:
                    f.write(f"# Files That Hit Quota Limit ({len(quota_exceeded_files)} total)\n")
                    f.write(f"# Generated: {datetime.now().isoformat()}\n\n")
                    for file_name in quota_exceeded_files:
                        f.write(f"{file_name}\n")
            
            # Save failed files with errors
            if failed_files:
                with open(output_dir / "failed_files.txt", 'w', encoding='utf-8') as f:
                    f.write(f"# Files That Failed Processing ({len(failed_files)} total)\n")
                    f.write(f"# Generated: {datetime.now().isoformat()}\n\n")
                    for file_name, error in failed_files:
                        f.write(f"{file_name}: {error}\n")
                        
            logger.info(f"📋 Processing logs saved to {output_dir}")
            
        except Exception as e:
            logger.error(f"❌ Failed to save processing logs: {e}")
    
    # STEP 11: Graph Data Output
    def save_graph_data(self, output_path: str = None) -> bool:
        if output_path is None:
            output_path = os.path.join("workspace/graph_data", "graph-data-initial.json")
        try:
            # Ensure output directory exists
            output_dir = Path(output_path).parent
            output_dir.mkdir(parents=True, exist_ok=True)
            
            # Compile final data from global entity registry
            final_nodes = []
            
            for semantic_key, entity_info in self.global_entity_registry.items():
                entity_id = entity_info["id"]
                
                # Create Neo4j node
                node = {
                    "id": entity_id,
                    "elementId": entity_id,
                    "labels": [entity_info["type"]], 
                    "properties": {
                        "name": entity_info["text"],
                        "content": entity_info.get("content", ""),
                        "source": entity_info.get("source_file", ""),
                        "confidence": entity_info["confidence"],
                        "created_date": datetime.now().strftime("%Y-%m-%d"),
                        "extraction_method": self.llm_provider
                    }
                }
                final_nodes.append(node)
            
            # Use relationships
            final_relationships = self.graph_data["relationships"]
            
            # Prepare final graph data
            final_graph = {
                "nodes": final_nodes,
                "relationships": final_relationships,
                "metadata": {
                    "node_count": len(final_nodes),
                    "relationship_count": len(final_relationships),
                    "generated_at": datetime.now().isoformat(),
                    "generator": "Allycat GraphBuilder",
                    "llm_provider": self.llm_provider,
                    "model": self.model_name,
                    "format_version": "neo4j-2025"
                }
            }
            
            # Save final graph data
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(final_graph, f, indent=2, ensure_ascii=False)
            
            # Calculate final output size
            output_size = os.path.getsize(output_path)
            output_size_mb = output_size / (1024 * 1024)
            
            logger.info(f"✅ Neo4j graph data saved to {output_path} ({output_size_mb:.2f} MB)")
            logger.info(f"Final stats: {len(final_nodes)} nodes, {len(final_relationships)} relationships")
            return True
            
        except Exception as e:
            logger.error(f"❌ Error saving graph data: {e}")
            return False
    
# STEP 12: Main Entry Point
def main():
    """Main function to run the content analysis pipeline."""
    logger.info(" Starting Content Analysis Pipeline (Cloud-based APIs)")
    
    # Choose LLM provider from environment or default to cerebras
    llm_provider = os.getenv("GRAPH_LLM_PROVIDER", "cerebras").lower()
    logger.info(f" Using LLM provider: {llm_provider.upper()}")
    
    # Validate provider choice
    valid_providers = ["cerebras", "gemini"]
    if llm_provider not in valid_providers:
        logger.warning(f"⚠️ Invalid provider '{llm_provider}'. Using 'cerebras' (default)")
        llm_provider = "cerebras"
    
    try:
        analyzer = GraphBuilder(llm_provider=llm_provider)
        
        # Normal processing
        summary = analyzer.process_all_md_files()
        
        if summary["status"] == "no_files":
            logger.warning("⚠️ No files to process")
            return 1
        
        if analyzer.save_graph_data():
            logger.info("✅ Content Analysis completed successfully!")
            logger.info(f" Results: {summary['successful']}/{summary['total_files']} files processed")
            logger.info(f"Graph: {summary['unique_entities']} nodes, {summary['total_relationships']} relationships")
            logger.info(f"Model used: {analyzer.model_name} via {llm_provider.upper()}")
            return 0
        else:
            logger.error("❌ Failed to save graph data")
            return 1
            
    except Exception as e:
        logger.error(f"❌ Pipeline failed: {e}")
        return 1

if __name__ == "__main__":
    exit(main())