Update app.py

app.py

@@ -1,388 +1,591 @@
-import streamlit as st
-import os
-from groq import Groq
-from datetime import datetime
-
-# Set page config FIRST
-st.set_page_config(page_title="AI Medical Consultancy", layout="wide")
-
-# Custom CSS for styling
-st.markdown("""
-    <style>
-    /* Main container styling */
-    .stApp {
-        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
-        font-family: 'Arial', sans-serif;
-    }
-
-    /* Headers styling */
-    h1, h2, h3 {
-        color: #2c3e50 !important;
-        border-bottom: 2px solid #3498db;
-        padding-bottom: 0.3em;
-    }
-
-    /* Form containers */
-    .stForm {
-        background: rgba(255, 255, 255, 0.9);
-        border-radius: 15px;
-        padding: 2rem;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-        margin: 1rem 0;
-    }
-
-    /* Input fields */
-    .stTextInput input, .stNumberInput input, .stSelectbox select, .stTextArea textarea {
-        border: 2px solid #3498db !important;
-        border-radius: 8px !important;
-        padding: 0.8rem !important;
-        transition: all 0.3s ease;
-    }
-
-    .stTextInput input:focus, .stNumberInput input:focus, 
-    .stSelectbox select:focus, .stTextArea textarea:focus {
-        border-color: #2980b9 !important;
-        box-shadow: 0 0 8px rgba(52, 152, 219, 0.3) !important;
-    }
-
-    /* Buttons styling */
-    .stButton>button {
-        background: linear-gradient(135deg, #3498db 0%, #2980b9 100%) !important;
-        color: white !important;
-        border: none !important;
-        border-radius: 8px !important;
-        padding: 0.8rem 1.5rem !important;
-        font-size: 1rem !important;
-        transition: transform 0.2s ease;
-    }
-
-    .stButton>button:hover {
-        transform: translateY(-2px);
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    }
-
-    /* Progress indicator */
-    .progress-bar {
-        display: flex;
-        justify-content: space-between;
-        margin: 2rem 0;
-        padding: 1rem;
-        background: rgba(255, 255, 255, 0.9);
-        border-radius: 10px;
-    }
-
-    .step {
-        flex: 1;
-        text-align: center;
-        padding: 0.5rem;
-        font-weight: bold;
-        color: #7f8c8d;
-    }
-
-    .step.active {
-        color: #3498db;
-        border-bottom: 3px solid #3498db;
-    }
-
-    /* Chat bubbles */
-    .dr-message {
-        background: #3498db;
-        color: white;
-        border-radius: 15px;
-        padding: 1rem;
-        margin: 1rem 0;
-        max-width: 80%;
-        width: fit-content;
-    }
-
-    /* Emergency alert */
-    .emergency-alert {
-        background: #e74c3c;
-        color: white;
-        padding: 2rem;
-        border-radius: 15px;
-        animation: pulse 1.5s infinite;
-    }
-
-    @keyframes pulse {
-        0% { transform: scale(1); }
-        50% { transform: scale(1.02); }
-        100% { transform: scale(1); }
-    }
-
-    /* Download button */
-    .download-btn {
-        background: linear-gradient(135deg, #2ecc71 0%, #27ae60 100%) !important;
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# Initialize session state variables
-if 'current_step' not in st.session_state:
-    st.session_state.current_step = 0
-if 'symptom_details' not in st.session_state:
-    st.session_state.symptom_details = []  # Initialize as an empty list
-if 'patient_info' not in st.session_state:
-    st.session_state.patient_info = {}
-
-def initialize_groq_client():
-    try:
-        # Try to get the API key from Streamlit secrets
-        api_key = None
-        try:
-            api_key = st.secrets.get("GROQ_API_KEY", os.getenv("GROQ_API_KEY"))
-        except FileNotFoundError:
-            st.warning("No `secrets.toml` file found. Please create one in the `.streamlit` folder.")
-        
-        # If not found, prompt the user to enter it
-        if not api_key:
-            api_key = st.text_input("Enter your Groq API Key:", type="password")
-            if not api_key:
-                st.warning("Please provide a valid Groq API key to proceed.")
-                return False
-        
-        # Initialize the Groq client
-        client = Groq(api_key=api_key)
-        st.session_state.client = client
-        return True
-    except Exception as e:
-        st.error(f"Error initializing Groq client: {str(e)}")
-        return False
-
-def symptom_interrogation_step():
-    client = st.session_state.client
-    main_symptom = st.session_state.patient_info['main_symptom']
-    step = len(st.session_state.symptom_details)  # Use the number of collected details as the step
-
-    if step == 0:
-        # First question: ask about the main symptom
-        medical_focus = {
-            'pain': "location/radiation/provoking factors",
-            'fever': "pattern/associated symptoms/response to meds",
-            'gi': "bowel changes/ingestion timing/associated symptoms",
-            'respiratory': "exertion relationship/sputum/triggers"
-        }
-        focus = medical_focus.get(main_symptom.lower(), 
-                "temporal pattern/severity progression/associated symptoms")
-
-        prompt = f"""As an ER physician, ask ONE high-yield question about {main_symptom}
-        focusing on {focus} to differentiate serious causes. Your task is to have a polite and simple conversation with a patient. 
-        Start by asking ONE specific follow-up question about their initial symptom: {main_symptom}.
-        Ask only one question at a time to avoid overwhelming the patient.
-        Keep your language clear, professional, and easy to understand."""
-        
-        messages = [
-            {"role": "system", "content": "Ask focused clinical questions. One at a time."},
-            {"role": "user", "content": prompt}
-        ]
-    else:
-        # Subsequent questions: use the last Q&A to generate the next question
-        last_qa = st.session_state.symptom_details[-1]
-        prompt = f"""Last Q&A: {last_qa['question']} → {last_qa['answer']}
-        Based on this, ask the NEXT most critical question to differentiate between
-        possible causes of {main_symptom}. Consider red flags and likelihood."""
-        messages = [{"role": "user", "content": prompt}]
-
-    try:
-        response = client.chat.completions.create(
-            messages=messages,
-            model="mixtral-8x7b-32768",
-            temperature=0.3
-        )
-        question = response.choices[0].message.content.strip()
-        if not question.endswith('?'):
-            question += '?'
-        st.session_state.current_question = question
-    except Exception as e:
-        st.error(f"Error generating question: {str(e)}")
-        st.stop()
-
-def handle_symptom_interrogation():
-    st.header("Symptom Analysis")
-    
-    if st.session_state.current_step == 1:
-        symptom_interrogation_step()
-        st.session_state.current_step = 2
-    
-    if 'current_question' in st.session_state:
-        with st.form("symptom_qna"):
-            st.markdown(f'<div class="dr-message">👨‍⚕️ {st.session_state.current_question}</div>', unsafe_allow_html=True)
-            answer = st.text_input("Your answer:", key=f"answer_{len(st.session_state.symptom_details)}")
-            
-            if st.form_submit_button("Next"):
-                if answer:
-                    st.session_state.symptom_details.append({
-                        "question": st.session_state.current_question,
-                        "answer": answer
-                    })
-                    del st.session_state.current_question
-                    
-                    # Check for emergency after 3 questions
-                    if len(st.session_state.symptom_details) >= 3:
-                        last_answer = st.session_state.symptom_details[-1]['answer']
-                        try:
-                            urgency_check = st.session_state.client.chat.completions.create(
-                                messages=[{"role": "user", "content": 
-                                        f"Does '{last_answer}' indicate immediate emergency? Yes/No"}],
-                                model="mixtral-8x7b-32768",
-                                temperature=0
-                            ).choices[0].message.content
-                            
-                            if 'YES' in urgency_check.upper():
-                                st.markdown('<div class="emergency-alert">🚨 Emergency detected! Please seek immediate medical attention.</div>', unsafe_allow_html=True)
-                                st.session_state.current_step = 4
-                                return
-                        except Exception as e:
-                            st.error(f"Error checking urgency: {str(e)}")
-                    
-                    if len(st.session_state.symptom_details) < 7:
-                        st.session_state.current_step = 1
-                        st.rerun()
-                    else:
-                        st.session_state.current_step = 3
-                        st.rerun()
-                else:
-                    st.warning("Please provide an answer")
-
-def collect_basic_info():
-    st.header("Patient Information")
-    with st.form("basic_info"):
-        st.session_state.patient_info['name'] = st.text_input("Full Name")
-        st.session_state.patient_info['age'] = st.number_input("Age", min_value=0, max_value=120)
-        st.session_state.patient_info['gender'] = st.selectbox("Gender", ["Male", "Female", "Other"])
-        st.session_state.patient_info['main_symptom'] = st.text_input("Main Symptom")
-        
-        if st.form_submit_button("Next"):
-            if all([st.session_state.patient_info.get(k) for k in ['name', 'age', 'gender', 'main_symptom']]):
-                st.session_state.current_step = 1
-                st.rerun()
-            else:
-                st.warning("Please fill all required fields")
-
-def collect_medical_history():
-    st.header("Medical History")
-    with st.form("medical_history"):
-        st.session_state.patient_info['medical_history'] = st.text_area("Relevant Medical History")
-        st.session_state.patient_info['medications'] = st.text_area("Current Medications")
-        st.session_state.patient_info['allergies'] = st.text_input("Known Allergies")
-        st.session_state.patient_info['last_meal'] = st.text_input("Last Meal Time")
-        st.session_state.patient_info['recent_travel'] = st.text_input("Recent Travel History")
-        
-        if st.form_submit_button("Submit"):
-            st.session_state.current_step = 4
-            st.rerun()
-
-def generate_risk_assessment():
-    st.header("Risk Assessment")
-    
-    try:
-        symptom_log = "\n".join(
-            [f"Q: {q['question']}\nA: {q['answer']}"
-            for q in st.session_state.symptom_details]
-        )
-
-        patient_profile = f"""
-**Patient Profile**
-Name: {st.session_state.patient_info['name']}
-Age: {st.session_state.patient_info['age']}
-Gender: {st.session_state.patient_info['gender']}
-
-**Primary Complaint**
-{st.session_state.patient_info['main_symptom']}
-
-**Symptom Interrogation**
-{symptom_log}
-
-**Medical History**
-{st.session_state.patient_info.get('medical_history', 'None reported')}
-
-**Current Medications**
-{st.session_state.patient_info.get('medications', 'None')}
-
-**Allergies**
-{st.session_state.patient_info.get('allergies', 'None reported')}
-
-**Recent Context**
-Last Meal: {st.session_state.patient_info.get('last_meal', 'Unknown')}
-Recent Travel: {st.session_state.patient_info.get('recent_travel', 'None')}
-        """
-
-        analysis_prompt = f"""STRICTLY follow these instructions:
-1. Analyze this case: {patient_profile}
-2. Output ONLY this exact format WITHOUT ANY additional text:
-[Age]-year-old [gender] with [symptom details]
-
-Example Output:
-"45-year-old man with severe chest pain radiating to the jaw"
-
-Your Output MUST BE:"""
-
-        response = st.session_state.client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are a medical AI that outputs ONLY patient descriptions."},
-                {"role": "user", "content": analysis_prompt}
-            ],
-            model="mixtral-8x7b-32768",
-            temperature=0.3,
-            max_tokens=100
-        )
-
-        risk_prompt = response.choices[0].message.content.strip('"')
-        
-        st.subheader("Clinical Summary")
-        st.markdown(f"```\n{risk_prompt}\n```")
-        
-        # Create download button
-        timestamp = datetime.now().strftime('%Y%m%d%H%M')
-        filename = f"{st.session_state.patient_info['name'].replace(' ', '_')}_assessment_{timestamp}.txt"
-        st.download_button(
-            label="Download Assessment",
-            data=risk_prompt,
-            file_name=filename,
-            mime="text/plain"
-        )
-        
-    except Exception as e:
-        st.error(f"Error generating risk assessment: {str(e)}")
-
-def main():
-    st.title("🏥 AI Medical Consultancy")
-    
-    # Progress indicator
-    steps_titles = ["Patient Info", "Symptoms", "Medical History", "Assessment"]
-    progress_html = """
-    <div class="progress-bar">
-        <div class="step {}">{}</div>
-        <div class="step {}">{}</div>
-        <div class="step {}">{}</div>
-        <div class="step {}">{}</div>
-    </div>
-    """.format(
-        'active' if st.session_state.current_step >= 0 else '',
-        '1. Patient Info',
-        'active' if st.session_state.current_step >= 1 else '',
-        '2. Symptoms',
-        'active' if st.session_state.current_step >= 3 else '',
-        '3. History',
-        'active' if st.session_state.current_step >= 4 else '',
-        '4. Report'
-    )
-    st.markdown(progress_html, unsafe_allow_html=True)
-    
-    if not initialize_groq_client():
-        return
-    
-    steps = {
-        0: collect_basic_info,
-        1: handle_symptom_interrogation,
-        2: handle_symptom_interrogation,
-        3: collect_medical_history,
-        4: generate_risk_assessment
-    }
-    
-    current_step = st.session_state.get('current_step', 0)
-    if current_step in steps:
-        steps[current_step]()
-
-if __name__ == "__main__":
+import streamlit as st
+import os
+from groq import Groq
+from datetime import datetime
+from transformers import pipeline
+import pandas as pd
+import re
+
+# Set page config FIRST
+st.set_page_config(page_title="AI Medical Consultancy", layout="wide")
+
+# Custom CSS for styling
+st.markdown("""
+    <style>
+    /* Color Variables */
+    :root {
+        --primary: #3498db;       /* Blue */
+        --secondary: #2c3e50;     /* Dark accent */
+        --accent: #f1c40f;        /* Yellow */
+        --success: #2ecc71;       /* Positive actions */
+        --light: #ffffff;         /* White backgrounds */
+        --dark: #000000;          /* Black text/elements */
+    }
+
+    /* Main container styling */
+    .stApp {
+        background: linear-gradient(135deg, #3498db 0%, #e0e0e0 100%);
+        font-family: 'Arial', sans-serif;
+    }
+
+    /* Headers styling */
+    h1, h2, h3 {
+        color: var(--dark) !important;
+        border-bottom: 3px solid var(--primary);
+        padding-bottom: 0.3em;
+    }
+
+    /* Form containers */
+    .stForm {
+        background: #000000;
+        border: 1px solid rgba(44, 62, 80, 0.2);
+        border-radius: 15px;
+        padding: 2rem;
+        box-shadow: 0 8px 30px rgba(0, 0, 0, 0.12);
+        margin: 1rem 0;
+    }
+
+    /* Input fields */
+    .stTextInput input, .stNumberInput input, 
+    .stSelectbox select, .stTextArea textarea {
+        border: 2px solid #00FFFF !important;
+        border-radius: 10px !important;
+        padding: 1rem !important;
+        background: #00FFFF !important;
+        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
+        color: var(--dark) !important;
+    }
+
+    .stTextInput input:focus, .stNumberInput input:focus, 
+    .stSelectbox select:focus, .stTextArea textarea:focus {
+        border-color: var(--primary) !important;
+        box-shadow: 0 0 12px rgba(52, 152, 219, 0.2) !important;
+        background: white !important;
+        color: var(--dark) !important;
+    }
+
+    /* Buttons styling */
+    .stButton>button {
+        background: linear-gradient(135deg, var(--primary) 0%, var(--accent) 100%) !important;
+        color: var(--dark) !important;
+        border: none !important;
+        border-radius: 10px !important;
+        padding: 1rem 2rem !important;
+        font-size: 1rem !important;
+        transition: all 0.3s ease;
+        position: relative;
+        overflow: hidden;
+    }
+
+    .stButton>button:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 8px 15px rgba(52, 152, 219, 0.3);
+        opacity: 0.95;
+    }
+
+    .stButton>button:active {
+        transform: translateY(0);
+        opacity: 1;
+    }
+
+    /* Progress indicator */
+    .progress-bar {
+        display: flex;
+        justify-content: space-between;
+        margin: 2rem 0;
+        padding: 1rem;
+        background: rgba(255, 255, 255, 0.9);
+        border-radius: 10px;
+        color: var(--dark) !important;
+    }
+
+    .step {
+        flex: 1;
+        text-align: center;
+        padding: 1rem;
+        font-weight: 600;
+        color: #95a5a6;
+        position: relative;
+    }
+
+    .step.active {
+        color: var(--primary);
+    }
+
+    .step.active:after {
+        content: '';
+        position: absolute;
+        bottom: -1px;
+        left: 50%;
+        transform: translateX(-50%);
+        width: 40%;
+        height: 3px;
+        background: var(--primary);
+    }
+
+    /* Chat bubbles */
+    .dr-message {
+        background: linear-gradient(135deg, var(--primary) 0%, #2980b9 100%);
+        color: white;
+        border-radius: 20px 20px 20px 4px;
+        padding: 1.2rem 1.5rem;
+        margin: 1rem 0;
+        max-width: 80%;
+        width: fit-content;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
+    }
+
+    .user-message {
+        background: linear-gradient(135deg, #f1c40f 0%, #e1b800 100%);
+        margin-left: auto;
+        border-radius: 20px 20px 4px 20px;
+        color: var(--dark) !important;
+    }
+
+    /* Emergency alert */
+    .emergency-alert {
+        background: linear-gradient(135deg, var(--accent) 0%, #c0392b 100%);
+        color: white;
+        padding: 2rem;
+        border-radius: 15px;
+        animation: pulse 1.5s infinite;
+        text-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
+    }
+
+    @keyframes pulse {
+        0% { transform: scale(1); }
+        50% { transform: scale(1.02); }
+        100% { transform: scale(1); }
+    }
+
+    /* Download button */
+    .download-btn {
+        background: linear-gradient(135deg, var(--success) 0%, #27ae60 100%) !important;
+    }
+
+    /* Enhanced Data Visualization Contrast */
+    .stDataFrame {
+        border: 1px solid rgba(0, 0, 0, 0.1);
+        border-radius: 12px;
+        overflow: hidden;
+        background: #f0f0f0;
+        color: var(--dark) !important;
+    }
+
+    /* Tabbed Interface Styling */
+    .stTabs [role="tablist"] {
+        gap: 10px;
+        padding: 8px;
+        background: rgba(240, 240, 240, 0.9);
+        border-radius: 12px;
+        color: var(--dark) !important;
+    }
+
+    .stTabs [role="tab"] {
+        background: #ffffff !important;
+        border-radius: 8px !important;
+        transition: all 0.3s ease;
+        color: var(--dark) !important;
+    }
+
+    .stTabs [role="tab"][aria-selected="true"] {
+        background: var(--primary) !important;
+        color: white !important;
+        transform: scale(1.05);
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+
+# Initialize session state variables
+if 'current_step' not in st.session_state:
+    st.session_state.current_step = 0
+if 'symptom_details' not in st.session_state:
+    st.session_state.symptom_details = []  # Initialize as an empty list
+if 'patient_info' not in st.session_state:
+    st.session_state.patient_info = {}
+
+# Initialize models and dataset with caching
+@st.cache_resource
+def load_ner_model():
+    return pipeline("token-classification", 
+                    model="d4data/biomedical-ner-all",
+                    aggregation_strategy='first')
+
+@st.cache_resource
+def load_dataset():
+    try:
+        data = pd.read_csv('DATASET.csv', encoding='latin1')
+        data['Symptom_lower'] = data['Symptom'].str.strip().str.lower()
+        return data
+    except FileNotFoundError:
+        st.error("DATASET.csv not found. Please ensure the file is in the correct directory.")
+        return None
+
+pipe = load_ner_model()
+data = load_dataset()
+if data is None:
+    st.stop()
+
+SEVERITY_KEYWORDS = {
+    'normal': ['mild', 'occasional', 'controlled', 'temporary', 'fleeting'],
+    'moderate': ['moderate', 'persistent', 'frequent', 'prolonged', 'bloating'],
+    'severe': ['severe', 'extreme', 'crushing', 'radiating', 'blood', 'inability', 
+              'sweating', 'vomiting', 'fever', 'swelling', 'radiates', 'persistent >3 days']
+}
+
+def merge_entities(entities, text):
+    """Merge entities using dataset symptom names for multi-word matching"""
+    merged = []
+    detected = set()
+    text_lower = text.lower()
+    
+    symptoms_sorted = data['Symptom'].str.lower().sort_values(
+        key=lambda x: x.str.len(), ascending=False).tolist()
+    
+    for symptom in symptoms_sorted:
+        if symptom in text_lower and symptom not in detected:
+            start = text_lower.find(symptom)
+            end = start + len(symptom)
+            merged.append({
+                'word': text[start:end],
+                'start': start,
+                'end': end,
+                'entity_group': 'Sign_symptom'
+            })
+            detected.add(symptom)
+    return merged
+
+def determine_condition(text, entity):
+    context_window = re.findall(r'\w+', text[max(0, entity['start']-20):entity['end']+20])
+    context = ' '.join(context_window).lower()
+    
+    for level, keywords in SEVERITY_KEYWORDS.items():
+        if any(k in context for k in keywords):
+            return level.capitalize()
+    return 'Normal'
+
+def get_risk_score(symptom, condition):
+    symptom_clean = symptom.strip().lower()
+    matches = data[
+        data['Symptom_lower'].str.contains(rf'\b{symptom_clean}\b', regex=True, flags=re.IGNORECASE)
+    ]
+    if not matches.empty:
+        matches = matches[matches['Condition'] == condition]
+        return matches['Risk Score'].values[0] if not matches.empty else 0
+    return 0
+
+def calculate_risk_score(text):
+    entities = pipe(text)
+    merged_entities = merge_entities(entities, text)
+
+    score = 0
+    breakdown = []
+    for ent in merged_entities:
+        if ent['entity_group'] == "Sign_symptom":
+            symptom = ent['word'].strip().lower()
+            condition = determine_condition(text, ent)
+            risk = get_risk_score(symptom, condition)
+            score += risk
+            breakdown.append({
+                'symptom': symptom.capitalize(),
+                'condition': condition,
+                'risk': risk
+            })
+    return score, breakdown
+
+def initialize_groq_client():
+    try:
+        # Try to get the API key from Streamlit secrets
+        api_key = None
+        try:
+            api_key = st.secrets.get("GROQ_API_KEY", os.getenv("GROQ_API_KEY"))
+        except FileNotFoundError:
+            st.warning("No secrets.toml file found. Please create one in the .streamlit folder.")
+        
+        # If not found, prompt the user to enter it
+        if not api_key:
+            api_key = st.text_input("Enter your Groq API Key:", type="password")
+            if not api_key:
+                st.warning("Please provide a valid Groq API key to proceed.")
+                return False
+        
+        # Initialize the Groq client
+        client = Groq(api_key=api_key)
+        st.session_state.client = client
+        return True
+    except Exception as e:
+        st.error(f"Error initializing Groq client: {str(e)}")
+        return False
+
+def symptom_interrogation_step():
+    client = st.session_state.client
+    main_symptom = st.session_state.patient_info['main_symptom']
+    step = len(st.session_state.symptom_details)  # Use the number of collected details as the step
+
+    if step == 0:
+        # First question: ask about the main symptom
+        medical_focus = {
+            'pain': "location/radiation/provoking factors",
+            'fever': "pattern/associated symptoms/response to meds",
+            'gi': "bowel changes/ingestion timing/associated symptoms",
+            'respiratory': "exertion relationship/sputum/triggers"
+        }
+        focus = medical_focus.get(main_symptom.lower(), 
+                "temporal pattern/severity progression/associated symptoms")
+
+        prompt = f"""As an ER physician, ask ONE high-yield question about {main_symptom}
+        focusing on {focus} to differentiate serious causes. Your task is to have a polite and simple conversation with a patient. 
+        Start by asking ONE specific follow-up question about their initial symptom: {main_symptom}.
+        Ask only one question at a time to avoid overwhelming the patient.
+        Keep your language clear, professional, and easy to understand. 
+        
+        Dont display possibe symptoms or why you are asking questions."""
+        
+        messages = [
+            {"role": "system", "content": "Ask focused clinical questions. One at a time."},
+            {"role": "user", "content": prompt}
+        ]
+    else:
+        # Subsequent questions: use the last Q&A to generate the next question
+        last_qa = st.session_state.symptom_details[-1]
+        prompt = f"""Last Q&A: {last_qa['question']} → {last_qa['answer']}
+        Based on this, ask the NEXT most critical question to differentiate between
+        possible causes of {main_symptom}. Consider red flags and likelihood."""
+        messages = [{"role": "user", "content": prompt}]
+
+    try:
+        response = client.chat.completions.create(
+            messages=messages,
+            model="mixtral-8x7b-32768",
+            temperature=0.3
+        )
+        question = response.choices[0].message.content.strip()
+        if not question.endswith('?'):
+            question += '?'
+        st.session_state.current_question = question
+    except Exception as e:
+        st.error(f"Error generating question: {str(e)}")
+        st.stop()
+
+def handle_symptom_interrogation():
+    st.header("Symptom Analysis")
+    
+    if st.session_state.current_step == 1:
+        symptom_interrogation_step()
+        st.session_state.current_step = 2
+    
+    if 'current_question' in st.session_state:
+        with st.form("symptom_qna"):
+            st.markdown(f'<div class="dr-message">👨‍⚕ {st.session_state.current_question}</div>', unsafe_allow_html=True)
+            answer = st.text_input("Your answer:", key=f"answer_{len(st.session_state.symptom_details)}")
+            
+            if st.form_submit_button("Next"):
+                if answer:
+                    st.session_state.symptom_details.append({
+                        "question": st.session_state.current_question,
+                        "answer": answer
+                    })
+                    del st.session_state.current_question
+                    
+                    # Check for emergency after 3 questions
+                    if len(st.session_state.symptom_details) >= 3:
+                        last_answer = st.session_state.symptom_details[-1]['answer']
+                        try:
+                            urgency_check = st.session_state.client.chat.completions.create(
+                                messages=[{"role": "user", "content": 
+                                        f"Does '{last_answer}' indicate immediate emergency? Yes/No"}],
+                                model="mixtral-8x7b-32768",
+                                temperature=0
+                            ).choices[0].message.content
+                            
+                            if 'YES' in urgency_check.upper():
+                                st.markdown('<div class="emergency-alert">🚨 Emergency detected! Please seek immediate medical attention.</div>', unsafe_allow_html=True)
+                                st.session_state.current_step = 4
+                                return
+                        except Exception as e:
+                            st.error(f"Error checking urgency: {str(e)}")
+                    
+                    if len(st.session_state.symptom_details) < 7:
+                        st.session_state.current_step = 1
+                        st.rerun()
+                    else:
+                        st.session_state.current_step = 3
+                        st.rerun()
+                else:
+                    st.warning("Please provide an answer")
+
+def collect_basic_info():
+    st.header("Patient Information")
+    with st.form("basic_info"):
+        st.session_state.patient_info['name'] = st.text_input("Full Name")
+        st.session_state.patient_info['age'] = st.number_input("Age", min_value=0, max_value=120)
+        st.session_state.patient_info['gender'] = st.selectbox("Gender", ["Male", "Female", "Other"])
+        st.session_state.patient_info['main_symptom'] = st.text_input("Main Symptom")
+        
+        if st.form_submit_button("Next"):
+            if all([st.session_state.patient_info.get(k) for k in ['name', 'age', 'gender', 'main_symptom']]):
+                st.session_state.current_step = 1
+                st.rerun()
+            else:
+                st.warning("Please fill all required fields")
+
+def collect_medical_history():
+    st.header("Medical History")
+    with st.form("medical_history"):
+        st.session_state.patient_info['medical_history'] = st.text_area("Relevant Medical History")
+        st.session_state.patient_info['medications'] = st.text_area("Current Medications")
+        st.session_state.patient_info['allergies'] = st.text_input("Known Allergies")
+        st.session_state.patient_info['last_meal'] = st.text_input("Last Meal Time")
+        st.session_state.patient_info['recent_travel'] = st.text_input("Recent Travel History")
+        
+        if st.form_submit_button("Submit"):
+            st.session_state.current_step = 4
+            st.rerun()
+
+def generate_risk_assessment():
+    st.header("Risk Assessment")
+    
+    try:
+        symptom_log = "\n".join(
+            [f"Q: {q['question']}\nA: {q['answer']}"
+            for q in st.session_state.symptom_details]
+        )
+
+        patient_profile = f"""
+*Patient Profile*
+Name: {st.session_state.patient_info['name']}
+Age: {st.session_state.patient_info['age']}
+Gender: {st.session_state.patient_info['gender']}
+
+*Primary Complaint*
+{st.session_state.patient_info['main_symptom']}
+
+*Symptom Interrogation*
+{symptom_log}
+
+*Medical History*
+{st.session_state.patient_info.get('medical_history', 'None reported')}
+
+*Current Medications*
+{st.session_state.patient_info.get('medications', 'None')}
+
+*Allergies*
+{st.session_state.patient_info.get('allergies', 'None reported')}
+
+*Recent Context*
+Last Meal: {st.session_state.patient_info.get('last_meal', 'Unknown')}
+Recent Travel: {st.session_state.patient_info.get('recent_travel', 'None')}
+        """
+
+        analysis_prompt = f"""STRICTLY follow these instructions:
+1. Analyze this case: {patient_profile}
+2. *Include ONLY symptoms the patient is actively experiencing*. Exclude all negated symptoms (e.g., "no fever," "denies breathlessness").
+3. Output *EXCLUSIVELY* in this format with NO additional text or explanations:
+[Age]-year-old [gender] with [specific, present symptoms].
+
+Example Output:
+"45-year-old man with severe chest pain radiating to the jaw"
+
+Your Output:"""
+
+        response = st.session_state.client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a medical AI that outputs ONLY patient descriptions."},
+                {"role": "user", "content": analysis_prompt}
+            ],
+            model="mixtral-8x7b-32768",
+            temperature=0.3,
+            max_tokens=100
+        )
+
+        risk_prompt = response.choices[0].message.content.strip('"')
+        
+        st.subheader("Clinical Summary")
+        st.markdown(f"\n{risk_prompt}\n")
+        
+        # Add risk scoring
+        risk_score, breakdown = calculate_risk_score(risk_prompt)
+        
+        # Display risk score with color coding
+        score_color = "#2ecc71" if risk_score < 40 else "#f1c40f" if risk_score < 70 else "#e74c3c"
+        st.markdown(f"""
+            <div style="background: {score_color}; 
+                        color: white; 
+                        padding: 1.5rem; 
+                        border-radius: 10px;
+                        text-align: center;
+                        margin: 2rem 0;">
+                <h2 style="color: white; margin: 0;">Risk Assessment Score</h2>
+                <h1 style="font-size: 3.5rem; margin: 0.5rem 0;">{risk_score}/100</h1>
+                <p>{'Low Risk' if risk_score < 40 else 'Moderate Risk' if risk_score < 70 else 'High Risk'}</p>
+            </div>
+        """, unsafe_allow_html=True)
+
+        # Display breakdown
+        with st.expander("View Detailed Risk Breakdown", expanded=True):
+            for item in breakdown:
+                st.markdown(f"""
+                    <div style="background: rgba(236, 240, 241, 0.5);
+                                padding: 1rem;
+                                border-radius: 8px;
+                                margin: 0.5rem 0;">
+                        <strong>{item['symptom']}</strong> ({item['condition']})
+                        <div style="float: right;">+{item['risk']} pts</div>
+                    </div>
+                """, unsafe_allow_html=True)
+
+        # Create download button
+        timestamp = datetime.now().strftime('%Y%m%d%H%M')
+        filename = f"{st.session_state.patient_info['name'].replace(' ', '')}_assessment{timestamp}.txt"
+        st.download_button(
+            label="Download Assessment",
+            data=risk_prompt,
+            file_name=filename,
+            mime="text/plain"
+        )
+        
+    except Exception as e:
+        st.error(f"Error generating risk assessment: {str(e)}")
+
+def main():
+    st.title("🏥 AI Medical Consultancy")
+    
+    # Progress indicator
+    steps_titles = ["Patient Info", "Symptoms", "Medical History", "Assessment"]
+    progress_html = """
+    <div class="progress-bar">
+        <div class="step {}">{}</div>
+        <div class="step {}">{}</div>
+        <div class="step {}">{}</div>
+        <div class="step {}">{}</div>
+    </div>
+    """.format(
+        'active' if st.session_state.current_step >= 0 else '',
+        '1. Patient Info',
+        'active' if st.session_state.current_step >= 1 else '',
+        '2. Symptoms',
+        'active' if st.session_state.current_step >= 3 else '',
+        '3. History',
+        'active' if st.session_state.current_step >= 4 else '',
+        '4. Report'
+    )
+    st.markdown(progress_html, unsafe_allow_html=True)
+    
+    if not initialize_groq_client():
+        return
+    
+    steps = {
+        0: collect_basic_info,
+        1: handle_symptom_interrogation,
+        2: handle_symptom_interrogation,
+        3: collect_medical_history,
+        4: generate_risk_assessment
+    }
+    
+    current_step = st.session_state.get('current_step', 0)
+    if current_step in steps:
+        steps[current_step]()
+
+if _name_ == "_main_":
     main()