rag_system.py

"""
Main RAG System - Orchestrates All Components
Integrates document processing, vector database, query parsing, and LLM reasoning
"""

import os
import json
import logging
import time
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path

# Import our custom components
from document_processer import AdvancedDocumentProcessor, DocumentChunk
from vector_database import VectorDatabase, SearchResult
from query_parser import AdvancedQueryParser, ParsedQuery
from llm_reasoning import AdvancedLLMReasoning, ReasoningResult

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class QueryResult:
    """Represents the complete result of a query processing"""
    query: str
    parsed_query: ParsedQuery
    search_results: List[SearchResult]
    reasoning_result: ReasoningResult
    processing_time: float
    timestamp: datetime
    audit_trail: Dict[str, Any]

class AdvancedRAGSystem:
    """Advanced RAG system that orchestrates all components"""
    
    def __init__(self, 
                 model_path: str = "./mistral-7b-instruct-v0.1.Q4_K_M.gguf",
                 use_gpu: bool = True,
                 vector_db_path: str = "./vector_db"):
        
        self.model_path = model_path
        self.use_gpu = use_gpu
        self.vector_db_path = vector_db_path
        
        # Initialize components
        self._initialize_components()
        
        # Audit trail storage
        self.audit_log = []
        
        logger.info("Advanced RAG System initialized successfully")
    
    def _initialize_components(self):
        """Initialize all system components"""
        try:
            # Initialize document processor
            self.document_processor = AdvancedDocumentProcessor(
                ocr_language='eng',
                chunk_size=1000,
                chunk_overlap=200
            )
            
            # Initialize vector database
            self.vector_database = VectorDatabase(
                embedding_model="all-MiniLM-L6-v2",
                collection_name="documents",
                persist_directory=self.vector_db_path,
                use_gpu=self.use_gpu
            )
            
            # Initialize query parser (using NLTK instead of spaCy)
            self.query_parser = AdvancedQueryParser(
                use_gpu=self.use_gpu
            )
            
            # Initialize LLM reasoning engine
            self.reasoning_engine = AdvancedLLMReasoning(
                model_path=self.model_path,
                use_gpu=self.use_gpu,
                max_tokens=2048
            )
            
            logger.info("All components initialized successfully")
            
        except Exception as e:
            logger.error(f"Error initializing components: {e}")
            raise
    
    def ingest_document(self, file_path: str, use_ocr: bool = False) -> List[DocumentChunk]:
        """Ingest and process a document"""
        try:
            logger.info(f"Starting document ingestion: {file_path}")
            
            # Process document
            chunks = self.document_processor.process_document(file_path, use_ocr)
            logger.info(f"Document processor created {len(chunks)} chunks")
            
            if not chunks:
                logger.warning("No chunks created by document processor")
                return []
            
            # Add to vector database
            logger.info(f"Adding {len(chunks)} chunks to vector database...")
            success = self.vector_database.add_documents(chunks)
            logger.info(f"Vector database add_documents returned: {success}")
            
            if success:
                logger.info(f"Successfully ingested {len(chunks)} chunks from {file_path}")
                
                # Add to audit trail
                self._add_audit_entry({
                    'action': 'document_ingestion',
                    'file_path': file_path,
                    'chunks_processed': len(chunks),
                    'use_ocr': use_ocr,
                    'timestamp': datetime.now().isoformat(),
                    'status': 'success'
                })
                
                return chunks
            else:
                logger.error(f"Failed to add documents to vector database, but returning chunks anyway")
                # Return chunks even if vector database fails, so the user can still see the processing worked
                return chunks
                
        except Exception as e:
            logger.error(f"Error ingesting document {file_path}: {e}")
            
            # Add error to audit trail
            self._add_audit_entry({
                'action': 'document_ingestion',
                'file_path': file_path,
                'error': str(e),
                'timestamp': datetime.now().isoformat(),
                'status': 'error'
            })
            
            raise
    
    def process_query(self, query: str, n_results: int = 5) -> QueryResult:
        """Process a natural language query"""
        try:
            start_time = time.time()
            
            logger.info(f"Processing query: {query}")
            
            # Step 1: Parse the query
            parsed_query = self.query_parser.parse_query(query)
            
            # Step 2: Search for relevant documents
            search_results = self.vector_database.hybrid_search(
                query=parsed_query.enhanced_query,
                n_results=n_results,
                semantic_weight=0.7,
                keyword_weight=0.3
            )
            
            # Step 3: Prepare context for reasoning
            context = self._prepare_context_for_reasoning(search_results)
            
            # Step 4: Analyze with LLM reasoning
            reasoning_result = self.reasoning_engine.analyze_query(
                query=query,
                context=context,
                query_type=parsed_query.query_type
            )
            
            processing_time = time.time() - start_time
            
            # Step 5: Create audit trail
            audit_trail = self._create_audit_trail(
                query, parsed_query, search_results, reasoning_result, processing_time
            )
            
            # Step 6: Build result
            result = QueryResult(
                query=query,
                parsed_query=parsed_query,
                search_results=search_results,
                reasoning_result=reasoning_result,
                processing_time=processing_time,
                timestamp=datetime.now(),
                audit_trail=audit_trail
            )
            
            # Add to audit log
            self._add_audit_entry(audit_trail)
            
            logger.info(f"Query processed successfully in {processing_time:.2f}s")
            return result
            
        except Exception as e:
            logger.error(f"Error processing query: {e}")
            
            # Create fallback result
            return self._create_fallback_result(query, str(e))
    
    def _prepare_context_for_reasoning(self, search_results: List[SearchResult]) -> List[Dict[str, Any]]:
        """Prepare search results for LLM reasoning"""
        try:
            context = []
            
            for result in search_results:
                context_item = {
                    'content': result.content,
                    'source_file': result.source_file,
                    'similarity_score': result.similarity_score,
                    'section_type': result.section_type,
                    'metadata': result.metadata
                }
                
                # Add table data if present
                if result.table_data:
                    context_item['table_data'] = result.table_data
                
                context.append(context_item)
            
            return context
            
        except Exception as e:
            logger.error(f"Error preparing context: {e}")
            return []
    
    def _create_audit_trail(self, 
                           query: str, 
                           parsed_query: ParsedQuery,
                           search_results: List[SearchResult],
                           reasoning_result: ReasoningResult,
                           processing_time: float) -> Dict[str, Any]:
        """Create comprehensive audit trail"""
        try:
            audit_trail = {
                'action': 'query_processing',
                'query': query,
                'parsed_query': {
                    'query_type': parsed_query.query_type,
                    'intent': parsed_query.intent,
                    'confidence': parsed_query.confidence,
                    'entities': parsed_query.entities,
                    'keywords': parsed_query.keywords
                },
                'search_results': {
                    'count': len(search_results),
                    'top_results': [
                        {
                            'content_preview': result.content[:100] + "...",
                            'source_file': result.source_file,
                            'similarity_score': result.similarity_score,
                            'section_type': result.section_type
                        }
                        for result in search_results[:3]
                    ]
                },
                'reasoning_result': {
                    'decision': reasoning_result.decision,
                    'confidence_score': reasoning_result.confidence_score,
                    'relevant_clauses': reasoning_result.relevant_clauses,
                    'amount': reasoning_result.amount,
                    'waiting_period': reasoning_result.waiting_period
                },
                'processing_time': processing_time,
                'timestamp': datetime.now().isoformat(),
                'status': 'success'
            }
            
            return audit_trail
            
        except Exception as e:
            logger.error(f"Error creating audit trail: {e}")
            return {
                'action': 'query_processing',
                'query': query,
                'error': str(e),
                'timestamp': datetime.now().isoformat(),
                'status': 'error'
            }
    
    def _create_fallback_result(self, query: str, error: str) -> QueryResult:
        """Create a fallback result when processing fails"""
        try:
            # Create basic parsed query
            parsed_query = ParsedQuery(
                original_query=query,
                enhanced_query=query,
                query_type='general_inquiry',
                entities={},
                intent='information_seeking',
                confidence=0.0,
                keywords=[],
                synonyms=[],
                context={},
                timestamp=datetime.now()
            )
            
            # Create fallback reasoning result
            reasoning_result = ReasoningResult(
                decision='pending',
                confidence_score=0.0,
                justification=f'Processing failed: {error}',
                relevant_clauses=[],
                reasoning_steps=['Processing failed'],
                source_references=[]
            )
            
            return QueryResult(
                query=query,
                parsed_query=parsed_query,
                search_results=[],
                reasoning_result=reasoning_result,
                processing_time=0.0,
                timestamp=datetime.now(),
                audit_trail={
                    'action': 'query_processing',
                    'query': query,
                    'error': error,
                    'timestamp': datetime.now().isoformat(),
                    'status': 'error'
                }
            )
            
        except Exception as e:
            logger.error(f"Error creating fallback result: {e}")
            raise
    
    def _add_audit_entry(self, entry: Dict[str, Any]):
        """Add entry to audit log"""
        try:
            self.audit_log.append(entry)
            
            # Keep audit log size manageable
            if len(self.audit_log) > 1000:
                self.audit_log = self.audit_log[-500:]
                
        except Exception as e:
            logger.error(f"Error adding audit entry: {e}")
    
    def get_audit_trail(self) -> List[Dict[str, Any]]:
        """Get the complete audit trail"""
        return self.audit_log.copy()
    
    def save_audit_trail(self, file_path: str) -> bool:
        """Save audit trail to file"""
        try:
            with open(file_path, 'w') as f:
                json.dump(self.audit_log, f, indent=2)
            
            logger.info(f"Audit trail saved to: {file_path}")
            return True
            
        except Exception as e:
            logger.error(f"Error saving audit trail: {e}")
            return False
    
    def get_system_statistics(self) -> Dict[str, Any]:
        """Get comprehensive system statistics"""
        try:
            # Get vector database statistics
            db_stats = self.vector_database.get_document_statistics()
            
            # Get audit trail statistics
            audit_stats = {
                'total_entries': len(self.audit_log),
                'successful_queries': len([e for e in self.audit_log if e.get('status') == 'success']),
                'failed_queries': len([e for e in self.audit_log if e.get('status') == 'error']),
                'document_ingestions': len([e for e in self.audit_log if e.get('action') == 'document_ingestion']),
                'query_processings': len([e for e in self.audit_log if e.get('action') == 'query_processing'])
            }
            
            # Get component information
            component_info = {
                'document_processor': 'AdvancedDocumentProcessor',
                'vector_database': 'AdvancedVectorDatabase',
                'query_parser': 'AdvancedQueryParser',
                'reasoning_engine': 'AdvancedLLMReasoning',
                'model_path': self.model_path,
                'use_gpu': self.use_gpu
            }
            
            stats = {
                'vector_database': db_stats,
                'audit_trail': audit_stats,
                'components': component_info,
                'timestamp': datetime.now().isoformat()
            }
            
            return stats
            
        except Exception as e:
            logger.error(f"Error getting system statistics: {e}")
            return {}
    
    def clear_system(self) -> bool:
        """Clear all data from the system"""
        try:
            # Clear vector database
            self.vector_database.clear_database()
            
            # Clear audit log
            self.audit_log = []
            
            logger.info("System cleared successfully")
            return True
            
        except Exception as e:
            logger.error(f"Error clearing system: {e}")
            return False
    
    def export_system_data(self, export_path: str) -> bool:
        """Export system data for backup or analysis"""
        try:
            # Get system statistics
            stats = self.get_system_statistics()
            
            # Add audit trail
            export_data = {
                'statistics': stats,
                'audit_trail': self.audit_log,
                'export_timestamp': datetime.now().isoformat()
            }
            
            with open(export_path, 'w') as f:
                json.dump(export_data, f, indent=2)
            
            logger.info(f"System data exported to: {export_path}")
            return True
            
        except Exception as e:
            logger.error(f"Error exporting system data: {e}")
            return False
    
    def validate_system(self) -> Dict[str, Any]:
        """Validate system components and return status"""
        try:
            validation_results = {
                'document_processor': True,
                'vector_database': True,
                'query_parser': True,
                'reasoning_engine': True,
                'overall_status': True,
                'errors': []
            }
            
            # Test document processor
            try:
                # This is a basic test - in practice you might want more comprehensive tests
                pass
            except Exception as e:
                validation_results['document_processor'] = False
                validation_results['errors'].append(f"Document processor: {e}")
            
            # Test vector database
            try:
                stats = self.vector_database.get_document_statistics()
            except Exception as e:
                validation_results['vector_database'] = False
                validation_results['errors'].append(f"Vector database: {e}")
            
            # Test query parser
            try:
                test_parsed = self.query_parser.parse_query("test query")
            except Exception as e:
                validation_results['query_parser'] = False
                validation_results['errors'].append(f"Query parser: {e}")
            
            # Test reasoning engine
            try:
                # Basic test - check if model file exists
                if not os.path.exists(self.model_path):
                    validation_results['reasoning_engine'] = False
                    validation_results['errors'].append("LLM model file not found")
            except Exception as e:
                validation_results['reasoning_engine'] = False
                validation_results['errors'].append(f"Reasoning engine: {e}")
            
            # Overall status
            validation_results['overall_status'] = all([
                validation_results['document_processor'],
                validation_results['vector_database'],
                validation_results['query_parser'],
                validation_results['reasoning_engine']
            ])
            
            return validation_results
            
        except Exception as e:
            logger.error(f"Error validating system: {e}")
            return {
                'overall_status': False,
                'errors': [f"Validation failed: {e}"]
            }

# Example usage
if __name__ == "__main__":
    # Initialize RAG system
    rag_system = AdvancedRAGSystem(use_gpu=True)
    
    # Test system validation
    validation = rag_system.validate_system()
    print(f"System validation: {validation['overall_status']}")
    
    if validation['overall_status']:
        print("✅ All components are working correctly")
    else:
        print("❌ Some components have issues:")
        for error in validation['errors']:
            print(f"  - {error}")
    
    # Get system statistics
    stats = rag_system.get_system_statistics()
    print(f"\nSystem statistics: {stats}")