app.py

import os
import json
import glob
from pathlib import Path
import torch
import streamlit as st
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
import numpy as np  
from sentence_transformers import util
import time

# Set device for model (CUDA if available)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load environment variables - works for both local and Hugging Face Spaces
load_dotenv()

# Set up the clinical assistant LLM
# Try to get API key from Hugging Face Spaces secrets first, then fall back to .env file
try:
    # For Hugging Face Spaces
    from huggingface_hub.inference_api import InferenceApi
    import os
    groq_api_key = os.environ.get('GROQ_API_KEY')
    
    # If not found in environment, try to get from st.secrets (Streamlit Cloud/Spaces)
    if not groq_api_key and hasattr(st, 'secrets') and 'GROQ_API_KEY' in st.secrets:
        groq_api_key = st.secrets['GROQ_API_KEY']
        
    if not groq_api_key:
        st.warning("API Key is not set in the secrets. Using a placeholder for UI demonstration.")
        # For UI demonstration without API key
        class MockLLM:
            def invoke(self, prompt):
                return {"answer": "This is a placeholder response. Please set up your GROQ_API_KEY to get real responses."}
        llm = MockLLM()
    else:    
        llm = ChatGroq(groq_api_key=groq_api_key, model_name="llama-3.3-70b-versatile")
        
except Exception as e:
    st.error(f"Error setting up LLM: {str(e)}")
    class MockLLM:
        def invoke(self, prompt):
            return {"answer": f"Error setting up LLM: {str(e)}. Please check your API key configuration."}
    llm = MockLLM()

# Set up embeddings for clinical context (Bio_ClinicalBERT)
embeddings = HuggingFaceEmbeddings(
    model_name="emilyalsentzer/Bio_ClinicalBERT",
    model_kwargs={"device": device}
)

def load_clinical_data():
    """Load both flowcharts and patient cases"""
    docs = []
    
    # Get the absolute path to the current script
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Try to handle potential errors with file loading
    try:
        # Load diagnosis flowcharts
        flowchart_dir = os.path.join(current_dir, "Diagnosis_flowchart")
        if os.path.exists(flowchart_dir):
            for fpath in glob.glob(os.path.join(flowchart_dir, "*.json")):
                try:
                    with open(fpath, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        content = f"""
                        DIAGNOSTIC FLOWCHART: {Path(fpath).stem}
                        Diagnostic Path: {data.get('diagnostic', 'N/A')}
                        Key Criteria: {data.get('knowledge', 'N/A')}
                        """
                        docs.append(Document(
                            page_content=content,
                            metadata={"source": fpath, "type": "flowchart"}
                        ))
                except Exception as e:
                    st.warning(f"Error loading flowchart file {fpath}: {str(e)}")
        else:
            st.warning(f"Flowchart directory not found at {flowchart_dir}")

        # Load patient cases
        finished_dir = os.path.join(current_dir, "Finished")
        if os.path.exists(finished_dir):
            for category_dir in glob.glob(os.path.join(finished_dir, "*")):
                if os.path.isdir(category_dir):
                    for case_file in glob.glob(os.path.join(category_dir, "*.json")):
                        try:
                            with open(case_file, 'r', encoding='utf-8') as f:
                                case_data = json.load(f)
                                notes = "\n".join(
                                    f"{k}: {v}" for k, v in case_data.items() if k.startswith("input")
                                )
                                docs.append(Document(
                                    page_content=f"""
                                    PATIENT CASE: {Path(case_file).stem}
                                    Category: {Path(category_dir).name}
                                    Notes: {notes}
                                    """,
                                    metadata={"source": case_file, "type": "patient_case"}
                                ))
                        except Exception as e:
                            st.warning(f"Error loading case file {case_file}: {str(e)}")
        else:
            st.warning(f"Finished directory not found at {finished_dir}")
            
        # If no documents were loaded, add a sample document for testing
        if not docs:
            st.warning("No clinical data files found. Using sample data for demonstration.")
            docs.append(Document(
                page_content="""SAMPLE CLINICAL DATA: This is sample data for demonstration purposes.
                This application requires clinical data files to be present in the correct directories.
                Please ensure the Diagnosis_flowchart and Finished directories exist with proper JSON files.""",
                metadata={"source": "sample", "type": "sample"}
            ))
    except Exception as e:
        st.error(f"Error loading clinical data: {str(e)}")
        # Add a fallback document
        docs.append(Document(
            page_content="Error loading clinical data. This is a fallback document for demonstration purposes.",
            metadata={"source": "error", "type": "error"}
        ))
    return docs

def build_vectorstore():
    """Build and return the vectorstore using FAISS"""
    documents = load_clinical_data()
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = splitter.split_documents(documents)
    vectorstore = FAISS.from_documents(splits, embeddings)
    return vectorstore

# Path for saving/loading the vectorstore
def get_vectorstore_path():
    """Get the path for saving/loading the vectorstore"""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(current_dir, "vectorstore")

# Initialize vectorstore with disk persistence
@st.cache_resource(show_spinner="Loading clinical knowledge base...")
def get_vectorstore():
    """Get or create the vectorstore with disk persistence"""
    vectorstore_path = get_vectorstore_path()
    
    # Try to load from disk first
    try:
        if os.path.exists(vectorstore_path):
            st.info("Loading vectorstore from disk...")
            # Set allow_dangerous_deserialization to True since we trust our own vectorstore files
            return FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True)
    except Exception as e:
        st.warning(f"Could not load vectorstore from disk: {str(e)}. Building new vectorstore.")
    
    # If loading fails or doesn't exist, build a new one
    st.info("Building new vectorstore...")
    vectorstore = build_vectorstore()
    
    # Save to disk for future use
    try:
        os.makedirs(vectorstore_path, exist_ok=True)
        vectorstore.save_local(vectorstore_path)
        st.success("Vectorstore saved to disk for future use")
    except Exception as e:
        st.warning(f"Could not save vectorstore to disk: {str(e)}")
    
    return vectorstore

def run_rag_chat(query, vectorstore):
    """Run the Retrieval-Augmented Generation (RAG) for clinical questions"""
    try:
        retriever = vectorstore.as_retriever()

        prompt_template = ChatPromptTemplate.from_template("""
        You are a clinical assistant AI. Based on the following clinical context, provide a reasoned and medically sound answer to the question.

        <context>
        {context}
        </context>

        Question: {input}

        Answer:
        """)

        retrieved_docs = retriever.invoke(query, k=3)
        retrieved_context = "\n".join([doc.page_content for doc in retrieved_docs])

        # Create document chain first
        document_chain = create_stuff_documents_chain(llm, prompt_template)
        
        # Then create retrieval chain
        chain = create_retrieval_chain(retriever, document_chain)

        # Invoke the chain
        response = chain.invoke({"input": query})
        
        # Add retrieved documents to response for transparency
        response["context"] = retrieved_docs
        
        return response
    except Exception as e:
        st.error(f"Error in RAG processing: {str(e)}")
        # Return a fallback response
        return {
            "answer": f"I encountered an error processing your query: {str(e)}",
            "context": [],
            "input": query
        }

def calculate_hit_rate(retriever, query, expected_docs, k=3):
    """Calculate the hit rate for top-k retrieved documents"""
    retrieved_docs = retriever.get_relevant_documents(query, k=k)
    retrieved_contents = [doc.page_content for doc in retrieved_docs]
    
    hits = 0
    for expected in expected_docs:
        if any(expected in retrieved for retrieved in retrieved_contents):
            hits += 1
    
    return hits / len(expected_docs) if expected_docs else 0.0

def evaluate_rag_response(response, embeddings):
    """Evaluate the RAG response for faithfulness and hit rate"""
    scores = {}

    # Faithfulness: Answer-Context Similarity
    answer_embed = embeddings.embed_query(response["answer"])
    context_embeds = [embeddings.embed_query(doc.page_content) for doc in response["context"]]
    similarities = [util.cos_sim(answer_embed, ctx_embed).item() for ctx_embed in context_embeds]
    scores["faithfulness"] = float(np.mean(similarities)) if similarities else 0.0

    # Custom Hit Rate Calculation
    retriever = response["retriever"]
    scores["hit_rate"] = calculate_hit_rate(
        retriever,
        query=response["input"],
        expected_docs=[doc.page_content for doc in response["context"]],
        k=3
    )
    
    return scores

def main():
    """Main function to run the Streamlit app"""
    # Set page configuration
    st.set_page_config(
        page_title="DiReCT - Clinical Diagnostic Assistant",
        page_icon="🩺",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Load vectorstore only once using session state
    if 'vectorstore' not in st.session_state:
        with st.spinner("Loading clinical knowledge base... This may take a minute."):
            try:
                st.session_state.vectorstore = get_vectorstore()
                # Use custom styled message without the success icon
                st.markdown("<div style='padding:10px 15px;background-color:rgba(40,167,69,0.2);border-radius:5px;border-left:5px solid rgba(40,167,69,0.8);'>Clinical knowledge base loaded successfully!</div>", unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Error loading knowledge base: {str(e)}")
                st.session_state.vectorstore = None

    # Custom CSS for modern look with dark theme compatibility
    st.markdown("""
    <style>
    .stApp {max-width: 1200px; margin: 0 auto;}
    .css-18e3th9 {padding-top: 2rem;}
    .stButton>button {background-color: #3498db; color: white;}
    .stButton>button:hover {background-color: #2980b9;}
    .chat-message {border-radius: 10px; padding: 10px; margin-bottom: 10px;}
    .chat-message-user {background-color: rgba(52, 152, 219, 0.2); color: inherit;}
    .chat-message-assistant {background-color: rgba(240, 240, 240, 0.2); color: inherit;}
    .source-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-bottom: 10px; border-left: 5px solid #3498db;}
    .metrics-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-top: 20px;}
    .features-container {display: flex; flex-wrap: wrap; gap: 20px; justify-content: center; margin-top: 30px;}
    .feature-item {flex: 1 1 calc(50% - 20px); min-width: 300px; display: flex; align-items: center; padding: 20px; border-radius: 10px; background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); transition: transform 0.3s, box-shadow 0.3s; border: 1px solid rgba(255, 255, 255, 0.1);}
    .feature-item:hover {transform: translateY(-5px); box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);}
    .feature-icon {width: 60px; height: 60px; border-radius: 50%; background: linear-gradient(135deg, #3498db, #2980b9); display: flex; align-items: center; justify-content: center; margin-right: 20px; box-shadow: 0 5px 15px rgba(52, 152, 219, 0.3);}
    .feature-icon i {font-size: 24px; color: white;}
    .feature-content {flex: 1;}
    .feature-content h3 {margin-top: 0; margin-bottom: 10px; color: inherit;}
    .feature-content p {margin: 0; font-size: 0.9em; color: inherit; opacity: 0.8;}
    .input-container {margin-bottom: 20px; padding: 15px; border-radius: 10px; background-color: rgba(255, 255, 255, 0.05); border: 1px solid rgba(255, 255, 255, 0.1);}
    </style>
    """, unsafe_allow_html=True)

    # App states
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'page' not in st.session_state:
        st.session_state.page = 'cover'

    # Sidebar
    with st.sidebar:
        st.image("https://img.icons8.com/color/96/000000/caduceus.png", width=80)
        st.title("DiReCT")
        st.markdown("### Diagnostic Reasoning for Clinical Text")
        st.markdown("---")
        
        if st.button("Home", key="home_btn"):
            st.session_state.page = 'cover'
        if st.button("Diagnostic Assistant", key="assistant_btn"):
            st.session_state.page = 'chat'
        if st.button("About", key="about_btn"):
            st.session_state.page = 'about'
            
        st.markdown("---")
        st.markdown("### Model Information")
        st.markdown("**Embedding Model:** Bio_ClinicalBERT")
        st.markdown("**LLM:** Llama-3.3-70B")
        st.markdown("**Vector Store:** FAISS")

    # Cover page
    if st.session_state.page == 'cover':
        # Hero section with animation
        col1, col2 = st.columns([2, 1])
        with col1:
            st.markdown("<h1 style='font-size:3.5em;'>DiReCT</h1>", unsafe_allow_html=True)
            st.markdown("<h2 style='font-size:1.8em;color:#3498db;'>Diagnostic Reasoning for Clinical Text</h2>", unsafe_allow_html=True)
            st.markdown("""<p style='font-size:1.2em;'>A powerful RAG-based clinical diagnostic assistant that leverages the MIMIC-IV-Ext dataset to provide accurate medical insights and diagnostic reasoning.</p>""", unsafe_allow_html=True)
            
            st.markdown("""<br>""", unsafe_allow_html=True)
            if st.button("Get Started", key="get_started"):
                st.session_state.page = 'chat'
                st.rerun()
        
        with col2:
            # Animated medical icon
            st.markdown("""
            <div style='display:flex;justify-content:center;align-items:center;height:100%;'>
                <img src="https://img.icons8.com/color/240/000000/healthcare-and-medical.png" style='max-width:90%;'>
            </div>
            """, unsafe_allow_html=True)
        
        # Modern Features section with Streamlit native components
        st.markdown("<br><br>", unsafe_allow_html=True)
        st.markdown("<h2 style='text-align:center;'>Key Features</h2>", unsafe_allow_html=True)
        
        # Create a 2x2 grid for features using Streamlit columns
        col1, col2 = st.columns(2)
        
        # Feature 1
        with col1:
            st.markdown("""
            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
                        padding: 20px; border-radius: 10px; height: 100%; 
                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
                <h3>🔍 Intelligent Retrieval</h3>
                <p>Finds the most relevant clinical information from the MIMIC-IV-Ext dataset</p>
            </div>
            """, unsafe_allow_html=True)
        
        # Feature 2
        with col2:
            st.markdown("""
            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
                        padding: 20px; border-radius: 10px; height: 100%; 
                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
                <h3>🧠 Advanced Reasoning</h3>
                <p>Applies clinical knowledge to generate accurate diagnostic insights</p>
            </div>
            """, unsafe_allow_html=True)
        
        # Feature 3
        with col1:
            st.markdown("""
            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
                        padding: 20px; border-radius: 10px; height: 100%; 
                        border: 1px solid rgba(255, 255, 255, 0.1);">
                <h3>📄 Source Transparency</h3>
                <p>Provides references to all clinical sources used in generating responses</p>
            </div>
            """, unsafe_allow_html=True)
        
        # Feature 4
        with col2:
            st.markdown("""
            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
                        padding: 20px; border-radius: 10px; height: 100%; 
                        border: 1px solid rgba(255, 255, 255, 0.1);">
                <h3>🌓 Dark/Light Theme Compatible</h3>
                <p>Optimized interface that works seamlessly in both dark and light themes</p>
            </div>
            """, unsafe_allow_html=True)

    # Chat interface
    elif st.session_state.page == 'chat':
        # Initialize session state for input if not exists
        if 'user_input' not in st.session_state:
            st.session_state.user_input = ""
        
        # Header with clear button
        col1, col2 = st.columns([3, 1])
        with col1:
            st.markdown("<h1>Clinical Diagnostic Assistant</h1>", unsafe_allow_html=True)
        with col2:
            # Add a clear button in the header
            if st.button("🗑️ Clear Chat"):
                st.session_state.chat_history = []
                st.session_state.user_input = ""
                st.rerun()
        
        st.markdown("Ask any clinical diagnostic question and get insights based on medical knowledge and patient cases.")
        
        # Fixed input area at the top
        with st.container():
            st.markdown("<div class='input-container'>", unsafe_allow_html=True)
            user_input = st.text_area("Ask a clinical question:", st.session_state.user_input, height=100, key="question_input")
            col1, col2 = st.columns([1, 5])
            with col1:
                submit_button = st.button("Submit")
            st.markdown("</div>", unsafe_allow_html=True)
        
        # Create a container for chat history
        chat_container = st.container()
        
        # Process query
        if submit_button and user_input:
            if st.session_state.vectorstore is None:
                st.error("Knowledge base not loaded. Please refresh the page and try again.")
            else:
                with st.spinner("Analyzing clinical data..."):
                    try:
                        # Add a small delay for UX
                        time.sleep(0.5)
                        
                        # Run RAG
                        response = run_rag_chat(user_input, st.session_state.vectorstore)
                        response["retriever"] = st.session_state.vectorstore.as_retriever()
                        
                        # Clear previous chat history and only keep the current response
                        st.session_state.chat_history = [(user_input, response)]
                        
                        # Clear the input field
                        st.session_state.user_input = ""
                        
                        # Rerun to update UI
                        st.rerun()
                    except Exception as e:
                        st.error(f"Error processing query: {str(e)}")
        
        # Display chat history in the container
        with chat_container:
            for i, (query, response) in enumerate(st.session_state.chat_history):
                st.markdown(f"<div class='chat-message chat-message-user'><b>🧑‍⚕️ You:</b> {query}</div>", unsafe_allow_html=True)
                
                st.markdown(f"<div class='chat-message chat-message-assistant'><b>🩺 DiReCT:</b> {response['answer']}</div>", unsafe_allow_html=True)
                
                with st.expander("View Sources"):
                    for doc in response["context"]:
                        st.markdown(f"<div class='source-box'>"
                                  f"<b>Source:</b> {Path(doc.metadata['source']).stem}<br>"
                                  f"<b>Type:</b> {doc.metadata['type']}<br>"
                                  f"<b>Content:</b> {doc.page_content[:300]}...</div>", 
                                  unsafe_allow_html=True)
                
                # Show evaluation metrics if available
                try:
                    eval_scores = evaluate_rag_response(response, embeddings)
                    with st.expander("View Evaluation Metrics"):
                        col1, col2 = st.columns(2)
                        with col1:
                            st.metric("Hit Rate (Top-3)", f"{eval_scores['hit_rate']:.2f}")
                        with col2:
                            st.metric("Faithfulness", f"{eval_scores['faithfulness']:.2f}")
                except Exception as e:
                    st.warning(f"Evaluation metrics unavailable: {str(e)}")
    
    # About page
    elif st.session_state.page == 'about':
        st.markdown("<h1>About DiReCT</h1>", unsafe_allow_html=True)
        
        st.markdown("""
        ### Project Overview
        
        DiReCT (Diagnostic Reasoning for Clinical Text) is a Retrieval-Augmented Generation (RAG) system designed to assist medical professionals with diagnostic reasoning based on clinical notes and medical knowledge.
        
        ### Data Sources
        
        This application uses the MIMIC-IV-Ext dataset, which contains de-identified clinical notes and medical records. The system processes:
        
        - Diagnostic flowcharts
        - Patient cases
        - Clinical guidelines
        
        ### Technical Implementation
        
        - **Embedding Model**: Bio_ClinicalBERT for domain-specific text understanding
        - **Vector Database**: FAISS for efficient similarity search
        - **LLM**: Llama-3.3-70B for generating medically accurate responses
        - **Framework**: Built with LangChain and Streamlit
        
        ### Evaluation Metrics
        
        The system evaluates responses using:
        
        - **Hit Rate**: Measures how many relevant documents were retrieved
        - **Faithfulness**: Measures how well the response aligns with the retrieved context
        
        ### Ethical Considerations
        
        This system is designed as a clinical decision support tool and not as a replacement for professional medical judgment. All patient data used has been properly de-identified in compliance with healthcare privacy regulations.
        """)
        
        st.markdown("<br>", unsafe_allow_html=True)
        st.markdown("### Developers")
        st.markdown("This project was developed as part of an academic assignment on RAG systems for clinical applications.")

if __name__ == "__main__":
    main()