# src/utils/document_processor.py
from typing import List, Dict, Optional, Union
import PyPDF2
import docx
import pandas as pd
import json
from pathlib import Path
import hashlib
import mimetypes # Add this instead
from bs4 import BeautifulSoup
import csv
from datetime import datetime
import threading
from queue import Queue
import tiktoken
from langchain.text_splitter import RecursiveCharacterTextSplitter
import logging
from bs4.element import ProcessingInstruction
from config.config import Settings
from .enhanced_excel_processor import EnhancedExcelProcessor
class DocumentProcessor:
def __init__(
self,
chunk_size: Optional[int] = None,
chunk_overlap: Optional[int] = None,
max_file_size: Optional[int] = None,
supported_formats: Optional[List[str]] = None
):
"""
Initialize DocumentProcessor with configurable parameters
Args:
chunk_size (Optional[int]): Size of text chunks
chunk_overlap (Optional[int]): Overlap between chunks
max_file_size (Optional[int]): Maximum file size in bytes
supported_formats (Optional[List[str]]): List of supported file extensions
"""
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(levelname)s - %(message)s'
)
# Get settings with validation
default_settings = Settings.get_document_processor_settings()
# Use provided values or defaults from settings
self.chunk_size = chunk_size if chunk_size is not None else default_settings[
'chunk_size']
self.chunk_overlap = chunk_overlap if chunk_overlap is not None else default_settings[
'chunk_overlap']
self.max_file_size = max_file_size if max_file_size is not None else default_settings[
'max_file_size']
self.supported_formats = supported_formats if supported_formats is not None else default_settings[
'supported_formats']
# Validate settings
self._validate_settings()
# Initialize existing components
self.processing_queue = Queue()
self.processed_docs = {}
self._initialize_text_splitter()
self.excel_processor = EnhancedExcelProcessor()
# Check for required packages (keep existing functionality)
try:
import striprtf.striprtf
except ImportError:
logging.warning(
"Warning: striprtf package not found. RTF support will be limited.")
try:
from bs4 import BeautifulSoup
import lxml
except ImportError:
logging.warning(
"Warning: beautifulsoup4 or lxml package not found. XML support will be limited.")
def _validate_settings(self):
"""Validate and adjust settings if necessary"""
# Ensure chunk_size is positive and reasonable
self.chunk_size = max(100, self.chunk_size)
# Ensure chunk_overlap is less than chunk_size
self.chunk_overlap = min(self.chunk_overlap, self.chunk_size - 50)
# Ensure max_file_size is reasonable (minimum 1MB)
self.max_file_size = max(1024 * 1024, self.max_file_size)
# Ensure supported_formats contains valid extensions
if not self.supported_formats:
# Fallback to default supported formats if empty
self.supported_formats = Settings.DOCUMENT_PROCESSOR['supported_formats']
# Ensure all formats start with a dot
self.supported_formats = [
f".{fmt.lower().lstrip('.')}" if not fmt.startswith(
'.') else fmt.lower()
for fmt in self.supported_formats
]
def _initialize_text_splitter(self):
"""Initialize the text splitter with custom settings"""
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=self.chunk_size,
chunk_overlap=self.chunk_overlap,
length_function=len,
# Modify separators to better handle markdown while maintaining overlap
separators=["\n\n", "\n", " ", ""],
keep_separator=True,
add_start_index=True,
strip_whitespace=False # Keep whitespace to maintain markdown formatting
)
def _find_break_point(self, text: str, prev_chunk: str) -> int:
"""
Find suitable breaking point that maintains document structure
Args:
text (str): Text to find break point in (the overlap portion)
prev_chunk (str): The complete previous chunk for context
Returns:
int: Position of suitable break point
"""
# Get the context of how the previous chunk ends
prev_chunk_lines = prev_chunk.split('\n')
# Special handling for markdown tables
if '|' in prev_chunk:
# Check if we're in the middle of a table
table_rows = [
line for line in prev_chunk_lines if line.strip().startswith('|')]
if table_rows:
# Find where the current table starts in the text
table_start = text.find('|')
if table_start >= 0:
# Find the next row boundary
next_row = text.find('\n', table_start)
if next_row >= 0:
return next_row + 1 # Include the newline
# Define break point markers in order of preference
break_markers = [
('\n\n', True), # Paragraph breaks (keep marker)
('\n', True), # Line breaks (keep marker)
('. ', True), # Sentence endings (keep marker)
(', ', True), # Clause breaks (keep marker)
(' ', False) # Word breaks (don't keep marker)
]
# Check the structure of the previous chunk end
last_line = prev_chunk_lines[-1] if prev_chunk_lines else ""
# Look for each type of break point
for marker, keep_marker in break_markers:
if marker in text:
# Try to find a break point that maintains document structure
marker_positions = [i for i in range(
len(text)) if text[i:i+len(marker)] == marker]
for pos in reversed(marker_positions):
# Check if this break point would maintain document structure
if self._is_valid_break_point(text, pos, last_line):
return pos + (len(marker) if keep_marker else 0)
# If no suitable break point found, default to exact position
return min(len(text), self.chunk_overlap)
def _is_valid_break_point(self, text: str, position: int, last_line: str) -> bool:
"""
Check if a break point would maintain document structure
Args:
text (str): Text being checked
position (int): Potential break position
last_line (str): Last line of previous chunk
Returns:
bool: True if break point is valid
"""
# Don't break in the middle of markdown formatting
markdown_markers = ['*', '_', '`', '[', ']', '(', ')', '#']
if position > 0 and position < len(text) - 1:
if text[position-1] in markdown_markers or text[position+1] in markdown_markers:
return False
# Don't break in the middle of a table cell
if '|' in last_line:
cell_count = last_line.count('|')
text_before_break = text[:position]
if text_before_break.count('|') % cell_count != 0:
return False
# Don't break URLs or code blocks
url_patterns = ['http://', 'https://', '```', '`']
for pattern in url_patterns:
if pattern in text[:position] and pattern not in text[position:]:
return False
return True
def _validate_chunks(self, original_text: str, chunks: List[str]) -> bool:
"""Validate that chunks maintain document integrity"""
try:
# Remove overlap to check content
reconstructed = chunks[0]
for chunk in chunks[1:]:
if len(chunk) > self.chunk_overlap:
reconstructed += chunk[self.chunk_overlap:]
# Clean both texts for comparison (remove extra whitespace)
clean_original = ' '.join(original_text.split())
clean_reconstructed = ' '.join(reconstructed.split())
return clean_original == clean_reconstructed
except Exception as e:
logging.error(f"Error validating chunks: {str(e)}")
return False
def _extract_content(self, file_path: Path) -> str:
"""Extract content from different file formats"""
suffix = file_path.suffix.lower()
try:
if suffix == '.pdf':
return self._extract_pdf(file_path)
elif suffix == '.docx':
return self._extract_docx(file_path)
elif suffix == '.csv':
return self._extract_csv(file_path)
elif suffix == '.json':
return self._extract_json(file_path)
elif suffix == '.html':
return self._extract_html(file_path)
elif suffix == '.txt' or suffix == '.md':
return self._extract_text(file_path)
elif suffix == '.xml':
return self._extract_xml(file_path)
elif suffix == '.rtf':
return self._extract_rtf(file_path)
elif suffix in ['.xlsx', '.xls']:
return self._extract_excel(file_path)
else:
raise ValueError(f"Unsupported format: {suffix}")
except Exception as e:
raise Exception(
f"Error extracting content from {file_path}: {str(e)}")
def _extract_text(self, file_path: Path) -> str:
"""Extract content from text-based files"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
return f.read()
except UnicodeDecodeError:
with open(file_path, 'r', encoding='latin-1') as f:
return f.read()
def _extract_pdf(self, file_path: Path) -> str:
"""Extract text from PDF with advanced features"""
text = ""
with open(file_path, 'rb') as file:
reader = PyPDF2.PdfReader(file)
metadata = reader.metadata
for page in reader.pages:
text += page.extract_text() + "\n\n"
# Extract images if available
if '/XObject' in page['/Resources']:
for obj in page['/Resources']['/XObject'].get_object():
if page['/Resources']['/XObject'][obj]['/Subtype'] == '/Image':
pass
return text.strip()
def _extract_docx(self, file_path: Path) -> str:
"""Extract text from DOCX with formatting"""
doc = docx.Document(file_path)
full_text = []
for para in doc.paragraphs:
full_text.append(para.text)
for table in doc.tables:
for row in table.rows:
row_text = [cell.text for cell in row.cells]
full_text.append(" | ".join(row_text))
return "\n\n".join(full_text)
def _extract_csv(self, file_path: Path) -> str:
"""Convert CSV to structured text"""
df = pd.read_csv(file_path)
return df.to_string()
def _extract_json(self, file_path: Path) -> str:
"""Convert JSON to readable text"""
with open(file_path) as f:
data = json.load(f)
return json.dumps(data, indent=2)
def _extract_html(self, file_path: Path) -> str:
"""Extract text from HTML with structure preservation"""
with open(file_path) as f:
soup = BeautifulSoup(f, 'html.parser')
for script in soup(["script", "style"]):
script.decompose()
text = soup.get_text(separator='\n')
lines = [line.strip() for line in text.splitlines() if line.strip()]
return "\n\n".join(lines)
def _extract_xml(self, file_path: Path) -> str:
"""Extract text from XML with structure preservation"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
soup = BeautifulSoup(f, 'xml')
for pi in soup.find_all(text=lambda text: isinstance(text, ProcessingInstruction)):
pi.extract()
text = soup.get_text(separator='\n')
lines = [line.strip()
for line in text.splitlines() if line.strip()]
return "\n\n".join(lines)
except Exception as e:
raise Exception(f"Error processing XML file: {str(e)}")
def _extract_rtf(self, file_path: Path) -> str:
"""Extract text from RTF files"""
try:
import striprtf.striprtf as striprtf
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
rtf_text = f.read()
plain_text = striprtf.rtf_to_text(rtf_text)
lines = [line.strip()
for line in plain_text.splitlines() if line.strip()]
return "\n\n".join(lines)
except ImportError:
raise ImportError("striprtf package is required for RTF support.")
except Exception as e:
raise Exception(f"Error processing RTF file: {str(e)}")
def _extract_excel(self, file_path: Path) -> str:
"""Extract content from Excel files with enhanced processing"""
try:
# Use enhanced Excel processor
processed_content = self.excel_processor.process_excel(file_path)
# If processing fails, fall back to basic processing
if not processed_content:
logging.warning(
f"Enhanced Excel processing failed for {file_path}, falling back to basic processing")
return self._basic_excel_extract(file_path)
return processed_content
except Exception as e:
logging.error(f"Error in enhanced Excel processing: {str(e)}")
# Fall back to basic Excel processing
return self._basic_excel_extract(file_path)
def _basic_excel_extract(self, file_path: Path) -> str:
"""Basic Excel extraction as fallback"""
try:
excel_file = pd.ExcelFile(file_path)
sheets_data = []
for sheet_name in excel_file.sheet_names:
df = pd.read_excel(excel_file, sheet_name=sheet_name)
sheet_content = f"\nSheet: {sheet_name}\n"
sheet_content += "=" * (len(sheet_name) + 7) + "\n"
if df.empty:
sheet_content += "Empty Sheet\n"
else:
sheet_content += df.fillna('').to_string(
index=False,
max_rows=None,
max_cols=None,
line_width=120
) + "\n"
sheets_data.append(sheet_content)
return "\n\n".join(sheets_data)
except Exception as e:
raise Exception(f"Error in basic Excel processing: {str(e)}")
def _get_mime_type(self, file_path: Path) -> str:
"""
Get MIME type for a file based on its extension
Args:
file_path (Path): Path to the file
Returns:
str: MIME type of the file
"""
# Standard MIME mappings for supported formats
MIME_MAPPINGS = {
'.pdf': 'application/pdf',
'.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
'.doc': 'application/msword',
'.csv': 'text/csv',
'.json': 'application/json',
'.html': 'text/html',
'.txt': 'text/plain',
'.md': 'text/markdown',
'.xml': 'text/xml',
'.rtf': 'application/rtf',
'.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'.xls': 'application/vnd.ms-excel'
}
suffix = file_path.suffix.lower()
# Verify the file format is supported
if suffix not in self.supported_formats:
logging.warning(f"Unsupported file format: {suffix}")
return 'application/octet-stream'
# Return known MIME type or fall back to mimetypes module
if suffix in MIME_MAPPINGS:
return MIME_MAPPINGS[suffix]
mime_type = mimetypes.guess_type(str(file_path))[0]
return mime_type if mime_type else 'application/octet-stream'
def _generate_metadata(
self,
file_path: Path,
content: str,
additional_metadata: Optional[Dict] = None
) -> Dict:
"""Generate comprehensive metadata"""
file_stat = file_path.stat()
metadata = {
'filename': file_path.name,
'file_type': file_path.suffix,
'file_size': file_stat.st_size,
'created_at': datetime.fromtimestamp(file_stat.st_ctime),
'modified_at': datetime.fromtimestamp(file_stat.st_mtime),
'content_hash': self._calculate_hash(content),
'mime_type': self._get_mime_type(file_path),
'word_count': len(content.split()),
'character_count': len(content),
'processing_timestamp': datetime.utcnow().isoformat()
}
# Add Excel-specific metadata if applicable
if file_path.suffix.lower() in ['.xlsx', '.xls']:
try:
if hasattr(self.excel_processor, 'get_metadata'):
excel_metadata = self.excel_processor.get_metadata()
metadata.update({'excel_metadata': excel_metadata})
except Exception as e:
logging.warning(f"Could not extract Excel metadata: {str(e)}")
if additional_metadata:
metadata.update(additional_metadata)
return metadata
# def _generate_metadata(
# self,
# file_path: Path,
# content: str,
# additional_metadata: Optional[Dict] = None
# ) -> Dict:
# """Generate comprehensive metadata"""
# file_stat = file_path.stat()
# metadata = {
# 'filename': file_path.name,
# 'file_type': file_path.suffix,
# 'file_size': file_stat.st_size,
# 'created_at': datetime.fromtimestamp(file_stat.st_ctime),
# 'modified_at': datetime.fromtimestamp(file_stat.st_mtime),
# 'content_hash': self._calculate_hash(content),
# 'mime_type': magic.from_file(str(file_path), mime=True),
# 'word_count': len(content.split()),
# 'character_count': len(content),
# 'processing_timestamp': datetime.now().isoformat()
# }
# # Add Excel-specific metadata if applicable
# if file_path.suffix.lower() in ['.xlsx', '.xls']:
# try:
# if hasattr(self.excel_processor, 'get_metadata'):
# excel_metadata = self.excel_processor.get_metadata()
# metadata.update({'excel_metadata': excel_metadata})
# except Exception as e:
# logging.warning(f"Could not extract Excel metadata: {str(e)}")
# if additional_metadata:
# metadata.update(additional_metadata)
# return metadata
def _calculate_hash(self, text: str) -> str:
"""Calculate SHA-256 hash of text"""
return hashlib.sha256(text.encode()).hexdigest()
def _process_chunks(self, text: str) -> List[str]:
"""Process text into chunks with proper overlap"""
chunks = self.text_splitter.split_text(text)
# Ensure minimum chunk size and handle overlaps
processed_chunks = []
for i, chunk in enumerate(chunks):
if i > 0:
# Add overlap from previous chunk
overlap_start = max(
0, len(processed_chunks[-1]) - self.chunk_overlap)
chunk = processed_chunks[-1][overlap_start:] + chunk
if len(chunk) > self.chunk_size:
# Split oversized chunks
sub_chunks = self.text_splitter.split_text(chunk)
processed_chunks.extend(sub_chunks)
else:
processed_chunks.append(chunk)
return processed_chunks
async def process_document(self, file_path: Union[str, Path]) -> Dict:
"""Process document with chunk overlapping"""
file_path = Path(file_path)
if not self._validate_file(file_path):
raise ValueError(f"Invalid file: {file_path}")
content = self._extract_content(file_path)
chunks = self._process_chunks(content)
return {
'content': content,
'chunks': chunks,
'metadata': self._generate_metadata(file_path, content)
}
def _calculate_overlap_size(self, chunk1: str, chunk2: str) -> int:
"""Calculate the size of overlap between two chunks"""
min_len = min(len(chunk1), len(chunk2))
for i in range(min_len, 0, -1):
if chunk1[-i:] == chunk2[:i]:
return i
return 0
def _validate_file(self, file_path: Path) -> bool:
"""Validate file type, size, and content"""
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
if file_path.suffix.lower() not in self.supported_formats:
raise ValueError(f"Unsupported file format: {file_path.suffix}")
if file_path.stat().st_size > self.max_file_size:
raise ValueError(f"File too large: {file_path}")
if file_path.stat().st_size == 0:
raise ValueError(f"Empty file: {file_path}")
return True
def _generate_statistics(self, content: str, chunks: List[str]) -> Dict:
"""Generate document statistics"""
return {
'total_chunks': len(chunks),
'average_chunk_size': sum(len(chunk) for chunk in chunks) / len(chunks),
'token_estimate': len(content.split()),
'unique_words': len(set(content.lower().split())),
'sentences': len([s for s in content.split('.') if s.strip()]),
}
async def batch_process(
self,
file_paths: List[Union[str, Path]],
parallel: bool = True
) -> Dict[str, Dict]:
"""Process multiple documents in parallel"""
results = {}
if parallel:
threads = []
for file_path in file_paths:
thread = threading.Thread(
target=self._process_and_store,
args=(file_path, results)
)
threads.append(thread)
thread.start()
for thread in threads:
thread.join()
else:
for file_path in file_paths:
await self._process_and_store(file_path, results)
return results
async def _process_and_store(
self,
file_path: Union[str, Path],
results: Dict
):
"""Process a single document and store results"""
try:
result = await self.process_document(file_path)
results[str(file_path)] = result
except Exception as e:
results[str(file_path)] = {'error': str(e)}