src/test.py

import os

# ✅ Fix cache issues: set Hugging Face model cache to a writable directory
os.environ["HF_HOME"] = "/tmp"
os.environ["TRANSFORMERS_CACHE"] = "/tmp"

import streamlit as st
from datetime import datetime, timedelta, time
import pandas as pd
from io import StringIO
from twilio.rest import Client
from transformers import pipeline

# ---- TWILIO CONFIGURATION ----
TWILIO_ACCOUNT_SID = os.environ.get('TWILIO_ACCOUNT_SID')
TWILIO_AUTH_TOKEN = os.environ.get('TWILIO_AUTH_TOKEN')
TWILIO_PHONE_NUMBER = os.environ.get('TWILIO_PHONE_NUMBER')

employee_phone_numbers = {
    "E001": "+18777804236",
    "E002": "+18777804236",
}

client = Client(TWILIO_ACCOUNT_SID, TWILIO_AUTH_TOKEN)

# ---- GPT-2 PIPELINE FIXED ----
@st.cache_resource
def load_gpt2():
    return pipeline("text-generation", model="gpt2", model_kwargs={"cache_dir": "/tmp"}, max_length=512)

gpt2_pipeline = load_gpt2()

# ---- STREAMLIT SETUP ----
st.set_page_config(page_title="AI Shift Filler for Workforce Management")
st.title("🏥 AI Shift Shortage Detection and Auto-Fulfillment")

# ---- USE CASE DESCRIPTION ----
st.markdown("""
### 👷 Workforce Management Use Case

This AI Agent automatically:
- 📉 Detects staff shortages
- 🧠 Identifies qualified, available employees
- 📩 Sends SMS/email shift offers
- ✅ Fills the open shift upon acceptance
""")

# ---- EMPLOYEE & SHIFT DATA ----
employee_data = """
ID,Name,Skills,Certifications,Available,OvertimeHours
E001,Alice,"ICU","ACLS",True,5
E002,Bob,"ER","BLS",False,12
E003,Charlie,"ICU","ACLS",True,8
"""

shift_data = """
ShiftID,Department,RequiredSkill,RequiredCert,ShiftTime
S101,ICU,ICU,ACLS,2025-06-04 07:00
"""

df_employees = pd.read_csv(StringIO(employee_data))
df_shifts = pd.read_csv(StringIO(shift_data))
df_employees['Skills'] = df_employees['Skills'].apply(lambda x: x.split(","))
df_employees['Certifications'] = df_employees['Certifications'].apply(lambda x: x.split(","))

# ---- MATCHING LOGIC ----
def find_eligible_employees(shift, employees):
    return employees[
        (employees['Skills'].apply(lambda s: shift['RequiredSkill'] in s)) &
        (employees['Certifications'].apply(lambda c: shift['RequiredCert'] in c)) &
        (employees['Available']) &
        (employees['OvertimeHours'] < 10)
    ]

# ---- PROCESS EACH SHIFT ----
results = []

for _, shift in df_shifts.iterrows():
    eligible = find_eligible_employees(shift, df_employees)
    if not eligible.empty:
        for _, emp in eligible.iterrows():
            phone_number = employee_phone_numbers.get(emp['ID'], None)
            if phone_number:
                sms_body = (
                    f"Shift Alert! Dear {emp['Name']}, "
                    f"a shift in {shift['Department']} at {shift['ShiftTime']} is available. "
                    f"Please reply to accept or decline."
                )
                try:
                    client.messages.create(
                        body=sms_body,
                        from_=TWILIO_PHONE_NUMBER,
                        to=phone_number
                    )
                    results.append((emp['Name'], shift['ShiftID'], "SMS Sent"))
                except Exception as e:
                    results.append((emp['Name'], shift['ShiftID'], f"Failed: {str(e)}"))
            else:
                results.append((emp['Name'], shift['ShiftID'], "No phone number"))
    else:
        results.append(("No eligible staff", shift['ShiftID'], "Shift Unfilled"))

# ---- DISPLAY RESULTS ----
st.subheader("📬 Shift Assignment Summary")
result_df = pd.DataFrame(results, columns=["Employee", "ShiftID", "Status"])
st.dataframe(result_df)

# ---- GPT-2 INSIGHT ----
st.subheader("🧠 GPT-2 Insights on Staffing")

if st.button("Generate Shift Fulfillment Recommendations"):
    shift_info = ""
    for _, row in df_shifts.iterrows():
        shift_info += (
            f"Shift {row['ShiftID']} in {row['Department']} at {row['ShiftTime']} requires {row['RequiredSkill']} with {row['RequiredCert']}.\n"
        )

    gpt_prompt = f"""
As an AI assistant, suggest strategies to handle shift shortages based on the following:
{shift_info}
Recommendations:
"""

    with st.spinner("AI Agent (GPT-2) is analyzing..."):
        try:
            result = gpt2_pipeline(gpt_prompt)[0]['generated_text']
            st.success("✅ GPT-2 Suggestion Ready")
            st.text_area("📋 GPT-2 Output", result, height=300)
        except Exception as e:
            st.error(f"Error using GPT-2: {str(e)}")