streamlit_app.py

import streamlit as st
import logging
import time
from datetime import datetime

from config import Config
from model import medical_chatbot
from data_loader import medical_data_loader

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

# Page configuration
st.set_page_config(
    page_title="🏥 Medical Chatbot",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(90deg, #4CAF50, #45a049);
        padding: 1rem;
        border-radius: 10px;
        color: white;
        text-align: center;
        margin-bottom: 2rem;
    }
    .disclaimer {
        background-color: #fff3cd;
        border: 1px solid #ffecb5;
        color: #856404;
        padding: 1rem;
        border-radius: 8px;
        margin: 1rem 0;
    }
    .chat-message {
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
        border-left: 4px solid #4CAF50;
        background-color: #f8f9fa;
    }
    .user-message {
        background-color: #e3f2fd;
        border-left-color: #2196F3;
    }
    .bot-message {
        background-color: #f1f8e9;
        border-left-color: #4CAF50;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
if "chatbot_initialized" not in st.session_state:
    st.session_state.chatbot_initialized = False
if "initialization_time" not in st.session_state:
    st.session_state.initialization_time = None

def initialize_chatbot():
    """Initialize the chatbot if not already done."""
    if not st.session_state.chatbot_initialized:
        with st.spinner("🔄 Initializing Medical Chatbot... This may take a few minutes..."):
            try:
                start_time = time.time()
                medical_chatbot.create_qa_chain()
                end_time = time.time()
                
                st.session_state.chatbot_initialized = True
                st.session_state.initialization_time = end_time - start_time
                st.success(f"✅ Medical Chatbot initialized successfully in {st.session_state.initialization_time:.1f} seconds!")
                return True
            except Exception as e:
                st.error(f"❌ Error initializing chatbot: {e}")
                logger.error(f"Chatbot initialization error: {e}")
                return False
    return True

def process_query(user_input):
    """Process user query and get response."""
    if not st.session_state.chatbot_initialized:
        st.error("❌ Chatbot not initialized. Please wait for initialization to complete.")
        return None
    
    try:
        with st.spinner("🤔 Processing your question..."):
            result = medical_chatbot.query(user_input)
            return result
    except Exception as e:
        st.error(f"❌ Error processing query: {e}")
        logger.error(f"Query processing error: {e}")
        return None

def main():
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>🏥 Medical Information Chatbot</h1>
        <p>Your AI-powered healthcare information assistant</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.header("⚙️ Settings")
        
        # Model selection
        config = Config()
        available_models = list(config.LLM_PROVIDERS.keys())
        selected_model = st.selectbox(
            "🤖 Select LLM Provider",
            available_models,
            index=available_models.index(config.DEFAULT_LLM_PROVIDER)
        )
        
        # Initialize button
        if st.button("🚀 Initialize Chatbot", type="primary"):
            st.session_state.chatbot_initialized = False
            initialize_chatbot()
        
        # Status
        st.subheader("📊 Status")
        if st.session_state.chatbot_initialized:
            st.success("✅ Chatbot Ready")
            if st.session_state.initialization_time:
                st.info(f"⏱️ Init Time: {st.session_state.initialization_time:.1f}s")
        else:
            st.warning("⏳ Chatbot Not Initialized")
        
        # Statistics
        if st.session_state.messages:
            st.subheader("📈 Chat Statistics")
            total_messages = len(st.session_state.messages)
            user_messages = len([m for m in st.session_state.messages if m["role"] == "user"])
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("Total Messages", total_messages)
            with col2:
                st.metric("Q&A Pairs", user_messages)
        
        # Clear chat button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.messages = []
            st.rerun()
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Chat interface
        st.subheader("💬 Chat with Medical Assistant")
        
        # Display chat messages
        chat_container = st.container()
        with chat_container:
            for message in st.session_state.messages:
                if message["role"] == "user":
                    st.markdown(f"""
                    <div class="chat-message user-message">
                        <strong>👤 You:</strong><br>
                        {message["content"]}
                    </div>
                    """, unsafe_allow_html=True)
                else:
                    st.markdown(f"""
                    <div class="chat-message bot-message">
                        <strong>🩺 Medical Assistant:</strong><br>
                        {message["content"]}
                    </div>
                    """, unsafe_allow_html=True)
                    
                    # Show sources if available
                    if "sources" in message:
                        with st.expander(f"📚 Sources ({len(message['sources'])} documents)"):
                            for i, source in enumerate(message["sources"], 1):
                                st.write(f"**Source {i}:**")
                                st.write(source["content"])
                                if source.get("metadata"):
                                    st.caption(f"Metadata: {source['metadata']}")
                                st.divider()
        
        # Input area
        user_input = st.text_area(
            "Ask your medical question:",
            placeholder="e.g., What are the symptoms of diabetes?",
            key="user_input",
            height=100
        )
        
        col_a, col_b = st.columns([1, 4])
        with col_a:
            send_button = st.button("📤 Send", type="primary")
        
        # Process input
        if send_button and user_input.strip():
            if not st.session_state.chatbot_initialized:
                initialize_chatbot()
            
            if st.session_state.chatbot_initialized:
                # Add user message
                st.session_state.messages.append({
                    "role": "user",
                    "content": user_input,
                    "timestamp": datetime.now()
                })
                
                # Process and add bot response
                result = process_query(user_input)
                if result:
                    bot_message = {
                        "role": "assistant",
                        "content": result.get("result", "I apologize, but I couldn't generate a response."),
                        "timestamp": datetime.now()
                    }
                    
                    # Add sources if available
                    if result.get("source_documents"):
                        sources = []
                        for doc in result["source_documents"][:3]:
                            sources.append({
                                "content": doc.page_content[:300] + "..." if len(doc.page_content) > 300 else doc.page_content,
                                "metadata": doc.metadata
                            })
                        bot_message["sources"] = sources
                    
                    st.session_state.messages.append(bot_message)
                
                # Clear input and rerun
                st.rerun()
    
    with col2:
        # Information panel
        st.subheader("ℹ️ Information")
        
        # Sample questions
        st.markdown("**💡 Sample Questions:**")
        sample_questions = [
            "What are the symptoms of diabetes?",
            "How does the cardiovascular system work?",
            "What is hypertension?",
            "What are the causes of headaches?",
            "How can I maintain a healthy diet?"
        ]
        
        for question in sample_questions:
            if st.button(f"💭 {question}", key=f"sample_{hash(question)}"):
                st.session_state["user_input"] = question
                st.rerun()
        
        # Model information
        with st.expander("🤖 Model Information"):
            st.write(f"**LLM Provider:** {selected_model}")
            st.write(f"**Embedding Model:** {config.EMBEDDING_MODEL}")
            st.write(f"**Retrieval Documents:** {config.RETRIEVAL_K}")
    
    # Disclaimer
    st.markdown("""
    <div class="disclaimer">
        <strong>⚠️ Important Disclaimer:</strong><br>
        This chatbot provides general medical information for educational purposes only. 
        It is not a substitute for professional medical advice, diagnosis, or treatment. 
        Always consult qualified healthcare professionals for medical concerns.
    </div>
    """, unsafe_allow_html=True)
    
    # Footer
    st.markdown("---")
    st.markdown("🏥 **Medical Chatbot** | Powered by LangChain & Hugging Face | Built with ❤️ for Healthcare")

if __name__ == "__main__":
    main()