cli.py

#!/usr/bin/env python3
"""
CLI Interface for RAG System
Supports Arabic/English PDF parsing with accurate retrieval
"""

import os
import sys
from pathlib import Path
from typing import Optional
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.prompt import Prompt, Confirm
from rich.table import Table
from rich import print as rprint
from dotenv import load_dotenv

from pdf_parser import PDFParser
from vectorstore import VectorStore
from rag_engine import RAGEngine


class RAGCLI:
    """CLI interface for the RAG system"""
    
    def __init__(self):
        self.console = Console()
        self.pdf_parser = PDFParser()
        self.vectorstore: Optional[VectorStore] = None
        self.rag_engine: Optional[RAGEngine] = None
        
        # Load configuration
        self._load_config()
        
        # Initialize components
        self._initialize_components()
    
    def _load_config(self):
        """Load configuration from environment"""
        load_dotenv()
        
        self.ollama_base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        self.ollama_llm_model = os.getenv("OLLAMA_LLM_MODEL", "llama3.1")
        self.ollama_embedding_model = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
        self.chroma_persist_directory = os.getenv("CHROMA_PERSIST_DIRECTORY", "./chroma_db")
        self.chunk_size = int(os.getenv("CHUNK_SIZE", "1000"))
        self.chunk_overlap = int(os.getenv("CHUNK_OVERLAP", "200"))
    
    def _initialize_components(self):
        """Initialize vector store and RAG engine"""
        try:
            self.vectorstore = VectorStore(
                ollama_base_url=self.ollama_base_url,
                embedding_model=self.ollama_embedding_model,
                persist_directory=self.chroma_persist_directory,
                chunk_size=self.chunk_size,
                chunk_overlap=self.chunk_overlap
            )
            
            self.rag_engine = RAGEngine(
                vectorstore=self.vectorstore,
                ollama_base_url=self.ollama_base_url,
                llm_model=self.ollama_llm_model
            )
            
            self.console.print("[green]✓[/green] RAG system initialized successfully")
        except Exception as e:
            self.console.print(f"[red]Error initializing system: {str(e)}[/red]")
            sys.exit(1)
    
    def show_banner(self):
        """Display welcome banner"""
        banner = """
╔═══════════════════════════════════════════════════════════╗
║         RAG System - Arabic/English PDF Support          ║
║                   Powered by Ollama                       ║
╚═══════════════════════════════════════════════════════════╝
        """
        self.console.print(banner, style="bold cyan")
    
    def show_menu(self):
        """Display main menu"""
        table = Table(title="Main Menu", show_header=True, header_style="bold magenta")
        table.add_column("Option", style="cyan", width=10)
        table.add_column("Description", style="white")
        
        table.add_row("1", "Add PDF document")
        table.add_row("2", "Ask a question")
        table.add_row("3", "View statistics")
        table.add_row("4", "Clear conversation history")
        table.add_row("5", "Clear all documents")
        table.add_row("6", "Exit")
        
        self.console.print(table)
    
    def add_pdf(self):
        """Add a PDF document to the system"""
        pdf_path = Prompt.ask("Enter PDF file path")
        
        if not os.path.exists(pdf_path):
            self.console.print(f"[red]Error: File not found: {pdf_path}[/red]")
            return
        
        try:
            # Show PDF info
            with self.console.status("[bold green]Extracting PDF information..."):
                pdf_info = self.pdf_parser.get_pdf_info(pdf_path)
            
            info_table = Table(title="PDF Information", show_header=False)
            info_table.add_column("Property", style="cyan")
            info_table.add_column("Value", style="white")
            
            for key, value in pdf_info.items():
                info_table.add_row(key.capitalize(), str(value))
            
            self.console.print(info_table)
            
            # Confirm processing
            if not Confirm.ask("Process this PDF?", default=True):
                return
            
            # Parse PDF
            with self.console.status("[bold green]Parsing PDF and extracting text..."):
                pages_content = self.pdf_parser.extract_text_with_pages(pdf_path)
            
            self.console.print(f"[green]✓[/green] Extracted text from {len(pages_content)} pages")
            
            # Add to vector store
            with self.console.status("[bold green]Creating embeddings and indexing..."):
                source_name = Path(pdf_path).name
                self.vectorstore.add_documents(pages_content, source_name)
            
            self.console.print(f"[green]✓[/green] Successfully indexed {source_name}")
            
        except Exception as e:
            self.console.print(f"[red]Error processing PDF: {str(e)}[/red]")
    
    def ask_question(self):
        """Interactive question-answering mode"""
        self.console.print("\n[bold cyan]Question Mode[/bold cyan] (type 'back' to return to menu)\n")
        
        while True:
            question = Prompt.ask("[bold yellow]Your question[/bold yellow]")
            
            if question.lower() in ['back', 'exit', 'quit']:
                break
            
            if not question.strip():
                continue
            
            # Get answer
            with self.console.status("[bold green]Searching and generating answer..."):
                result = self.rag_engine.query(question, k=4, include_sources=True)
            
            # Display answer
            answer_panel = Panel(
                result["answer"],
                title="[bold green]Answer[/bold green]",
                border_style="green"
            )
            self.console.print(answer_panel)
            
            # Display page numbers
            if result["pages"]:
                pages_str = ", ".join([f"Page {p}" for p in result["pages"]])
                self.console.print(f"\n[cyan]📄 Relevant pages: {pages_str}[/cyan]")
            
            # Show detailed sources if available
            if result["sources"]:
                show_sources = Confirm.ask("\nShow detailed sources?", default=False)
                if show_sources:
                    sources_table = Table(title="Sources", show_header=True)
                    sources_table.add_column("Page", style="cyan", width=8)
                    sources_table.add_column("Content Preview", style="white")
                    sources_table.add_column("Relevance", style="yellow", width=10)
                    
                    for source in result["sources"]:
                        preview = source["content"][:100] + "..." if len(source["content"]) > 100 else source["content"]
                        relevance = f"{(1 - source['score']):.2%}"  # Convert distance to similarity
                        sources_table.add_row(
                            str(source["page"]),
                            preview,
                            relevance
                        )
                    
                    self.console.print(sources_table)
            
            self.console.print()  # Empty line for spacing
    
    def show_stats(self):
        """Display system statistics"""
        stats = self.vectorstore.get_stats()
        history_count = len(self.rag_engine.get_history())
        
        stats_table = Table(title="System Statistics", show_header=False)
        stats_table.add_column("Metric", style="cyan")
        stats_table.add_column("Value", style="white")
        
        stats_table.add_row("Documents in vector store", str(stats.get("documents", 0)))
        stats_table.add_row("Conversation history", f"{history_count} exchanges")
        stats_table.add_row("LLM Model", self.ollama_llm_model)
        stats_table.add_row("Embedding Model", self.ollama_embedding_model)
        stats_table.add_row("Chunk Size", str(self.chunk_size))
        
        self.console.print(stats_table)
    
    def clear_history(self):
        """Clear conversation history"""
        if Confirm.ask("Clear conversation history?", default=False):
            self.rag_engine.clear_history()
            self.console.print("[green]✓[/green] Conversation history cleared")
    
    def clear_documents(self):
        """Clear all documents from vector store"""
        if Confirm.ask("⚠️  Clear ALL documents? This cannot be undone!", default=False):
            self.vectorstore.clear_store()
            self.console.print("[green]✓[/green] All documents cleared")
    
    def run(self):
        """Run the CLI application"""
        self.show_banner()
        
        while True:
            self.console.print()  # Empty line
            self.show_menu()
            
            choice = Prompt.ask("Select an option", choices=["1", "2", "3", "4", "5", "6"])
            
            if choice == "1":
                self.add_pdf()
            elif choice == "2":
                self.ask_question()
            elif choice == "3":
                self.show_stats()
            elif choice == "4":
                self.clear_history()
            elif choice == "5":
                self.clear_documents()
            elif choice == "6":
                if Confirm.ask("Exit?", default=True):
                    self.console.print("\n[cyan]Goodbye! 👋[/cyan]\n")
                    break


def main():
    """Main entry point"""
    try:
        cli = RAGCLI()
        cli.run()
    except KeyboardInterrupt:
        print("\n\nInterrupted by user. Goodbye! 👋\n")
        sys.exit(0)
    except Exception as e:
        print(f"\n[red]Fatal error: {str(e)}[/red]\n")
        sys.exit(1)


if __name__ == "__main__":
    main()