tabs/query.py

"""
Query documents tab functionality for the Gradio app
"""
import gradio as gr

def query_documents(question, language, global_vars):
    """Handle document queries"""
    rag_system = global_vars.get('rag_system')
    vectorstore = global_vars.get('vectorstore')
    
    if not rag_system:
        return "❌ Please initialize systems first using the 'Initialize System' tab!"
    
    if not vectorstore:
        return "❌ Please upload and process documents first using the 'Upload Documents' tab!"
    
    if not question.strip():
        return "❌ Please enter a question."
    
    try:
        print(f"🔍 Processing query: {question}")
        result = rag_system.query(question, language)
        
        # Format response
        answer = result["answer"]
        sources = result.get("source_documents", [])
        model_used = result.get("model_used", "SEA-LION")
        
        # Add model information
        response = f"**Model Used:** {model_used}\n\n"
        response += f"**Answer:**\n{answer}\n\n"
        
        if sources:
            response += "**📚 Sources:**\n"
            for i, doc in enumerate(sources[:3], 1):
                metadata = doc.metadata
                source_name = metadata.get('source', 'Unknown')
                university = metadata.get('university', 'Unknown')
                country = metadata.get('country', 'Unknown')
                doc_type = metadata.get('document_type', 'Unknown')
                
                response += f"{i}. **{source_name}**\n"
                response += f"   - University: {university}\n"
                response += f"   - Country: {country}\n"
                response += f"   - Type: {doc_type}\n"
                response += f"   - Preview: {doc.page_content[:150]}...\n\n"
        else:
            response += "\n*No specific sources found. This might be a general response.*"
        
        return response
        
    except Exception as e:
        return f"❌ Error querying documents: {str(e)}\n\nPlease check the console for more details."

def get_example_questions():
    """Return example questions for the interface"""
    return [
        "What are the admission requirements for Computer Science programs in Singapore?",
        "Which universities offer scholarships for international students?",
        "What are the tuition fees for MBA programs in Thailand?",
        "Find universities with engineering programs under $5000 per year",
        "What are the application deadlines for programs in Malaysia?",
        "Compare admission requirements between different ASEAN countries"
    ]

def create_query_tab(global_vars):
    """Create the Search & Query tab"""
    with gr.Tab("🔍 Search & Query", id="query"):
        gr.Markdown("""
        ### Step 3: Ask Questions
        Ask questions about the uploaded documents in your preferred language.
        The AI will provide detailed answers with source citations.
        """)
        
        with gr.Row():
            with gr.Column(scale=3):
                question_input = gr.Textbox(
                    label="💭 Your Question",
                    placeholder="Ask anything about the universities...",
                    lines=3
                )
            
            with gr.Column(scale=1):
                language_dropdown = gr.Dropdown(
                    choices=[
                        "English", "Chinese", "Malay", "Thai", 
                        "Indonesian", "Vietnamese", "Filipino"
                    ],
                    value="English",
                    label="🌍 Response Language"
                )
        
        query_btn = gr.Button(
            "🔍 Search Documents", 
            variant="primary", 
            size="lg"
        )
        
        answer_output = gr.Textbox(
            label="🤖 AI Response",
            interactive=False,
            lines=20,
            placeholder="Ask a question to get AI-powered answers..."
        )
        
        # Example questions section
        gr.Markdown("### 💡 Example Questions")
        example_questions = get_example_questions()
        
        with gr.Row():
            for i in range(0, len(example_questions), 2):
                with gr.Column():
                    if i < len(example_questions):
                        example_btn = gr.Button(
                            example_questions[i], 
                            size="sm",
                            variant="secondary"
                        )
                        example_btn.click(
                            lambda x=example_questions[i]: x,
                            outputs=question_input
                        )
                    
                    if i + 1 < len(example_questions):
                        example_btn2 = gr.Button(
                            example_questions[i + 1], 
                            size="sm",
                            variant="secondary"
                        )
                        example_btn2.click(
                            lambda x=example_questions[i + 1]: x,
                            outputs=question_input
                        )
        
        query_btn.click(
            lambda question, language: query_documents(question, language, global_vars),
            inputs=[question_input, language_dropdown],
            outputs=answer_output
        )