app.py

import streamlit as st
import groq
from dotenv import load_dotenv
import os

load_dotenv()

# Initialize the Groq client
client = groq.Groq()

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Available models
MODELS = [
    "mixtral-8x7b-32768",
    "gemma2-9b-it",
    "llama-3.2-1b-preview",
    
]

# Default system prompt
DEFAULT_SYSTEM_PROMPT = """You are an expert physiotherapist dedicated to helping users improve their well-being. 
Using the user's provided data,  
Send a brief, 2-line message to [patient name], personalized message that checks in on their condition or asks
 how they’ve been feeling lately, even though they haven’t taken a plan yet.
The message should be warm, empathetic, and tailored to their persona, 
city or country, and specific needs. Focus on being kind, supportive, and motivating, 
while offering a helpful tip or insight that shows your genuine care for their health. 
Build trust by emphasizing that you’re here to help them at their pace, without any pressure, so they feel confident and reassured in reaching out whenever they’re ready."""

# Sidebar for system prompt
st.sidebar.title("System Prompt")
system_prompt = st.sidebar.text_area("Edit the system prompt here:", value=DEFAULT_SYSTEM_PROMPT, height=300)

# Add a button to apply the system prompt
if st.sidebar.button("Apply System Prompt"):
    st.session_state.system_prompt = system_prompt
    # Clear the chat history when a new prompt is applied
    st.session_state.messages = [{"role": "system", "content": system_prompt}]
    st.sidebar.success("System prompt applied successfully! Chat history cleared.")

# Initialize system_prompt in session state if it doesn't exist
if "system_prompt" not in st.session_state:
    st.session_state.system_prompt = DEFAULT_SYSTEM_PROMPT

# Initialize chat history if it doesn't exist
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": st.session_state.system_prompt}]

# Streamlit app
st.title("Health Genie App")

# Model selection
selected_model = st.selectbox("Select a model", MODELS)

# Add a Clear Conversation button
if st.button("Clear Conversation"):
    st.session_state.messages = [{"role": "system", "content": st.session_state.system_prompt}]
    st.success("Conversation cleared!")

# Predefined patient profiles
patient_profiles = {
    "Patient 1": {
        "name": "John Doe",
        "age": 35, "gender": "Male", "height": 175, "weight": 80, "city": "New York", "country": "USA",
        "occupation": "Software Engineer", "chief_complaints": "Lower back pain", "pain_level": 6,
        "pain_duration": "2 weeks", "comorbidities": "None", "lifestyle": "Sedentary",
        "pain_history": "Started after long hours of sitting", "aggravating_factor": "Prolonged sitting",
        "relieving_factor": "Walking", "patient_goal": "Return to normal work routine",
        "joints": "Lumbar spine", "observations": "Reduced lumbar lordosis",
        "clinical_assessment": "Muscle spasm in lower back", "medical_reports": "X-ray shows no abnormalities",
        "provisional_diagnosis": "Mechanical low back pain", "treatment_plan": "Physical therapy and ergonomic adjustments",
        "precautions": "Avoid prolonged sitting", "remarks_physio": "Focus on core strengthening",
        "attitude": "Motivated", "patient_persona": "Tech-savvy, busy professional"
    },
    "Patient 2": {
        "name": "Jane Smith",
        "age": 55, "gender": "Female", "height": 165, "weight": 70, "city": "London", "country": "UK",
        "occupation": "Teacher", "chief_complaints": "Knee pain", "pain_level": 7,
        "pain_duration": "3 months", "comorbidities": "Hypertension", "lifestyle": "Moderately active",
        "pain_history": "Gradual onset, worsening over time", "aggravating_factor": "Climbing stairs",
        "relieving_factor": "Rest and ice", "patient_goal": "Walk without pain",
        "joints": "Right knee", "observations": "Slight swelling",
        "clinical_assessment": "Crepitus on movement", "medical_reports": "MRI shows mild osteoarthritis",
        "provisional_diagnosis": "Osteoarthritis of the knee", "treatment_plan": "Physical therapy and weight management",
        "precautions": "Low-impact exercises only", "remarks_physio": "Gait training and knee stabilization exercises",
        "attitude": "Concerned", "patient_persona": "Dedicated educator, worried about mobility"
    },
    "Patient 3": {
        "name": "Emily Brown",
        "age": 28, "gender": "Female", "height": 160, "weight": 55, "city": "Sydney", "country": "Australia",
        "occupation": "Graphic Designer", "chief_complaints": "Neck and shoulder pain", "pain_level": 5,
        "pain_duration": "1 month", "comorbidities": "Migraine", "lifestyle": "Active",
        "pain_history": "Started after increased workload", "aggravating_factor": "Long hours at computer",
        "relieving_factor": "Stretching", "patient_goal": "Work without discomfort",
        "joints": "Cervical spine, shoulder", "observations": "Forward head posture",
        "clinical_assessment": "Tight upper trapezius", "medical_reports": "No imaging done",
        "provisional_diagnosis": "Work-related musculoskeletal disorder", "treatment_plan": "Ergonomic assessment, posture correction",
        "precautions": "Regular breaks from computer work", "remarks_physio": "Focus on scapular stabilization",
        "attitude": "Proactive", "patient_persona": "Creative professional, health-conscious"
    },
    "Patient 4": {
        "name": "Michael Johnson",
        "age": 45, "gender": "Male", "height": 180, "weight": 90, "city": "Toronto", "country": "Canada",
        "occupation": "Construction Worker", "chief_complaints": "Shoulder pain", "pain_level": 8,
        "pain_duration": "6 weeks", "comorbidities": "Type 2 Diabetes", "lifestyle": "Physically demanding job",
        "pain_history": "Injury while lifting heavy object", "aggravating_factor": "Overhead activities",
        "relieving_factor": "Rest and NSAIDs", "patient_goal": "Return to work full capacity",
        "joints": "Right shoulder", "observations": "Limited range of motion",
        "clinical_assessment": "Positive impingement tests", "medical_reports": "Ultrasound shows rotator cuff tendinopathy",
        "provisional_diagnosis": "Rotator cuff tendinopathy", "treatment_plan": "Physical therapy, gradual return to work",
        "precautions": "Avoid heavy lifting temporarily", "remarks_physio": "Rotator cuff strengthening program",
        "attitude": "Frustrated", "patient_persona": "Hardworking, eager to return to full duties"
    },
    "Patient 5": {
        "name": "Anna Schmidt",
        "age": 62, "gender": "Female", "height": 170, "weight": 75, "city": "Berlin", "country": "Germany",
        "occupation": "Retired", "chief_complaints": "Hip pain", "pain_level": 6,
        "pain_duration": "4 months", "comorbidities": "Osteoporosis", "lifestyle": "Moderately active",
        "pain_history": "Gradual onset, worse in mornings", "aggravating_factor": "Prolonged walking",
        "relieving_factor": "Warm compress", "patient_goal": "Maintain independence in daily activities",
        "joints": "Left hip", "observations": "Antalgic gait",
        "clinical_assessment": "Reduced internal rotation", "medical_reports": "X-ray shows mild joint space narrowing",
        "provisional_diagnosis": "Early hip osteoarthritis", "treatment_plan": "Physical therapy, aquatic exercises",
        "precautions": "Fall prevention strategies", "remarks_physio": "Focus on hip mobility and strength",
        "attitude": "Determined", "patient_persona": "Active retiree, enjoys gardening"
    }
}

# Function to create a row of 3 inputs
def create_input_row(col1_input, col2_input, col3_input):
    col1, col2, col3 = st.columns(3)
    with col1:
        val1 = col1_input()
    with col2:
        val2 = col2_input()
    with col3:
        val3 = col3_input()
    return val1, val2, val3

# Create buttons for patient profiles
st.subheader("Quick Patient Profiles")
cols = st.columns(5)
for i, (profile_name, profile_data) in enumerate(patient_profiles.items()):
    if cols[i].button(profile_name):
        st.session_state.update(profile_data)
        user_name = st.session_state.get("name", "")  # Update the name field

# Create a form for user inputs
with st.form(key='patient_info_form'):
    # User input rows
    name, age, gender = create_input_row(
        lambda: st.text_input("Patient Name", value=st.session_state.get("name", "")),
        lambda: st.number_input("Age", min_value=0, max_value=120, value=st.session_state.get("age", 30)),
        lambda: st.selectbox("Gender", ["Male", "Female", "Other"], index=["Male", "Female", "Other"].index(st.session_state.get("gender", "Male")))
    )

    height, weight, city = create_input_row(
        lambda: st.number_input("Height (cm)", min_value=0, max_value=300, value=st.session_state.get("height", 170)),
        lambda: st.number_input("Weight (kg)", min_value=0, max_value=500, value=st.session_state.get("weight", 70)),
        lambda: st.text_input("City/Town", value=st.session_state.get("city", ""))
    )

    country, occupation, chief_complaints = create_input_row(
        lambda: st.text_input("Country", value=st.session_state.get("country", "")),
        lambda: st.text_input("Occupation", value=st.session_state.get("occupation", "")),
        lambda: st.text_area("Chief Complaints", value=st.session_state.get("chief_complaints", ""))
    )

    pain_level, pain_duration, comorbidities = create_input_row(
        lambda: st.number_input("Pain Level", 0, 10, value=st.session_state.get("pain_level", 5)),
        lambda: st.text_input("Pain Duration", value=st.session_state.get("pain_duration", "")),
        lambda: st.text_area("Co-morbidities", value=st.session_state.get("comorbidities", ""))
    )

    lifestyle, pain_history, aggravating_factor = create_input_row(
        lambda: st.text_area("Lifestyle", value=st.session_state.get("lifestyle", "")),
        lambda: st.text_area("History of this Pain", value=st.session_state.get("pain_history", "")),
        lambda: st.text_area("Aggravating Factor", value=st.session_state.get("aggravating_factor", ""))
    )

    relieving_factor, patient_goal, joints = create_input_row(
        lambda: st.text_area("Relieving Factor", value=st.session_state.get("relieving_factor", "")),
        lambda: st.text_area("Patient Eventual Goal", value=st.session_state.get("patient_goal", "")),
        lambda: st.text_area("Joints", value=st.session_state.get("joints", ""))
    )

    observations, clinical_assessment, medical_reports = create_input_row(
        lambda: st.text_area("Observations", value=st.session_state.get("observations", "")),
        lambda: st.text_area("Clinical Assessment", value=st.session_state.get("clinical_assessment", "")),
        lambda: st.text_area("Medical Reports Summary", value=st.session_state.get("medical_reports", ""))
    )

    provisional_diagnosis, treatment_plan, precautions = create_input_row(
        lambda: st.text_area("Provisional Diagnosis", value=st.session_state.get("provisional_diagnosis", "")),
        lambda: st.text_area("Treatment Plan Advised", value=st.session_state.get("treatment_plan", "")),
        lambda: st.text_area("Precautions/Advice", value=st.session_state.get("precautions", ""))
    )

    remarks_physio, attitude, patient_persona = create_input_row(
        lambda: st.text_area("Remarks for Physio", value=st.session_state.get("remarks_physio", "")),
        lambda: st.text_input("Attitude", value=st.session_state.get("attitude", "")),
        lambda: st.text_area("Patient Persona", value=st.session_state.get("patient_persona", ""))
    )

    # Add the Apply button at the end of the form
    apply_button = st.form_submit_button(label='Apply')

# Handle form submission
if apply_button:
    # Update session state with new values
    st.session_state.update({
        "name": name,
        "age": age,
        "gender": gender,
        "height": height,
        "weight": weight,
        "city": city,
        "country": country,
        "occupation": occupation,
        "chief_complaints": chief_complaints,
        "pain_level": pain_level,
        "pain_duration": pain_duration,
        "comorbidities": comorbidities,
        "lifestyle": lifestyle,
        "pain_history": pain_history,
        "aggravating_factor": aggravating_factor,
        "relieving_factor": relieving_factor,
        "patient_goal": patient_goal,
        "joints": joints,
        "observations": observations,
        "clinical_assessment": clinical_assessment,
        "medical_reports": medical_reports,
        "provisional_diagnosis": provisional_diagnosis,
        "treatment_plan": treatment_plan,
        "precautions": precautions,
        "remarks_physio": remarks_physio,
        "attitude": attitude,
        "patient_persona": patient_persona
    })
    st.success("Patient information updated successfully!")

# Display chat messages (excluding system message)
for message in st.session_state.messages[1:]:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
if prompt := st.chat_input("Type Anything to start"):
    # Prepare the user information
    user_info = f"""
    Name: {name}
    Age: {age}, Gender: {gender}, Height: {height} cm, Weight: {weight} kg
    Location: {city}, {country}
    Occupation: {occupation}
    Chief Complaints: {chief_complaints}
    Pain Level: {pain_level}, Pain Duration: {pain_duration}
    Co-morbidities: {comorbidities}
    Lifestyle: {lifestyle}
    Pain History: {pain_history}
    Aggravating Factor: {aggravating_factor}
    Relieving Factor: {relieving_factor}
    Patient Goal: {patient_goal}
    Joints: {joints}
    Observations: {observations}
    Clinical Assessment: {clinical_assessment}
    Medical Reports: {medical_reports}
    Provisional Diagnosis: {provisional_diagnosis}
    Treatment Plan: {treatment_plan}
    Precautions: {precautions}
    Remarks for Physio: {remarks_physio}
    Attitude: {attitude}
    Patient Persona: {patient_persona}
    """
    
    # Update the system prompt with user information
    updated_system_prompt = f"{st.session_state.system_prompt}\n\nCurrent Patient Information:\n{user_info}"
    
    # Update the first message in the conversation (system prompt)
    st.session_state.messages[0] = {"role": "system", "content": updated_system_prompt}
    
    # Append the user's question
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        for response in client.chat.completions.create(
            messages=st.session_state.messages,
            model=selected_model,
            stream=True,
        ):
            full_response += (response.choices[0].delta.content or "")
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
    st.session_state.messages.append({"role": "assistant", "content": full_response})

# Display a warning about API key
st.sidebar.warning("Make sure to set your Groq API key as an environment variable named GROQ_API_KEY")