app.py

#!/usr/bin/env python3
"""
VedaMD Enhanced: Sri Lankan Clinical Assistant
Main Gradio Application for Hugging Face Spaces Deployment

Enhanced Medical-Grade RAG System with:
✅ 5x Enhanced Retrieval (15+ documents vs previous 5)
✅ Medical Entity Extraction & Clinical Terminology  
✅ Clinical ModernBERT (768d medical embeddings)
✅ Medical Response Verification & Safety Protocols
✅ Advanced Re-ranking & Coverage Verification
✅ Source Traceability & Citation Support
"""

import os
import logging
import gradio as gr
from typing import List, Dict, Optional
import sys

# Add src directory to path for imports
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from src.enhanced_groq_medical_rag import EnhancedGroqMedicalRAG, EnhancedMedicalResponse

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize Enhanced Medical RAG System
logger.info("🏥 Initializing VedaMD Enhanced for Hugging Face Spaces...")
try:
    enhanced_rag_system = EnhancedGroqMedicalRAG()
    logger.info("✅ Enhanced Medical RAG system ready!")
except Exception as e:
    logger.error(f"❌ Failed to initialize system: {e}")
    raise

def process_enhanced_medical_query(message: str, history: List[List[str]]) -> str:
    """
    Process medical query with enhanced RAG system
    """
    try:
        if not message.strip():
            return "Please enter a medical question about Sri Lankan clinical guidelines."
        
        # Convert Gradio chat history to our format
        formatted_history = []
        if history:
            for chat_pair in history:
                if len(chat_pair) >= 2:
                    user_msg, assistant_msg = chat_pair[0], chat_pair[1]
                    if user_msg:
                        formatted_history.append({"role": "user", "content": user_msg})
                    if assistant_msg:
                        formatted_history.append({"role": "assistant", "content": assistant_msg})
        
        # Get enhanced response
        response: EnhancedMedicalResponse = enhanced_rag_system.query(
            query=message, 
            history=formatted_history
        )
        
        # Format enhanced response for display
        formatted_response = format_enhanced_medical_response(response)
        return formatted_response
        
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        return f"⚠️ **System Error**: {str(e)}\n\nPlease try again or contact support if the issue persists."

def format_enhanced_medical_response(response: EnhancedMedicalResponse) -> str:
    """
    Format the enhanced medical response for display, ensuring citations are always included.
    """
    formatted_parts = []
    
    # Main response from the LLM
    # The new prompt instructs the LLM to include markdown citations like [1], [2]
    # The final response text is now the primary source of the answer.
    final_response_text = response.answer
    formatted_parts.append(final_response_text)
    
    # Always add the clinical sources section if sources exist
    if response.sources:
        formatted_parts.append("\n\n---\n")
        formatted_parts.append("### 📋 **Clinical Sources**")
        # Create a numbered list of sources for clarity
        for i, source in enumerate(response.sources, 1):
            # Ensure we don't list more sources than were used for citations
            if f"[{i}]" in final_response_text:
                 formatted_parts.append(f"{i}. {source}")
            else:
                # If the LLM didn't cite this source, we can choose to omit it or list it as an uncited reference
                pass # For now, only show cited sources to keep the output clean.
    
    # Enhanced information section
    formatted_parts.append("\n\n---\n")
    formatted_parts.append("### 📊 **Enhanced Medical Analysis**")
    
    # Safety and verification info
    if response.verification_result:
        safety_emoji = "🛡️" if response.safety_status == "SAFE" else "⚠️"
        formatted_parts.append(f"**{safety_emoji} Medical Safety**: {response.safety_status}")
        formatted_parts.append(f"**🔍 Verification Score**: {response.verification_result.verification_score:.1%}")
        formatted_parts.append(f"**✅ Verified Claims**: {response.verification_result.verified_claims}/{response.verification_result.total_claims}")
    
    # Enhanced retrieval info
    formatted_parts.append(f"**🧠 Medical Entities Extracted**: {response.medical_entities_count}")
    formatted_parts.append(f"**🎯 Context Adherence**: {response.context_adherence_score:.1%}")
    formatted_parts.append(f"**📚 Sources Used**: {len(response.sources)}")
    if hasattr(response, 'query_time'): # Changed from processing_time to match the object attribute
        formatted_parts.append(f"**⚡ Processing Time**: {response.query_time:.2f}s")
    
    # Medical disclaimer
    formatted_parts.append("\n---\n")
    formatted_parts.append("*This information is for clinical reference based on Sri Lankan guidelines and does not replace professional medical judgment.*")
    
    return "\n".join(formatted_parts)

def create_enhanced_medical_interface():
    """
    Create the enhanced Gradio interface for Hugging Face Spaces
    """
    # Custom CSS for medical theme
    custom_css = """
    .gradio-container {
        max-width: 900px !important;
        margin: auto !important;
    }
    .medical-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 20px;
        border-radius: 10px;
        margin-bottom: 20px;
        text-align: center;
    }
    """
    
    with gr.Blocks(
        title="🏥 VedaMD Enhanced: Sri Lankan Clinical Assistant",
        theme=gr.themes.Soft(),
        css=custom_css
    ) as demo:
        
        # Header
        gr.HTML("""
        <div class="medical-header">
            <h1>🏥 VedaMD Enhanced: Sri Lankan Clinical Assistant</h1>
            <h3>Enhanced Medical-Grade AI with Advanced RAG & Safety Protocols</h3>
            <p>✅ 5x Enhanced Retrieval • ✅ Medical Verification • ✅ Clinical ModernBERT • ✅ Source Traceability</p>
        </div>
        """)
        
        # Description
        gr.Markdown("""
        **🩺 Advanced Medical AI Assistant** for Sri Lankan maternal health guidelines with **enhanced safety protocols**:

        🎯 **Enhanced Features:**
        - **5x Enhanced Retrieval**: 15+ documents analyzed vs previous 5
        - **Medical Entity Extraction**: Advanced clinical terminology recognition  
        - **Clinical ModernBERT**: Specialized 768d medical domain embeddings
        - **Medical Response Verification**: 100% source traceability validation
        - **Advanced Re-ranking**: Medical relevance scoring with coverage verification
        - **Safety Protocols**: Comprehensive medical claim verification before delivery

        **Ask me anything about Sri Lankan clinical guidelines with confidence!** 🇱🇰
        """)
        
        # Chat interface
        chatbot = gr.ChatInterface(
            fn=process_enhanced_medical_query,
            examples=[
                "What is the complete management protocol for severe preeclampsia in Sri Lankan guidelines?",
                "How should postpartum hemorrhage be managed according to our local clinical protocols?",
                "What medications are contraindicated during pregnancy based on Sri Lankan guidelines?",
                "What are the evidence-based recommendations for managing gestational diabetes?",
                "How should puerperal sepsis be diagnosed and treated according to our guidelines?",
                "What are the protocols for assisted vaginal delivery in complicated cases?",
                "How should intrapartum fever be managed based on Sri Lankan standards?"
            ],
            cache_examples=False
        )
        
        # Footer with technical info
        gr.Markdown("""
        ---
        **🔧 Technical Details**: Enhanced RAG with Clinical ModernBERT embeddings, medical entity extraction, 
        response verification, and multi-stage retrieval for comprehensive medical information coverage.
        
        **⚖️ Disclaimer**: This AI assistant is for clinical reference only and does not replace professional medical judgment.
        Always consult with qualified healthcare professionals for patient care decisions.
        """)
    
    return demo

# Create and launch the interface
if __name__ == "__main__":
    logger.info("🚀 Launching VedaMD Enhanced for Hugging Face Spaces...")
    
    # Create the interface
    demo = create_enhanced_medical_interface()
    
    # Launch with appropriate settings for HF Spaces
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        show_api=False
    )