MMo4/csit-ned-chatbot

1from pathlib import Path2from typing import List, Dict, Any, Optional, Tuple3import logging4from dataclasses import dataclass5import yaml6from src.utils.helpers import extract_yaml_metadata, validate_metadata, get_file_paths7 8logger = logging.getLogger(__name__)9 10@dataclass11class DocumentChunk:12    """Represents a loaded document chunk with metadata"""13    chunk_id: str14    title: str15    content: str16    metadata: Dict[str, Any]17    file_path: str18    19    def to_dict(self) -> Dict[str, Any]:20        """Convert to dictionary representation"""21        return {22            'chunk_id': self.chunk_id,23            'title': self.title,24            'content': self.content,25            'metadata': self.metadata,26            'file_path': self.file_path27        }28 29class KnowledgeBaseLoader:30    """Loads and processes markdown documents from the knowledge base"""31    32    def __init__(self, knowledge_base_path: str):33        self.knowledge_base_path = Path(knowledge_base_path)34        self.chunks: List[DocumentChunk] = []35        36        if not self.knowledge_base_path.exists():37            logger.warning(f"Knowledge base path does not exist: {knowledge_base_path}")38    39    def load_all_documents(self) -> List[DocumentChunk]:40        """Load all markdown documents from the knowledge base"""41        logger.info(f"Loading documents from: {self.knowledge_base_path}")42        43        # Get all markdown files44        md_files = get_file_paths(self.knowledge_base_path, "*.md")45        46        if not md_files:47            logger.warning("No markdown files found in knowledge base")48            return []49        50        logger.info(f"Found {len(md_files)} markdown files")51        52        loaded_chunks = []53        for file_path in md_files:54            try:55                chunk = self._load_single_document(file_path)56                if chunk:57                    loaded_chunks.append(chunk)58                    logger.debug(f"Loaded: {chunk.chunk_id}")59            except Exception as e:60                logger.error(f"Failed to load {file_path}: {e}")61                continue62        63        self.chunks = loaded_chunks64        logger.info(f"Successfully loaded {len(loaded_chunks)} document chunks")65        return loaded_chunks66    67    def _load_single_document(self, file_path: Path) -> Optional[DocumentChunk]:68        """Load a single markdown document"""69        try:70            with open(file_path, 'r', encoding='utf-8') as f:71                content = f.read()72            73            # Extract YAML metadata and content74            metadata, markdown_content = extract_yaml_metadata(content)75            76            if not metadata:77                logger.warning(f"No metadata found in {file_path}")78                return None79            80            # Validate metadata81            if not validate_metadata(metadata):82                logger.warning(f"Invalid metadata in {file_path}")83                return None84            85            # Create document chunk86            chunk = DocumentChunk(87                chunk_id=metadata['chunk_id'],88                title=metadata['title'],89                content=markdown_content,90                metadata=metadata,91                file_path=str(file_path)92            )93            94            return chunk95            96        except Exception as e:97            logger.error(f"Error loading {file_path}: {e}")98            return None99    100    def load_by_category(self, category: str) -> List[DocumentChunk]:101        """Load documents from a specific category"""102        if not self.chunks:103            self.load_all_documents()104        105        category_chunks = [106            chunk for chunk in self.chunks 107            if chunk.metadata.get('category') == category108        ]109        110        logger.info(f"Found {len(category_chunks)} chunks in category: {category}")111        return category_chunks112    113    def load_by_department(self, department: str) -> List[DocumentChunk]:114        """Load documents related to a specific department"""115        if not self.chunks:116            self.load_all_documents()117        118        department_chunks = [119            chunk for chunk in self.chunks 120            if department in chunk.metadata.get('departments', [])121        ]122        123        logger.info(f"Found {len(department_chunks)} chunks for department: {department}")124        return department_chunks125    126    def search_by_keywords(self, keywords: List[str]) -> List[DocumentChunk]:127        """Search documents by keywords in metadata"""128        if not self.chunks:129            self.load_all_documents()130        131        keyword_set = set(word.lower() for word in keywords)132        matching_chunks = []133        134        for chunk in self.chunks:135            chunk_keywords = set(chunk.metadata.get('search_keywords', []))136            chunk_topics = set(chunk.metadata.get('topics', []))137            138            # Check if any keyword matches139            if keyword_set & (chunk_keywords | chunk_topics):140                matching_chunks.append(chunk)141        142        logger.info(f"Found {len(matching_chunks)} chunks matching keywords: {keywords}")143        return matching_chunks144    145    def get_chunk_by_id(self, chunk_id: str) -> Optional[DocumentChunk]:146        """Get a specific chunk by its ID"""147        if not self.chunks:148            self.load_all_documents()149        150        for chunk in self.chunks:151            if chunk.chunk_id == chunk_id:152                return chunk153        154        return None155    156    def get_related_chunks(self, chunk_id: str) -> List[DocumentChunk]:157        """Get chunks related to a specific chunk"""158        chunk = self.get_chunk_by_id(chunk_id)159        if not chunk:160            return []161        162        related_ids = chunk.metadata.get('related_chunks', [])163        related_chunks = []164        165        for related_id in related_ids:166            related_chunk = self.get_chunk_by_id(related_id)167            if related_chunk:168                related_chunks.append(related_chunk)169        170        return related_chunks171    172    def validate_all_chunks(self) -> Dict[str, Any]:173        """Validate all loaded chunks and return statistics"""174        if not self.chunks:175            self.load_all_documents()176        177        validation_stats = {178            'total_chunks': len(self.chunks),179            'valid_chunks': 0,180            'invalid_chunks': 0,181            'missing_metadata': 0,182            'categories': {},183            'departments': {},184            'validation_errors': []185        }186        187        for chunk in self.chunks:188            try:189                if validate_metadata(chunk.metadata):190                    validation_stats['valid_chunks'] += 1191                    192                    # Count categories and departments193                    category = chunk.metadata.get('category')194                    if category:195                        validation_stats['categories'][category] = validation_stats['categories'].get(category, 0) + 1196                    197                    departments = chunk.metadata.get('departments', [])198                    for dept in departments:199                        validation_stats['departments'][dept] = validation_stats['departments'].get(dept, 0) + 1200                else:201                    validation_stats['invalid_chunks'] += 1202                    validation_stats['validation_errors'].append(f"Invalid metadata in {chunk.chunk_id}")203                    204            except Exception as e:205                validation_stats['invalid_chunks'] += 1206                validation_stats['validation_errors'].append(f"Error validating {chunk.chunk_id}: {e}")207        208        logger.info(f"Validation complete: {validation_stats['valid_chunks']} valid, {validation_stats['invalid_chunks']} invalid")209        return validation_stats210    211    def export_chunks_data(self) -> List[Dict[str, Any]]:212        """Export all chunks as list of dictionaries"""213        if not self.chunks:214            self.load_all_documents()215        216        return [chunk.to_dict() for chunk in self.chunks]217 218class MetadataManager:219    """Manager for knowledge base metadata files"""220    221    def __init__(self, metadata_path: str):222        self.metadata_path = Path(metadata_path)223        self.content_taxonomy = {}224        self.question_mapping = {}225    226    def load_content_taxonomy(self) -> Dict[str, Any]:227        """Load content taxonomy YAML file"""228        taxonomy_file = self.metadata_path / "content_taxonomy.yaml"229        230        if not taxonomy_file.exists():231            logger.warning(f"Content taxonomy file not found: {taxonomy_file}")232            return {}233        234        try:235            with open(taxonomy_file, 'r', encoding='utf-8') as f:236                self.content_taxonomy = yaml.safe_load(f) or {}237            logger.info("Content taxonomy loaded successfully")238            return self.content_taxonomy239        except Exception as e:240            logger.error(f"Failed to load content taxonomy: {e}")241            return {}242    243    def load_question_mapping(self) -> Dict[str, Any]:244        """Load question mapping YAML file"""245        mapping_file = self.metadata_path / "question_mapping.yaml"246        247        if not mapping_file.exists():248            logger.warning(f"Question mapping file not found: {mapping_file}")249            return {}250        251        try:252            with open(mapping_file, 'r', encoding='utf-8') as f:253                self.question_mapping = yaml.safe_load(f) or {}254            logger.info("Question mapping loaded successfully")255            return self.question_mapping256        except Exception as e:257            logger.error(f"Failed to load question mapping: {e}")258            return {}259    260    def get_category_info(self, category: str) -> Dict[str, Any]:261        """Get information about a specific category"""262        if not self.content_taxonomy:263            self.load_content_taxonomy()264        265        return self.content_taxonomy.get('categories', {}).get(category, {})266    267    def get_common_questions(self, category: str = None, department: str = None) -> List[str]:268        """Get common questions for category or department"""269        if not self.question_mapping:270            self.load_question_mapping()271        272        questions = []273        274        if category:275            questions.extend(self.question_mapping.get('by_category', {}).get(category, []))276        277        if department:278            questions.extend(self.question_mapping.get('by_department', {}).get(department, []))279        280        return questions281 282# Global instances283knowledge_loader = None284metadata_manager = None285 286def get_knowledge_loader(knowledge_base_path: str = None) -> KnowledgeBaseLoader:287    """Get or create the global knowledge base loader"""288    global knowledge_loader289    if knowledge_loader is None or knowledge_base_path:290        from src.config import settings291        path = knowledge_base_path or settings.knowledge_base_path292        knowledge_loader = KnowledgeBaseLoader(path)293    return knowledge_loader294 295def get_metadata_manager(metadata_path: str = None) -> MetadataManager:296    """Get or create the global metadata manager"""297    global metadata_manager298    if metadata_manager is None or metadata_path:299        from src.config import settings300        path = metadata_path or settings.metadata_path301        metadata_manager = MetadataManager(path)302    return metadata_manager