Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,602 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
 import streamlit as st
+import requests
+import pandas as pd
+from datetime import datetime
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+import plotly.express as px
+import plotly.graph_objects as go
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Configuration & Data Classes
+# ─────────────────────────────────────────────────────────────────────────────
+
+@dataclass
+class UserConfig:
+    id: int
+    name: str
+    email: str
+    age: int
+    gender: str
+    user_type: str
+    created_at: str
+    updated_at: str
+
+@dataclass
+class ChatMessage:
+    role: str
+    content: str
+    timestamp: Optional[datetime] = None
+
+class Config:
+    API_URL = os.getenv("API_URL", "https://clinical-agents-333016757590.us-central1.run.app/api/v1")
+    PAGE_TITLE = "AI Triage System"
+    PAGE_ICON = "🏥"
+    
+    # UI Constants
+    SIDEBAR_WIDTH = 300
+    CHAT_HEIGHT = 500
+    
+    # Colors
+    PRIMARY_COLOR = "#1f77b4"
+    SUCCESS_COLOR = "#2ca02c"
+    WARNING_COLOR = "#ff7f0e"
+    ERROR_COLOR = "#d62728"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Session State Manager
+# ─────────────────────────────────────────────────────────────────────────────
+
+class SessionStateManager:
+    @staticmethod
+    def init_state():
+        """Initialize session state with default values"""
+        defaults = {
+            "user_type": "patient",
+            "auth_done": False,
+            "user_id": None,
+            "user_data": None,
+            "messages": [],
+            "notes": [],
+            "chat_active": False,
+            "finished": False,
+            "current_assessment_id": None,
+            "show_help": False
+        }
+        
+        for key, val in defaults.items():
+            if key not in st.session_state:
+                st.session_state[key] = val
+    
+    @staticmethod
+    def get_user_config() -> Optional[UserConfig]:
+        """Get current user configuration"""
+        if st.session_state.auth_done and st.session_state.user_data:
+            data = st.session_state.user_data
+            return UserConfig(
+                id=data["id"],
+                name=data["name"],
+                email=data["email"],
+                age=data["age"],
+                gender=data["gender"],
+                user_type=data["user_type"],
+                created_at=data["created_at"],
+                updated_at=data["updated_at"]
+            )
+        return None
+    
+    @staticmethod
+    def reset_chat():
+        """Reset chat state"""
+        st.session_state.messages = []
+        st.session_state.chat_active = False
+        st.session_state.finished = False
+        st.session_state.current_assessment_id = None
+
+# ─────────────────────────────────────────────────────────────────────────────
+# API Service Layer
+# ─────────────────────────────────────────────────────────────────────────────
+
+class APIService:
+    @staticmethod
+    def login_user(name: str, email: str, age: int, gender: str, user_type: str) -> Tuple[bool, Dict]:
+        """Login user with the new API structure"""
+        try:
+            payload = {
+                "name": name,
+                "email": email,
+                "age": age,
+                "gender": gender,
+                "user_type": user_type
+            }
+            
+            resp = requests.post(
+                f"{Config.API_URL}/users/login",
+                json=payload,
+                timeout=10
+            )
+            resp.raise_for_status()
+            return True, resp.json()
+        except Exception as e:
+            return False, {"error": str(e)}
+    
+    @staticmethod
+    def get_user_by_id(user_id: int) -> Tuple[bool, Dict]:
+        """Get user information by ID"""
+        try:
+            resp = requests.get(f"{Config.API_URL}/users/{user_id}", timeout=10)
+            resp.raise_for_status()
+            return True, resp.json()
+        except Exception as e:
+            return False, {"error": str(e)}
+    
+    @staticmethod
+    def fetch_assessments() -> List[Dict]:
+        """Fetch all assessments from API"""
+        try:
+            resp = requests.get(f"{Config.API_URL}/assessments", timeout=10)
+            resp.raise_for_status()
+            return resp.json()
+        except Exception as e:
+            st.error(f"Failed to fetch assessments: {str(e)}")
+            return []
+    
+    @staticmethod
+    def send_chat_message(message: str, history: List[Dict], patient_id: int) -> Tuple[bool, Dict]:
+        """Send chat message to triage API"""
+        try:
+            payload = {
+                "message": message,
+                "history": history,
+                "patient_id": patient_id
+            }
+            
+            resp = requests.post(
+                f"{Config.API_URL}/triage/chat",
+                json=payload,
+                timeout=30
+            )
+            resp.raise_for_status()
+            return True, resp.json()
+        except Exception as e:
+            return False, {"error": str(e)}
+
+# ─────────────────────────────────────────────────────────────────────────────
+# UI Components
+# ─────────────────────────────────────────────────────────────────────────────
+
+class UIComponents:
+    @staticmethod
+    def render_header():
+        """Render main header with branding"""
+        st.markdown("""
+        <div style="text-align: center; padding: 1rem 0; background: linear-gradient(90deg, #1f77b4, #2ca02c); 
+                    border-radius: 10px; margin-bottom: 2rem;">
+            <h1 style="color: white; margin: 0; font-size: 2.5rem;">🏥 AI Triage System</h1>
+            <p style="color: #f0f0f0; margin: 0.5rem 0 0 0; font-size: 1.1rem;">
+                Intelligent healthcare assessment at your fingertips
+            </p>
+        </div>
+        """, unsafe_allow_html=True)
+    
+    @staticmethod
+    def render_sidebar_auth() -> bool:
+        """Render sidebar authentication section"""
+        with st.sidebar:
+            st.markdown("### 👤 Authentication")
+            
+            # User type selection
+            user_type = st.selectbox(
+                "Select your role:",
+                ["patient", "staff"],
+                key="user_type_select",
+                help="Choose whether you're a patient seeking assessment or medical staff"
+            )
+            st.session_state.user_type = user_type
+            
+            # Sign in form
+            with st.form("signin_form"):
+                st.markdown("**Sign In / Register**")
+                name = st.text_input("Full Name", placeholder="Enter your full name")
+                email = st.text_input("Email", placeholder="Enter your email address")
+                
+                # Additional fields for patients
+                if user_type == "patient":
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        age = st.number_input("Age", min_value=1, max_value=120, value=25)
+                    with col2:
+                        gender = st.selectbox("Gender", ["male", "female", "other"])
+                else:
+                    age = 30  # Default for staff
+                    gender = "not_specified"
+                
+                submitted = st.form_submit_button("🔐 Sign In", use_container_width=True)
+                
+                if submitted:
+                    if name.strip() and email.strip():
+                        with st.spinner("🔄 Signing in..."):
+                            success, response = APIService.login_user(
+                                name.strip(), 
+                                email.strip(), 
+                                age, 
+                                gender, 
+                                user_type
+                            )
+                            
+                            if success:
+                                st.session_state.user_data = response
+                                st.session_state.user_id = response["id"]
+                                st.session_state.auth_done = True
+                                st.success("✅ Successfully signed in!")
+                                st.rerun()
+                            else:
+                                st.error(f"❌ Login failed: {response.get('error', 'Unknown error')}")
+                    else:
+                        st.error("❌ Please enter both name and email")
+            
+            # Show current user info
+            if st.session_state.auth_done and st.session_state.user_data:
+                st.markdown("---")
+                st.markdown("**Current User:**")
+                user_data = st.session_state.user_data
+                st.info(f"""
+                👤 **{user_data['name']}**
+                📧 {user_data['email']}
+                🏷️ {user_data['user_type'].title()}
+                🆔 ID: {user_data['id']}
+                """)
+                
+                if st.button("🚪 Sign Out", use_container_width=True):
+                    for key in list(st.session_state.keys()):
+                        del st.session_state[key]
+                    st.rerun()
+        
+        return st.session_state.auth_done
+    
+    @staticmethod
+    def render_chat_interface(user_config: UserConfig):
+        """Render patient chat interface"""
+        st.markdown("### 💬 Chat Assessment")
+        
+        # Chat container
+        chat_container = st.container(height=Config.CHAT_HEIGHT)
+        
+        with chat_container:
+            # Display chat history
+            for i, msg in enumerate(st.session_state.messages):
+                with st.chat_message(msg["role"]):
+                    st.markdown(msg["content"])
+        
+        # Chat controls
+        col1, col2, col3 = st.columns([2, 1, 1])
+        
+        with col1:
+            if not st.session_state.chat_active:
+                if st.button("🚀 Start New Assessment", use_container_width=True, type="primary"):
+                    UIComponents._start_new_assessment(user_config)
+        
+        with col2:
+            if st.session_state.chat_active and st.button("🔄 Reset Chat", use_container_width=True):
+                SessionStateManager.reset_chat()
+                st.rerun()
+        
+        with col3:
+            if st.button("❓ Help", use_container_width=True):
+                st.session_state.show_help = not st.session_state.get("show_help", False)
+        
+        # Help section
+        if st.session_state.get("show_help", False):
+            st.markdown("""
+            ---
+            **ℹ️ How to use the AI Triage System:**
+            
+            1. **Start Assessment**: Click 'Start New Assessment' to begin
+            2. **Describe Symptoms**: Be detailed about your symptoms, when they started, and their severity
+            3. **Answer Questions**: The AI will ask follow-up questions to better understand your condition
+            4. **Get Results**: Receive your triage level and recommended next steps
+            
+            **Tips for better results:**
+            - Be honest and specific about your symptoms
+            - Include timeline information (when symptoms started)
+            - Mention any relevant medical history
+            - Don't hesitate to ask for clarification
+            """)
+        
+        # Chat input
+        if st.session_state.chat_active and not st.session_state.finished:
+            if user_input := st.chat_input("💭 Describe your symptoms or ask a question..."):
+                UIComponents._handle_user_message(user_input, user_config)
+    
+    @staticmethod
+    def _start_new_assessment(user_config: UserConfig):
+        """Start a new triage assessment"""
+        st.session_state.chat_active = True
+        st.session_state.messages = []
+        st.session_state.finished = False
+        
+        with st.spinner("🔄 Starting your assessment..."):
+            success, response = APIService.send_chat_message("", [], user_config.id)
+            
+            if success:
+                st.session_state.messages.append({
+                    "role": "assistant",
+                    "content": response["response"]
+                })
+                st.rerun()
+            else:
+                st.error(f"❌ Failed to start assessment: {response.get('error', 'Unknown error')}")
+                st.session_state.chat_active = False
+    
+    @staticmethod
+    def _handle_user_message(user_input: str, user_config: UserConfig):
+        """Handle user message in chat"""
+        # Add user message
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        with st.spinner("🤔 AI is analyzing your response..."):
+            # Convert messages to API format
+            history = [{"role": msg["role"], "content": msg["content"]} 
+                      for msg in st.session_state.messages]
+            
+            success, response = APIService.send_chat_message(user_input, history, user_config.id)
+            
+            if success:
+                st.session_state.messages.append({
+                    "role": "assistant",
+                    "content": response["response"]
+                })
+                
+                if response.get("finished", False):
+                    st.session_state.chat_active = False
+                    st.session_state.finished = True
+                    st.success("✅ Assessment completed successfully!")
+                    st.balloons()
+                
+                st.rerun()
+            else:
+                st.error(f"❌ Error: {response.get('error', 'Unknown error')}")
+
+class StaffDashboard:
+    @staticmethod
+    def render_dashboard(user_config: UserConfig):
+        """Render staff dashboard"""
+        st.markdown("### 📊 Staff Dashboard")
+        
+        # Fetch data
+        with st.spinner("📥 Loading assessment data..."):
+            assessments = APIService.fetch_assessments()
+        
+        if not assessments:
+            st.info("📭 No assessments available yet.")
+            return
+        
+        df = pd.DataFrame(assessments)
+        df["created_at"] = pd.to_datetime(df["created_at"])
+        
+        # Dashboard metrics
+        StaffDashboard._render_metrics(df)
+        
+        # Charts
+        col1, col2 = st.columns(2)
+        with col1:
+            StaffDashboard._render_esi_distribution(df)
+        with col2:
+            StaffDashboard._render_timeline_chart(df)
+        
+        # Data table
+        StaffDashboard._render_assessments_table(df)
+    
+    @staticmethod
+    def _render_metrics(df: pd.DataFrame):
+        """Render key metrics"""
+        col1, col2, col3, col4 = st.columns(4)
+        
+        with col1:
+            st.metric(
+                "📋 Total Assessments",
+                len(df),
+                delta=None
+            )
+        
+        with col2:
+            avg_esi = df['esi_level'].mean()
+            st.metric(
+                "⚡ Avg ESI Level",
+                f"{avg_esi:.1f}",
+                delta=None
+            )
+        
+        with col3:
+            emergency_cases = len(df[df['esi_level'] <= 2])
+            st.metric(
+                "🚨 Emergency Cases",
+                emergency_cases,
+                delta=f"{(emergency_cases/len(df)*100):.1f}% of total"
+            )
+        
+        with col4:
+            latest = df['created_at'].max()
+            hours_ago = (datetime.now() - latest.replace(tzinfo=None)).total_seconds() / 3600
+            st.metric(
+                "🕐 Last Assessment",
+                f"{hours_ago:.1f}h ago",
+                delta=None
+            )
+    
+    @staticmethod
+    def _render_esi_distribution(df: pd.DataFrame):
+        """Render ESI level distribution chart"""
+        st.markdown("**🎯 ESI Level Distribution**")
+        
+        esi_counts = df['esi_level'].value_counts().sort_index()
+        
+        colors = ['#d62728', '#ff7f0e', '#ffbb78', '#2ca02c', '#98df8a']
+        
+        fig = px.bar(
+            x=esi_counts.index,
+            y=esi_counts.values,
+            color=esi_counts.index,
+            color_continuous_scale='RdYlGn_r',
+            title="Distribution by ESI Level"
+        )
+        
+        fig.update_layout(
+            xaxis_title="ESI Level",
+            yaxis_title="Count",
+            showlegend=False,
+            height=300
+        )
+        
+        st.plotly_chart(fig, use_container_width=True)
+    
+    @staticmethod
+    def _render_timeline_chart(df: pd.DataFrame):
+        """Render assessments timeline"""
+        st.markdown("**📈 Assessment Timeline**")
+        
+        # Group by date
+        df_daily = df.groupby(df['created_at'].dt.date).size().reset_index()
+        df_daily.columns = ['date', 'count']
+        
+        fig = px.line(
+            df_daily,
+            x='date',
+            y='count',
+            title="Daily Assessment Volume",
+            markers=True
+        )
+        
+        fig.update_layout(
+            xaxis_title="Date",
+            yaxis_title="Number of Assessments",
+            height=300
+        )
+        
+        st.plotly_chart(fig, use_container_width=True)
+    
+    @staticmethod
+    def _render_assessments_table(df: pd.DataFrame):
+        """Render assessments data table"""
+        st.markdown("**📋 Recent Assessments**")
+        
+        # Prepare display dataframe
+        display_df = df.copy()
+        display_df['created_at'] = pd.to_datetime(display_df['created_at']).dt.strftime("%Y-%m-%d %H:%M:%S")
+        
+        # Add user information by fetching user details
+        if not display_df.empty:
+            # Create a dictionary to store user info to avoid multiple API calls
+            user_info_cache = {}
+            user_names = []
+            user_emails = []
+            
+            for user_id in display_df['user_id']:
+                if user_id not in user_info_cache:
+                    success, user_data = APIService.get_user_by_id(user_id)
+                    if success:
+                        user_info_cache[user_id] = user_data
+                    else:
+                        user_info_cache[user_id] = {"name": "Unknown", "email": "Unknown"}
+                
+                user_info = user_info_cache[user_id]
+                user_names.append(user_info.get("name", "Unknown"))
+                user_emails.append(user_info.get("email", "Unknown"))
+            
+            display_df['patient_name'] = user_names
+            display_df['patient_email'] = user_emails
+        
+        # Select columns to display
+        columns = ["id", "patient_name", "patient_email", "esi_level", "diagnosis", "notes", "created_at"]
+        
+        # Sort by creation time (newest first)
+        display_df = display_df.sort_values('created_at', ascending=False)
+        
+        st.dataframe(
+            display_df[columns],
+            use_container_width=True,
+            hide_index=True,
+            column_config={
+                "id": "Assessment ID",
+                "patient_name": "Patient Name",
+                "patient_email": "Patient Email",
+                "esi_level": st.column_config.NumberColumn(
+                    "ESI Level",
+                    help="Emergency Severity Index (1=Most urgent, 5=Least urgent)",
+                    min_value=1,
+                    max_value=5,
+                    format="%d"
+                ),
+                "diagnosis": "Diagnosis",
+                "notes": "Notes",
+                "created_at": "Created At"
+            }
+        )
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main Application
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    """Main application entry point"""
+    # Page configuration
+    st.set_page_config(
+        page_title=Config.PAGE_TITLE,
+        layout="wide",
+        page_icon=Config.PAGE_ICON,
+        initial_sidebar_state="expanded"
+    )
+    
+    # Initialize session state
+    SessionStateManager.init_state()
+    
+    # Custom CSS
+    st.markdown("""
+    <style>
+    .main > div {
+        padding-top: 2rem;
+    }
+    .stChatMessage {
+        padding: 1rem;
+        border-radius: 10px;
+        margin-bottom: 1rem;
+    }
+    .stButton > button {
+        border-radius: 20px;
+        border: none;
+        font-weight: 600;
+    }
+    .metric-container {
+        background: #f8f9fa;
+        padding: 1rem;
+        border-radius: 10px;
+        margin-bottom: 1rem;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    # Render header
+    UIComponents.render_header()
+    
+    # Handle authentication
+    if not UIComponents.render_sidebar_auth():
+        st.markdown("""
+        <div style="text-align: center; padding: 3rem; color: #666;">
+            <h3>👋 Welcome to AI Triage System</h3>
+            <p>Please sign in using the sidebar to get started with your health assessment.</p>
+        </div>
+        """, unsafe_allow_html=True)
+        st.stop()
+    
+    # Get user configuration
+    user_config = SessionStateManager.get_user_config()
+    
+    # Route to appropriate interface
+    if user_config.user_type == "patient":
+        UIComponents.render_chat_interface(user_config)
+    else:
+        StaffDashboard.render_dashboard(user_config)
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if __name__ == "__main__":
+    main()