main.py

#!/usr/bin/env python3
"""
Main Entry Point for Multimodal RAG System
Provides CLI interface and coordinates all system components
"""

import argparse
import logging
import sys
from pathlib import Path
import yaml
from typing import Dict, Any

# Import system components
from src.document_processor import DocumentProcessor
from src.image_processor import ImageProcessor
from src.audio_processor import AudioProcessor
from src.vector_store import VectorStore
from src.retrieval_engine import RetrievalEngine

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('multimodal_rag.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class MultimodalRAGSystem:
    """Main system coordinator for multimodal RAG"""
    
    def __init__(self, config_path: str = "config.yaml"):
        """Initialize the multimodal RAG system"""
        self.config = self._load_config(config_path)
        self._initialize_components()
        
    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """Load configuration from YAML file"""
        try:
            with open(config_path, 'r') as f:
                config = yaml.safe_load(f)
            logger.info(f"Loaded configuration from {config_path}")
            return config
        except FileNotFoundError:
            logger.error(f"Configuration file not found: {config_path}")
            sys.exit(1)
        except yaml.YAMLError as e:
            logger.error(f"Error parsing configuration file: {e}")
            sys.exit(1)
            
    def _initialize_components(self):
        """Initialize all system components"""
        try:
            logger.info("Initializing multimodal RAG system components...")
            
            # Initialize vector store
            self.vector_store = VectorStore(self.config)
            logger.info("✓ Vector store initialized")
            
            # Initialize processors
            self.doc_processor = DocumentProcessor(self.config)
            self.img_processor = ImageProcessor(self.config)
            self.audio_processor = AudioProcessor(self.config)
            logger.info("✓ Document, image, and audio processors initialized")
            
            # Initialize retrieval engine
            self.retrieval_engine = RetrievalEngine(self.config, self.vector_store)
            logger.info("✓ Retrieval engine initialized")
            
            logger.info("🚀 Multimodal RAG system ready!")
            
        except Exception as e:
            logger.error(f"Failed to initialize system: {e}")
            sys.exit(1)
            
    def process_files(self, directory_path: str, file_types: list = None):
        """Process all files in a directory"""
        directory = Path(directory_path)
        
        if not directory.exists():
            logger.error(f"Directory not found: {directory_path}")
            return
            
        logger.info(f"Processing files in: {directory_path}")
        
        total_chunks = 0
        
        # Process documents
        if not file_types or 'documents' in file_types:
            try:
                doc_chunks = self.doc_processor.batch_process(str(directory))
                if doc_chunks:
                    self.vector_store.add_documents(doc_chunks)
                    total_chunks += len(doc_chunks)
                    logger.info(f"✓ Processed {len(doc_chunks)} document chunks")
            except Exception as e:
                logger.error(f"Document processing failed: {e}")
                
        # Process images
        if not file_types or 'images' in file_types:
            try:
                img_chunks = self.img_processor.batch_process(str(directory))
                if img_chunks:
                    self.vector_store.add_documents(img_chunks)
                    total_chunks += len(img_chunks)
                    logger.info(f"✓ Processed {len(img_chunks)} image chunks")
            except Exception as e:
                logger.error(f"Image processing failed: {e}")
                
        # Process audio
        if not file_types or 'audio' in file_types:
            try:
                audio_chunks = self.audio_processor.batch_process(str(directory))
                if audio_chunks:
                    self.vector_store.add_documents(audio_chunks)
                    total_chunks += len(audio_chunks)
                    logger.info(f"✓ Processed {len(audio_chunks)} audio chunks")
            except Exception as e:
                logger.error(f"Audio processing failed: {e}")
                
        logger.info(f"🎉 Processing complete! Total chunks: {total_chunks}")
        
    def search(self, query: str, search_type: str = "text", top_k: int = 5):
        """Perform search and display results"""
        try:
            logger.info(f"Searching: '{query}' (type: {search_type})")
            
            results = self.retrieval_engine.search(
                query=query,
                search_type=search_type,
                top_k=top_k
            )
            
            if not results:
                print("No results found.")
                return
                
            print(f"\\n🔍 Found {len(results)} results:\\n")
            
            for i, result in enumerate(results, 1):
                metadata = result.get('metadata', {})
                content = result.get('content', '')
                
                print(f"[{i}] {metadata.get('file_name', 'Unknown')} "
                      f"(Score: {result['similarity_score']:.3f})")
                print(f"    Type: {metadata.get('modality', 'unknown')} | "
                      f"Format: {metadata.get('file_type', 'unknown')}")
                
                # Show content preview
                preview = content[:200] + "..." if len(content) > 200 else content
                print(f"    Content: {preview}")
                print()
                
        except Exception as e:
            logger.error(f"Search failed: {e}")
            
    def chat(self):
        """Start interactive chat session"""
        print("\\n💬 Multimodal RAG Chat Interface")
        print("Type your questions or 'quit' to exit.")
        print("-" * 50)
        
        while True:
            try:
                query = input("\\nYou: ").strip()
                
                if query.lower() in ['quit', 'exit', 'q']:
                    print("Goodbye!")
                    break
                    
                if not query:
                    continue
                    
                print("\\nAssistant: ", end="")
                
                # Generate response
                response = self.retrieval_engine.generate_response(
                    query=query,
                    top_k=self.config['search']['top_k']
                )
                
                print(response)
                
            except KeyboardInterrupt:
                print("\\n\\nGoodbye!")
                break
            except Exception as e:
                print(f"Error: {e}")
                
    def get_system_stats(self):
        """Display system statistics"""
        try:
            stats = self.vector_store.get_statistics()
            
            print("\\n📊 System Statistics:")
            print("-" * 30)
            print(f"Total Documents: {stats['total_documents']}")
            print(f"Vector Store: {stats['store_type']}")
            print(f"Collection: {stats['collection_name']}")
            print(f"Embedding Dimension: {stats['embedding_dimension']}")
            
            print("\\nBy Modality:")
            for modality, count in stats['modality_distribution'].items():
                print(f"  • {modality}: {count}")
                
            print("\\nBy File Type:")
            for file_type, count in stats['file_type_distribution'].items():
                print(f"  • {file_type}: {count}")
                
            print(f"\\nStorage Path: {stats['storage_path']}")
            
        except Exception as e:
            logger.error(f"Failed to get stats: {e}")

def main():
    """Main CLI entry point"""
    parser = argparse.ArgumentParser(
        description="Multimodal RAG System - Offline semantic search across documents, images, and audio"
    )
    
    parser.add_argument(
        "--config", 
        default="config.yaml",
        help="Path to configuration file"
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Available commands')
    
    # Process command
    process_parser = subparsers.add_parser('process', help='Process files')
    process_parser.add_argument('directory', help='Directory to process')
    process_parser.add_argument(
        '--types', 
        nargs='+', 
        choices=['documents', 'images', 'audio'],
        help='File types to process (default: all)'
    )
    
    # Search command
    search_parser = subparsers.add_parser('search', help='Search for content')
    search_parser.add_argument('query', help='Search query')
    search_parser.add_argument(
        '--type', 
        default='text',
        choices=['text', 'image', 'hybrid'],
        help='Search type'
    )
    search_parser.add_argument(
        '--top-k', 
        type=int, 
        default=5,
        help='Number of results to return'
    )
    
    # Chat command
    chat_parser = subparsers.add_parser('chat', help='Start interactive chat')
    
    # Stats command
    stats_parser = subparsers.add_parser('stats', help='Show system statistics')
    
    # UI command
    ui_parser = subparsers.add_parser('ui', help='Launch web interface')
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return
        
    # Initialize system
    rag_system = MultimodalRAGSystem(args.config)
    
    # Execute command
    if args.command == 'process':
        rag_system.process_files(args.directory, args.types)
        
    elif args.command == 'search':
        rag_system.search(args.query, args.type, args.top_k)
        
    elif args.command == 'chat':
        rag_system.chat()
        
    elif args.command == 'stats':
        rag_system.get_system_stats()
        
    elif args.command == 'ui':
        # Launch Streamlit UI
        import subprocess
        import sys
        
        try:
            subprocess.run([
                sys.executable, "-m", "streamlit", "run", 
                "src/ui_interface.py", "--server.port", "8501"
            ])
        except KeyboardInterrupt:
            print("\\nShutting down web interface...")

if __name__ == "__main__":
    main()