CultriX/Generate-Knowledge-Graphs

1import os2import json3from typing import List, Dict, Any4import pdfplumber5from docx import Document6from config.settings import Config7 8class DocumentProcessor:9    def __init__(self):10        self.config = Config()11    12    def validate_file_size(self, file_path: str) -> bool:13        """Validate file size is within limits."""14        file_size_mb = os.path.getsize(file_path) / (1024 * 1024)15        return file_size_mb <= self.config.MAX_FILE_SIZE_MB16    17    def load_document(self, file_path: str) -> str:18        """Load document content based on file extension."""19        if not self.validate_file_size(file_path):20            raise ValueError(f"File size exceeds {self.config.MAX_FILE_SIZE_MB}MB limit")21        22        file_ext = os.path.splitext(file_path)[1].lower()23        24        if file_ext == '.pdf':25            return self._load_pdf(file_path)26        elif file_ext == '.docx':27            return self._load_docx(file_path)28        elif file_ext == '.txt':29            return self._load_txt(file_path)30        elif file_ext == '.json':31            return self._load_json(file_path)32        else:33            raise ValueError(f"Unsupported file format: {file_ext}")34    35    def _load_pdf(self, file_path: str) -> str:36        """Load PDF content."""37        text = ""38        with pdfplumber.open(file_path) as pdf:39            for page in pdf.pages:40                page_text = page.extract_text()41                if page_text:42                    text += page_text + "\n"43        return text44    45    def _load_docx(self, file_path: str) -> str:46        """Load DOCX content."""47        doc = Document(file_path)48        text = ""49        for paragraph in doc.paragraphs:50            text += paragraph.text + "\n"51        return text52    53    def _load_txt(self, file_path: str) -> str:54        """Load TXT content."""55        with open(file_path, 'r', encoding='utf-8') as file:56            return file.read()57    58    def _load_json(self, file_path: str) -> str:59        """Load JSON content and convert to text."""60        with open(file_path, 'r', encoding='utf-8') as file:61            data = json.load(file)62            return json.dumps(data, indent=2)63    64    def chunk_text(self, text: str) -> List[str]:65        """Split text into overlapping chunks for processing."""66        if len(text) <= self.config.CHUNK_SIZE:67            return [text]68        69        chunks = []70        start = 071        72        while start < len(text):73            end = start + self.config.CHUNK_SIZE74            75            # Try to break at sentence boundaries76            if end < len(text):77                # Look for sentence endings78                sentence_end = text.rfind('.', start, end)79                if sentence_end == -1:80                    sentence_end = text.rfind('!', start, end)81                if sentence_end == -1:82                    sentence_end = text.rfind('?', start, end)83                84                if sentence_end != -1 and sentence_end > start + self.config.CHUNK_SIZE // 2:85                    end = sentence_end + 186            87            chunk = text[start:end].strip()88            if chunk:89                chunks.append(chunk)90            91            start = end - self.config.CHUNK_OVERLAP92            if start >= len(text):93                break94        95        return chunks96    97    def process_documents(self, file_paths: List[str], batch_mode: bool = False) -> List[Dict[str, Any]]:98        """Process multiple documents."""99        results = []100        101        for file_path in file_paths:102            try:103                content = self.load_document(file_path)104                chunks = self.chunk_text(content)105                106                results.append({107                    'file_path': file_path,108                    'content': content,109                    'chunks': chunks,110                    'status': 'success'111                })112                113                if not batch_mode:114                    break  # Process only one file if not in batch mode115                    116            except Exception as e:117                results.append({118                    'file_path': file_path,119                    'content': '',120                    'chunks': [],121                    'status': 'error',122                    'error': str(e)123                })124        125        return results126