MMo4/csit-ned-chatbot
0
1from typing import List, Dict, Any, Optional, Tuple2import re3import logging4from dataclasses import dataclass5from src.utils.helpers import clean_text, chunk_text, calculate_text_stats6from src.data.loader import DocumentChunk7 8logger = logging.getLogger(__name__)9 10@dataclass11class ProcessedChunk:12 """Represents a processed text chunk ready for embedding"""13 original_chunk_id: str14 chunk_index: int15 content: str16 metadata: Dict[str, Any]17 embedding_text: str18 word_count: int19 20 def to_vector_store_format(self) -> Tuple[str, Dict[str, Any]]:21 """Convert to format suitable for vector store"""22 doc_id = f"{self.original_chunk_id}_chunk_{self.chunk_index}"23 24 # Combine metadata with chunk-specific info25 vector_metadata = self.metadata.copy()26 27 # Convert lists to strings for ChromaDB compatibility28 for key, value in vector_metadata.items():29 if isinstance(value, list):30 vector_metadata[key] = ', '.join(str(item) for item in value)31 32 vector_metadata.update({33 'original_chunk_id': self.original_chunk_id,34 'chunk_index': self.chunk_index,35 'word_count': self.word_count,36 'doc_id': doc_id37 })38 39 return doc_id, vector_metadata40 41class TextProcessor:42 """Processes text content for optimal embedding and retrieval"""43 44 def __init__(self, chunk_size: int = 500, overlap: int = 50):45 self.chunk_size = chunk_size46 self.overlap = overlap47 48 def process_document_chunk(self, doc_chunk: DocumentChunk) -> List[ProcessedChunk]:49 """Process a document chunk into smaller embedding-ready chunks"""50 try:51 # Clean the content for processing52 cleaned_content = self._clean_markdown_content(doc_chunk.content)53 54 # Create embedding text (includes title and key metadata)55 embedding_text = self._create_embedding_text(doc_chunk, cleaned_content)56 57 # Check if content is small enough to be a single chunk58 if len(cleaned_content.split()) <= self.chunk_size:59 return [self._create_single_processed_chunk(doc_chunk, cleaned_content, embedding_text)]60 61 # Split into multiple chunks62 text_chunks = chunk_text(cleaned_content, self.chunk_size, self.overlap)63 64 processed_chunks = []65 for i, chunk_content in enumerate(text_chunks):66 # Create embedding text for each chunk67 chunk_embedding_text = self._create_chunk_embedding_text(68 doc_chunk, chunk_content, i69 )70 71 processed_chunk = ProcessedChunk(72 original_chunk_id=doc_chunk.chunk_id,73 chunk_index=i,74 content=chunk_content,75 metadata=doc_chunk.metadata,76 embedding_text=chunk_embedding_text,77 word_count=len(chunk_content.split())78 )79 80 processed_chunks.append(processed_chunk)81 82 logger.debug(f"Split {doc_chunk.chunk_id} into {len(processed_chunks)} chunks")83 return processed_chunks84 85 except Exception as e:86 logger.error(f"Error processing document chunk {doc_chunk.chunk_id}: {e}")87 return []88 89 def _create_single_processed_chunk(self, doc_chunk: DocumentChunk, 90 cleaned_content: str, embedding_text: str) -> ProcessedChunk:91 """Create a single processed chunk for small documents"""92 return ProcessedChunk(93 original_chunk_id=doc_chunk.chunk_id,94 chunk_index=0,95 content=cleaned_content,96 metadata=doc_chunk.metadata,97 embedding_text=embedding_text,98 word_count=len(cleaned_content.split())99 )100 101 def _clean_markdown_content(self, content: str) -> str:102 """Clean markdown content while preserving structure"""103 # Remove excessive whitespace but preserve paragraph breaks104 content = re.sub(r'\n{3,}', '\n\n', content)105 106 # Clean markdown syntax for better readability107 content = re.sub(r'^#{1,6}\s+', '', content, flags=re.MULTILINE) # Remove headers108 content = re.sub(r'\*\*(.*?)\*\*', r'\1', content) # Remove bold109 content = re.sub(r'\*(.*?)\*', r'\1', content) # Remove italic110 content = re.sub(r'`(.*?)`', r'\1', content) # Remove code ticks111 content = re.sub(r'\[(.*?)\]\(.*?\)', r'\1', content) # Remove links, keep text112 113 # Clean up bullet points114 content = re.sub(r'^[-*+]\s+', '• ', content, flags=re.MULTILINE)115 116 # Remove table formatting117 content = re.sub(r'\|.*?\|', '', content)118 content = re.sub(r'^[-+|:\s]+$', '', content, flags=re.MULTILINE)119 120 # Clean excessive whitespace121 content = re.sub(r'\s+', ' ', content)122 content = re.sub(r'\n\s*\n', '\n\n', content)123 124 return content.strip()125 126 def _create_embedding_text(self, doc_chunk: DocumentChunk, content: str) -> str:127 """Create optimized text for embedding generation"""128 parts = []129 130 # Add title and category info131 parts.append(f"Title: {doc_chunk.title}")132 133 category = doc_chunk.metadata.get('category', '')134 if category:135 parts.append(f"Category: {category}")136 137 # Add department information138 departments = doc_chunk.metadata.get('departments', [])139 if departments:140 parts.append(f"Programs: {', '.join(departments)}")141 142 # Add key topics143 topics = doc_chunk.metadata.get('topics', [])144 if topics:145 parts.append(f"Topics: {', '.join(topics)}")146 147 # Add main content148 parts.append(f"Content: {content}")149 150 return '\n'.join(parts)151 152 def _create_chunk_embedding_text(self, doc_chunk: DocumentChunk, 153 chunk_content: str, chunk_index: int) -> str:154 """Create embedding text for a specific chunk"""155 parts = []156 157 # Add document context158 parts.append(f"Document: {doc_chunk.title}")159 parts.append(f"Part {chunk_index + 1}")160 161 category = doc_chunk.metadata.get('category', '')162 if category:163 parts.append(f"Category: {category}")164 165 departments = doc_chunk.metadata.get('departments', [])166 if departments:167 parts.append(f"Programs: {', '.join(departments)}")168 169 # Add chunk content170 parts.append(f"Content: {chunk_content}")171 172 return '\n'.join(parts)173 174 def extract_key_phrases(self, text: str, max_phrases: int = 10) -> List[str]:175 """Extract key phrases from text for better searchability"""176 # Simple key phrase extraction (can be enhanced with NLP libraries)177 178 # Find noun phrases (simple pattern)179 noun_phrases = re.findall(r'\b[A-Z][a-z]+(?:\s+[a-z]+)*\b', text)180 181 # Find technical terms (capitalized words or abbreviations)182 technical_terms = re.findall(r'\b[A-Z]{2,}\b|\b[A-Z][a-z]*[A-Z][a-z]*\b', text)183 184 # Find important phrases (words near "important", "key", "main", etc.)185 importance_pattern = r'(?:important|key|main|primary|essential|critical|significant)\s+(\w+(?:\s+\w+){0,2})'186 important_phrases = re.findall(importance_pattern, text, re.IGNORECASE)187 188 # Combine and deduplicate189 all_phrases = noun_phrases + technical_terms + important_phrases190 unique_phrases = list(set(phrase.strip() for phrase in all_phrases if len(phrase.strip()) > 3))191 192 # Sort by length and return top phrases193 unique_phrases.sort(key=len, reverse=True)194 return unique_phrases[:max_phrases]195 196 def enhance_metadata(self, doc_chunk: DocumentChunk) -> Dict[str, Any]:197 """Enhance metadata with processed text statistics and extracted features"""198 enhanced_metadata = doc_chunk.metadata.copy()199 200 # Calculate text statistics201 text_stats = calculate_text_stats(doc_chunk.content)202 enhanced_metadata['text_stats'] = text_stats203 204 # Extract key phrases205 key_phrases = self.extract_key_phrases(doc_chunk.content)206 enhanced_metadata['extracted_phrases'] = key_phrases207 208 # Add processing metadata209 enhanced_metadata['processed_at'] = '2024-08-27' # Would use current timestamp210 enhanced_metadata['processor_version'] = '1.0'211 212 return enhanced_metadata213 214class ContentOptimizer:215 """Optimizes content for better retrieval and relevance"""216 217 def __init__(self):218 self.department_synonyms = {219 'BCIT': ['Computer Science', 'CS', 'BCIT'],220 'SE': ['Software Engineering', 'Software Engineering', 'SE'],221 'CIS': ['Computer Engineering', 'Computer Information Systems', 'CIS']222 }223 224 self.topic_synonyms = {225 'programming': ['coding', 'development', 'software development'],226 'theory': ['theoretical', 'academic', 'mathematical'],227 'practical': ['hands-on', 'applied', 'industry-focused'],228 'career': ['job', 'employment', 'profession', 'work'],229 'salary': ['pay', 'income', 'compensation', 'earnings']230 }231 232 def optimize_for_search(self, processed_chunk: ProcessedChunk) -> ProcessedChunk:233 """Optimize processed chunk for better search performance"""234 # Add synonym expansion to embedding text235 optimized_text = self._expand_synonyms(processed_chunk.embedding_text)236 237 # Add common questions context if available238 common_questions = processed_chunk.metadata.get('common_questions', [])239 if common_questions:240 questions_text = ' '.join(common_questions)241 optimized_text += f"\nRelated questions: {questions_text}"242 243 # Create optimized chunk244 optimized_chunk = ProcessedChunk(245 original_chunk_id=processed_chunk.original_chunk_id,246 chunk_index=processed_chunk.chunk_index,247 content=processed_chunk.content,248 metadata=processed_chunk.metadata,249 embedding_text=optimized_text,250 word_count=processed_chunk.word_count251 )252 253 return optimized_chunk254 255 def _expand_synonyms(self, text: str) -> str:256 """Expand text with relevant synonyms for better matching"""257 expanded_text = text258 259 # Add department synonyms260 for dept, synonyms in self.department_synonyms.items():261 if dept in text:262 expanded_text += f" {' '.join(synonyms)}"263 264 # Add topic synonyms265 for topic, synonyms in self.topic_synonyms.items():266 if topic in text.lower():267 expanded_text += f" {' '.join(synonyms)}"268 269 return expanded_text270 271# Global instances272text_processor = None273content_optimizer = None274 275def get_text_processor(chunk_size: int = 500, overlap: int = 50) -> TextProcessor:276 """Get or create the global text processor"""277 global text_processor278 if text_processor is None:279 text_processor = TextProcessor(chunk_size, overlap)280 return text_processor281 282def get_content_optimizer() -> ContentOptimizer:283 """Get or create the global content optimizer"""284 global content_optimizer285 if content_optimizer is None:286 content_optimizer = ContentOptimizer()287 return content_optimizer