Create document_processor.py

src/document_processor.py

@@ -0,0 +1,338 @@
+import os
+import io
+import sqlite3
+import pandas as pd
+from typing import List, Dict, Any
+from pathlib import Path
+
+# Document processing libraries
+import PyPDF2
+import pdfplumber
+from docx import Document
+import pytesseract
+from PIL import Image
+
+# ML libraries
+from sentence_transformers import SentenceTransformer
+
+from config import Config
+
+class DocumentProcessor:
+    """Handle document processing for various file types"""
+    
+    def __init__(self, config: Config = None):
+        self.config = config or Config()
+        
+        # Initialize embedding model
+        print(f"Loading embedding model: {self.config.EMBEDDING_MODEL}")
+        self.embedding_model = SentenceTransformer(self.config.EMBEDDING_MODEL)
+        
+        # Configure Tesseract if available
+        self._setup_tesseract()
+    
+    def _setup_tesseract(self):
+        """Setup Tesseract OCR configuration"""
+        try:
+            if os.path.exists(self.config.TESSERACT_CMD):
+                pytesseract.pytesseract.tesseract_cmd = self.config.TESSERACT_CMD
+            print("✅ Tesseract OCR configured successfully")
+        except Exception as e:
+            print(f"⚠️ Tesseract setup warning: {e}")
+    
+    def extract_text_from_pdf(self, file_path: str) -> str:
+        """Extract text from PDF using multiple methods"""
+        text = ""
+        
+        try:
+            # Primary method: pdfplumber
+            with pdfplumber.open(file_path) as pdf:
+                for page_num, page in enumerate(pdf.pages):
+                    try:
+                        page_text = page.extract_text()
+                        if page_text and page_text.strip():
+                            text += f"\n[Page {page_num + 1}]\n{page_text}\n"
+                    except Exception as e:
+                        print(f"Warning: Could not extract text from page {page_num + 1}: {e}")
+                        
+        except Exception as e:
+            print(f"pdfplumber failed, trying PyPDF2: {e}")
+            
+            # Fallback method: PyPDF2
+            try:
+                with open(file_path, 'rb') as file:
+                    pdf_reader = PyPDF2.PdfReader(file)
+                    for page_num, page in enumerate(pdf_reader.pages):
+                        try:
+                            page_text = page.extract_text()
+                            if page_text and page_text.strip():
+                                text += f"\n[Page {page_num + 1}]\n{page_text}\n"
+                        except Exception as e:
+                            print(f"Warning: Could not extract text from page {page_num + 1}: {e}")
+            except Exception as e:
+                print(f"PyPDF2 also failed: {e}")
+                raise ValueError(f"Could not extract text from PDF: {e}")
+        
+        if not text.strip():
+            raise ValueError("No text content found in PDF")
+            
+        return text
+    
+    def extract_text_from_docx(self, file_path: str) -> str:
+        """Extract text from Word document"""
+        try:
+            doc = Document(file_path)
+            text = ""
+            
+            # Extract paragraph text
+            for para_num, paragraph in enumerate(doc.paragraphs):
+                if paragraph.text.strip():
+                    text += f"{paragraph.text}\n"
+            
+            # Extract table text if any
+            for table_num, table in enumerate(doc.tables):
+                text += f"\n[Table {table_num + 1}]\n"
+                for row in table.rows:
+                    row_text = " | ".join([cell.text.strip() for cell in row.cells])
+                    if row_text.strip():
+                        text += f"{row_text}\n"
+            
+            if not text.strip():
+                raise ValueError("No text content found in Word document")
+                
+            return text
+            
+        except Exception as e:
+            raise ValueError(f"Could not process Word document: {e}")
+    
+    def extract_text_from_image(self, image_data: bytes) -> str:
+        """Extract text from image using OCR"""
+        try:
+            # Open image
+            image = Image.open(io.BytesIO(image_data))
+            
+            # Convert to RGB if necessary
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            # Perform OCR
+            text = pytesseract.image_to_string(
+                image,
+                lang=self.config.OCR_LANGUAGE,
+                config='--psm 6'  # Uniform block of text
+            )
+            
+            if not text.strip():
+                # Try different PSM mode
+                text = pytesseract.image_to_string(
+                    image,
+                    lang=self.config.OCR_LANGUAGE,
+                    config='--psm 3'  # Fully automatic page segmentation
+                )
+            
+            return text.strip()
+            
+        except Exception as e:
+            raise ValueError(f"OCR failed: {e}")
+    
+    def extract_text_from_csv(self, file_path: str) -> str:
+        """Extract text from CSV file"""
+        try:
+            # Try different encodings
+            encodings = ['utf-8', 'latin-1', 'cp1252', 'iso-8859-1']
+            df = None
+            
+            for encoding in encodings:
+                try:
+                    df = pd.read_csv(file_path, encoding=encoding)
+                    break
+                except UnicodeDecodeError:
+                    continue
+            
+            if df is None:
+                raise ValueError("Could not read CSV with any supported encoding")
+            
+            # Convert DataFrame to text
+            text = f"CSV Data from: {Path(file_path).name}\n\n"
+            text += f"Shape: {df.shape[0]} rows, {df.shape[1]} columns\n\n"
+            
+            # Add column information
+            text += "Columns:\n"
+            for col in df.columns:
+                text += f"- {col}\n"
+            text += "\n"
+            
+            # Add sample data (first few rows)
+            text += "Sample Data:\n"
+            text += df.head(10).to_string(index=False) + "\n\n"
+            
+            # Add summary statistics for numeric columns
+            numeric_cols = df.select_dtypes(include=['number']).columns
+            if len(numeric_cols) > 0:
+                text += "Numeric Summary:\n"
+                text += df[numeric_cols].describe().to_string() + "\n\n"
+            
+            return text
+            
+        except Exception as e:
+            raise ValueError(f"Could not process CSV file: {e}")
+    
+    def extract_text_from_db(self, file_path: str) -> str:
+        """Extract text from SQLite database"""
+        try:
+            conn = sqlite3.connect(file_path)
+            text = f"SQLite Database: {Path(file_path).name}\n\n"
+            
+            # Get all table names
+            cursor = conn.cursor()
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+            tables = cursor.fetchall()
+            
+            if not tables:
+                raise ValueError("No tables found in database")
+            
+            text += f"Tables found: {len(tables)}\n\n"
+            
+            for table_name_tuple in tables:
+                table_name = table_name_tuple[0]
+                text += f"=== Table: {table_name} ===\n"
+                
+                try:
+                    # Get table schema
+                    cursor.execute(f"PRAGMA table_info({table_name})")
+                    columns = cursor.fetchall()
+                    
+                    text += "Columns:\n"
+                    for col in columns:
+                        text += f"- {col[1]} ({col[2]})\n"
+                    
+                    # Get row count
+                    cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
+                    row_count = cursor.fetchone()[0]
+                    text += f"Row count: {row_count}\n\n"
+                    
+                    # Get sample data
+                    df = pd.read_sql_query(f"SELECT * FROM {table_name} LIMIT 10", conn)
+                    text += "Sample Data:\n"
+                    text += df.to_string(index=False) + "\n\n"
+                    
+                except Exception as e:
+                    text += f"Error reading table {table_name}: {e}\n\n"
+            
+            conn.close()
+            return text
+            
+        except Exception as e:
+            raise ValueError(f"Could not process SQLite database: {e}")
+    
+    def chunk_text(self, text: str, metadata: Dict[str, Any] = None) -> List[Dict[str, Any]]:
+        """Split text into chunks with overlap and metadata"""
+        if not text.strip():
+            return []
+        
+        # Clean text
+        text = self._clean_text(text)
+        
+        chunks = []
+        words = text.split()
+        
+        if len(words) <= self.config.CHUNK_SIZE:
+            # If text is smaller than chunk size, return as single chunk
+            chunks.append({
+                'text': text,
+                'metadata': metadata or {},
+                'chunk_index': 0,
+                'word_count': len(words)
+            })
+        else:
+            # Split into overlapping chunks
+            for i in range(0, len(words), self.config.CHUNK_SIZE - self.config.CHUNK_OVERLAP):
+                chunk_words = words[i:i + self.config.CHUNK_SIZE]
+                chunk_text = " ".join(chunk_words)
+                
+                chunk_metadata = (metadata or {}).copy()
+                chunk_metadata.update({
+                    'chunk_index': len(chunks),
+                    'word_count': len(chunk_words),
+                    'start_word': i,
+                    'end_word': i + len(chunk_words)
+                })
+                
+                chunks.append({
+                    'text': chunk_text,
+                    'metadata': chunk_metadata
+                })
+                
+                # Break if we've covered all words
+                if i + self.config.CHUNK_SIZE >= len(words):
+                    break
+        
+        return chunks
+    
+    def _clean_text(self, text: str) -> str:
+        """Clean and normalize text"""
+        # Remove excessive whitespace
+        import re
+        text = re.sub(r'\s+', ' ', text)
+        
+        # Remove special characters that might cause issues
+        text = re.sub(r'[^\w\s\.,!?;:()\-\'"$%&@#]', ' ', text)
+        
+        # Remove excessive punctuation
+        text = re.sub(r'[.]{3,}', '...', text)
+        text = re.sub(r'[-]{3,}', '---', text)
+        
+        return text.strip()
+    
+    def process_document(self, file_path: str, file_type: str) -> List[str]:
+        """Process document based on file type and return text chunks"""
+        try:
+            # Extract text based on file type
+            if file_type.lower() == '.pdf':
+                text = self.extract_text_from_pdf(file_path)
+            elif file_type.lower() == '.docx':
+                text = self.extract_text_from_docx(file_path)
+            elif file_type.lower() == '.txt':
+                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                    text = f.read()
+            elif file_type.lower() in ['.jpg', '.jpeg', '.png']:
+                with open(file_path, 'rb') as f:
+                    text = self.extract_text_from_image(f.read())
+            elif file_type.lower() == '.csv':
+                text = self.extract_text_from_csv(file_path)
+            elif file_type.lower() == '.db':
+                text = self.extract_text_from_db(file_path)
+            else:
+                raise ValueError(f"Unsupported file type: {file_type}")
+            
+            if not text or not text.strip():
+                raise ValueError("No text content extracted from file")
+            
+            # Create metadata
+            metadata = {
+                'filename': Path(file_path).name,
+                'file_type': file_type,
+                'file_size': os.path.getsize(file_path)
+            }
+            
+            # Chunk the text
+            chunks_data = self.chunk_text(text, metadata)
+            
+            # Return just the text chunks for backward compatibility
+            return [chunk['text'] for chunk in chunks_data]
+            
+        except Exception as e:
+            print(f"Error processing document {file_path}: {e}")
+            raise
+    
+    def get_supported_formats(self) -> Dict[str, str]:
+        """Get supported file formats"""
+        return {
+            '.pdf': 'PDF documents',
+            '.docx': 'Microsoft Word documents',
+            '.txt': 'Plain text files',
+            '.jpg': 'JPEG images (with OCR)',
+            '.jpeg': 'JPEG images (with OCR)',
+            '.png': 'PNG images (with OCR)',
+            '.csv': 'Comma-separated values',
+            '.db': 'SQLite databases'
+        }