app.py

#!/usr/bin/env python
"""Streamlit app for the Drug Interaction System."""

import os
import sys
import streamlit as st
import matplotlib.pyplot as plt
import io
import base64
import networkx as nx
import uuid

# Add the current directory to the Python path
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))

# Import the necessary components from your package
from src.models.chatbot import DrugInteractionChatbot

# Initialize the chatbot
@st.cache_resource
def get_chatbot():
    """Get or create the chatbot instance with caching."""
    return DrugInteractionChatbot()

# Set page config
st.set_page_config(
    page_title="Drug Interaction Assistant",
    page_icon="💊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Title and description
st.title("Drug Interaction Assistant")
st.markdown("""
This application helps you analyze drug interactions, get information about medications, 
and visualize drug interaction networks. Powered by biomedical language models.
""")

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Sidebar with information
with st.sidebar:
    st.header("About")
    st.markdown("""
    This Drug Interaction Assistant can:
    - Analyze potential interactions between medications
    - Provide detailed information about specific drugs
    - Analyze clinical notes for drug mentions and interactions
    - Generate visualizations of drug interaction networks
    """)
    
    st.header("Example Questions")
    st.markdown("""
    - "Can I take aspirin and warfarin together?"
    - "Tell me about metformin"
    - "Analyze this clinical note: Patient is taking..."
    - "Show me a visualization for warfarin"
    """)

# Main content area
col1, col2 = st.columns([2, 1])

with col1:
    # Chat interface
    st.header("Chat with the Assistant")
    
    # Display chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("Ask about drug interactions..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Get chatbot response
        chatbot = get_chatbot()
        response = chatbot.process_message(prompt)
        
        # Check if we need to generate a visualization
        visualization_needed = False
        drug_name = None
        
        if "interaction found between" in response:
            # Extract drug name from response
            import re
            match = re.search(r'interaction found between (.+?) and', response)
            if match:
                drug_name = match.group(1)
                visualization_needed = True
        
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": response})
        
        # Display assistant response
        with st.chat_message("assistant"):
            st.markdown(response)
            
            # Generate and display visualization if needed
            if visualization_needed and drug_name:
                try:
                    G, error = chatbot.processor.generate_network(drug_name)
                    if not error:
                        # Create visualization
                        plt.figure(figsize=(10, 8))
                        
                        # Get positions
                        pos = nx.spring_layout(G, seed=42)
                        
                        # Draw nodes
                        node_sizes = [G.nodes[node].get('size', 10) for node in G.nodes()]
                        node_colors = [G.nodes[node].get('color', 'blue') for node in G.nodes()]
                        nx.draw_networkx_nodes(G, pos, node_size=node_sizes, node_color=node_colors, alpha=0.8)
                        
                        # Draw edges with colors based on severity
                        edge_colors = []
                        edge_widths = []
                        for u, v, data in G.edges(data=True):
                            edge_colors.append(data.get('color', 'gray'))
                            edge_widths.append(data.get('weight', 1))
                            
                        nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, alpha=0.7)
                        
                        # Add labels
                        nx.draw_networkx_labels(G, pos, font_size=10, font_family="sans-serif")
                        
                        # Save to BytesIO
                        buf = io.BytesIO()
                        plt.axis('off')
                        plt.tight_layout()
                        plt.savefig(buf, format='png', dpi=150)
                        buf.seek(0)
                        plt.close()
                        
                        # Convert to base64 for display
                        img_str = base64.b64encode(buf.read()).decode('utf-8')
                        st.image(f"data:image/png;base64,{img_str}", caption=f"Interaction Network for {drug_name}")
                except Exception as e:
                    st.error(f"Error generating visualization: {str(e)}")

with col2:
    # Drug information section
    st.header("Drug Information")
    
    # Drug search
    drug_search = st.text_input("Search for a drug", key="drug_search")
    
    if drug_search:
        chatbot = get_chatbot()
        drug_info = chatbot.processor.get_drug_information(drug_search)
        
        if drug_info:
            st.subheader(drug_info.get("drug_name", drug_search))
            
            if drug_info.get("drug_class") and drug_info["drug_class"] != "Information not available":
                st.markdown(f"**Drug Class:** {drug_info['drug_class']}")
            
            if drug_info.get("mechanism") and drug_info["mechanism"] != "Information not available":
                st.markdown(f"**Mechanism of Action:** {drug_info['mechanism']}")
            
            if drug_info.get("indications") and drug_info["indications"][0] != "Information not available":
                st.markdown("**Common Indications:**")
                for indication in drug_info["indications"]:
                    st.markdown(f"- {indication}")
            
            if drug_info.get("side_effects") and drug_info["side_effects"][0] != "Information not available":
                st.markdown("**Common Side Effects:**")
                for effect in drug_info["side_effects"]:
                    st.markdown(f"- {effect}")
            
            if drug_info.get("common_interactions") and drug_info["common_interactions"][0] != "Information not available":
                st.markdown("**Common Interactions:**")
                for interaction in drug_info["common_interactions"]:
                    st.markdown(f"- {interaction}")
            
            if drug_info.get("contraindications") and drug_info["contraindications"][0] != "Information not available":
                st.markdown("**Contraindications:**")
                for contraindication in drug_info["contraindications"]:
                    st.markdown(f"- {contraindication}")
        else:
            st.warning(f"No information found for {drug_search}")
    
    # Clinical note analysis section
    st.header("Clinical Note Analysis")
    clinical_note = st.text_area("Enter clinical note to analyze", height=150)
    
    if clinical_note and st.button("Analyze Note"):
        chatbot = get_chatbot()
        results = chatbot.processor.extract_drugs_from_clinical_notes(clinical_note)
        
        # Display medications
        if results["medications"]:
            st.subheader("Medications Identified")
            for med in results["medications"]:
                name = med.get("name", "Unknown")
                dosage = med.get("dosage", "Not specified")
                frequency = med.get("frequency", "Not specified")
                
                if dosage != "Not specified" or frequency != "Not specified":
                    st.markdown(f"- **{name}**: {dosage} {frequency}")
                else:
                    st.markdown(f"- **{name}**")
        else:
            st.info("No medications were identified in the clinical notes.")
        
        # Display potential interactions
        if results.get("potential_interactions"):
            st.subheader("Potential Interactions")
            for interaction in results["potential_interactions"]:
                drug1 = interaction.get("drug1", "Unknown")
                drug2 = interaction.get("drug2", "Unknown")
                concern = interaction.get("concern", "Potential interaction")
                
                st.markdown(f"- **{drug1}** + **{drug2}**: {concern}")
        elif results.get("database_interactions"):
            st.subheader("Potential Interactions")
            for interaction in results["database_interactions"]:
                drug1 = interaction.get("drug1", "Unknown")
                drug2 = interaction.get("drug2", "Unknown")
                desc = interaction.get("description", "Potential interaction")
                severity = interaction.get("severity", "Unknown")
                
                st.markdown(f"- **{drug1}** + **{drug2}**: {desc} ({severity})")
        else:
            st.info("No potential interactions were identified.")

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center'>
    <p>Drug Interaction Assistant | Powered by Biomedical Language Models</p>
    <p><small>This information is for educational purposes only. Always consult a healthcare professional for medical advice.</small></p>
</div>
""", unsafe_allow_html=True)