ai-lab-transformers-playground/src/commands/qa.py

from src.cli.base import CLICommand
from src.cli.display import DisplayFormatter
from src.pipelines.qa import QuestionAnsweringSystem


class QACommand(CLICommand):
    """Interactive Question Answering command"""
    
    def __init__(self):
        self.qa_system = None
        self.current_context = None
        self.session_questions = []
    
    @property
    def name(self) -> str:
        return "qa"
    
    @property
    def description(self) -> str:
        return "Question Answering - Ask questions about a given text"
    
    def _initialize_qa_system(self):
        """Lazy initialization of the QA system"""
        if self.qa_system is None:
            print("🔄 Loading Question Answering model...")
            self.qa_system = QuestionAnsweringSystem()
            DisplayFormatter.show_success("QA model loaded!")
    
    def _show_instructions(self):
        """Show usage instructions and examples"""
        print("\n❓ Question Answering System")
        print("Ask questions about a text context and get precise answers.")
        print("\n📝 How it works:")
        print("  1. First, provide a context (text containing information)")
        print("  2. Then ask questions about that context")
        print("  3. The system extracts answers directly from the text")
        print("\n💡 Example context:")
        print("  'Albert Einstein was born in 1879 in Germany. He developed the theory of relativity.'")
        print("💡 Example questions:")
        print("  - When was Einstein born?")
        print("  - Where was Einstein born?")
        print("  - What theory did Einstein develop?")
        print("\n🎛️  Commands:")
        print("  'back' - Return to main menu")
        print("  'help' - Show these instructions")
        print("  'context' - Set new context")
        print("  'multi' - Ask multiple questions at once")
        print("  'session' - Review session history")
        print("  'settings' - Adjust confidence threshold")
        print("-" * 70)
    
    def _set_context(self):
        """Allow user to set or change the context"""
        print("\n📄 Set Context")
        print("Enter the text that will serve as context for your questions.")
        print("You can enter multiple lines. Type 'done' when finished.")
        print("-" * 50)
        
        lines = []
        while True:
            line = input("📝 ").strip()
            if line.lower() == 'done':
                break
            if line:
                lines.append(line)
        
        if not lines:
            DisplayFormatter.show_warning("No context provided")
            return False
        
        self.current_context = " ".join(lines)
        
        # Analyze context
        analysis = self.qa_system.interactive_qa(self.current_context)
        if "error" in analysis:
            DisplayFormatter.show_error(analysis["error"])
            return False
        
        formatted_analysis = DisplayFormatter.format_qa_context_analysis(analysis)
        print(formatted_analysis)
        
        return True
    
    def _ask_single_question(self):
        """Ask a single question about the current context"""
        if not self.current_context:
            DisplayFormatter.show_warning("Please set a context first using 'context' command")
            return
        
        question = input("\n❓ Your question: ").strip()
        
        if not question:
            DisplayFormatter.show_warning("Please enter a question")
            return
        
        DisplayFormatter.show_loading("Finding answer...")
        result = self.qa_system.answer(question, self.current_context)
        
        if "error" not in result:
            self.session_questions.append(result)
        
        formatted_result = DisplayFormatter.format_qa_result(result)
        print(formatted_result)
    
    def _multi_question_mode(self):
        """Allow asking multiple questions at once"""
        if not self.current_context:
            DisplayFormatter.show_warning("Please set a context first using 'context' command")
            return
        
        print("\n❓ Multiple Questions Mode")
        print("Enter your questions one by one. Type 'done' when finished.")
        print("-" * 50)
        
        questions = []
        while True:
            question = input(f"Question #{len(questions)+1}: ").strip()
            if question.lower() == 'done':
                break
            if question:
                questions.append(question)
        
        if not questions:
            DisplayFormatter.show_warning("No questions provided")
            return
        
        DisplayFormatter.show_loading(f"Processing {len(questions)} questions...")
        result = self.qa_system.answer_multiple(questions, self.current_context)
        
        if "error" not in result:
            self.session_questions.extend(result["results"])
        
        formatted_result = DisplayFormatter.format_qa_multiple_result(result)
        print(formatted_result)
    
    def _show_session_history(self):
        """Show the history of questions asked in this session"""
        if not self.session_questions:
            DisplayFormatter.show_warning("No questions asked in this session yet")
            return
        
        print(f"\n📚 Session History ({len(self.session_questions)} questions)")
        print("=" * 60)
        
        for i, qa in enumerate(self.session_questions, 1):
            confidence_emoji = "✅" if qa["is_confident"] else "⚠️"
            print(f"\n{i}. {qa['question']}")
            print(f"   {confidence_emoji} {qa['answer']} (confidence: {qa['confidence']:.1%})")
    
    def _adjust_settings(self):
        """Allow user to adjust QA settings"""
        current_threshold = self.qa_system.confidence_threshold
        print(f"\n⚙️  Current Settings:")
        print(f"Confidence threshold: {current_threshold:.2f}")
        print("\nLower threshold = more answers accepted (less strict)")
        print("Higher threshold = fewer answers accepted (more strict)")
        
        try:
            new_threshold = input(f"Enter new threshold (0.0-1.0, current: {current_threshold}): ").strip()
            if new_threshold:
                threshold = float(new_threshold)
                self.qa_system.set_confidence_threshold(threshold)
                DisplayFormatter.show_success(f"Threshold set to {threshold:.2f}")
        except ValueError:
            DisplayFormatter.show_error("Invalid threshold value")
    
    def run(self):
        """Run interactive Question Answering"""
        self._initialize_qa_system()
        self._show_instructions()
        
        while True:
            if self.current_context:
                context_preview = (self.current_context[:50] + "...") if len(self.current_context) > 50 else self.current_context
                prompt = f"\n💬 [{context_preview}] Ask a question: "
            else:
                prompt = "\n💬 Enter command or set context first: "
            
            user_input = input(prompt).strip()
            
            if user_input.lower() == 'back':
                break
            elif user_input.lower() == 'help':
                self._show_instructions()
                continue
            elif user_input.lower() == 'context':
                self._set_context()
                continue
            elif user_input.lower() == 'multi':
                self._multi_question_mode()
                continue
            elif user_input.lower() == 'session':
                self._show_session_history()
                continue
            elif user_input.lower() == 'settings':
                self._adjust_settings()
                continue
            
            if not user_input:
                DisplayFormatter.show_warning("Please enter a question or command")
                continue
            
            # If we have a context and user input is not a command, treat it as a question
            if self.current_context:
                DisplayFormatter.show_loading("Finding answer...")
                result = self.qa_system.answer(user_input, self.current_context)
                
                if "error" not in result:
                    self.session_questions.append(result)
                
                formatted_result = DisplayFormatter.format_qa_result(result)
                print(formatted_result)
            else:
                DisplayFormatter.show_warning("Please set a context first using 'context' command")