app.py

import streamlit as st

# 1. Set page config FIRST
st.set_page_config(
    page_title="TARS: Therapist Assistance and Response System", 
    page_icon="🧠"
)

import os
import torch
from transformers import pipeline

# Canadian Crisis Resources (example)
CRISIS_RESOURCES = {
    "Canada Suicide Prevention Service": "1-833-456-4566",
    "Crisis Services Canada": "1-833-456-4566",
    "Kids Help Phone": "1-800-668-6868",
    "First Nations and Inuit Hope for Wellness Help Line": "1-855-242-3310"
}

# Keywords for detecting potential self-harm or suicidal language
SUICIDE_KEYWORDS = [
    "suicide", "kill myself", "end my life", 
    "want to die", "hopeless", "no way out", 
    "better off dead", "pain is too much"
]

class AITherapistAssistant:
    def __init__(self, conversation_model_name="microsoft/phi-1_5", summary_model_name="facebook/bart-large-cnn"):
        """
        Initialize the conversation (LLM) model and the summarization model.
        
        If you truly have a different 'phi2' from Microsoft, replace 'microsoft/phi-1_5'
        with your private or custom Hugging Face repo name.
        """
        # Load conversation LLM (phi2 / phi-1_5)
        try:
            self.conversation_model = pipeline(
                "text-generation",
                model=conversation_model_name,
                device=0 if torch.cuda.is_available() else -1
            )
        except Exception as e:
            st.error(f"Error loading conversation model: {e}")
            self.conversation_model = None

        # Load summarization model (BART Large CNN as default)
        try:
            self.summary_model = pipeline(
                "summarization", 
                model=summary_model_name, 
                device=0 if torch.cuda.is_available() else -1
            )
        except Exception as e:
            st.error(f"Error loading summary model: {e}")
            self.summary_model = None
    
    def detect_crisis(self, message: str) -> bool:
        """Check if message contains suicidal or distress-related keywords."""
        message_lower = message.lower()
        return any(keyword in message_lower for keyword in SUICIDE_KEYWORDS)
    
    def generate_response(self, message: str) -> str:
        """Generate a supportive AI response from the conversation model."""
        if not self.conversation_model:
            return (
                "I'm having some technical difficulties right now. "
                "I'm truly sorry I can't provide support at the moment. "
                "Would you consider reaching out to a human counselor or support helpline?"
            )
        
        # Create a more structured and empathetic prompt
        prompt = (
            "You are a compassionate, empathetic AI therapist trained to provide supportive, "
            "non-judgmental responses. Your goal is to validate feelings, offer gentle support, "
            "and help the user feel heard and understood.\n\n"
            "User's message: {}\n\n"
            "Your compassionate response:".format(message)
        )
        
        try:
            outputs = self.conversation_model(
                prompt, 
                max_length=300,  # Increased length for more nuanced responses
                num_return_sequences=1,
                do_sample=True,
                top_p=0.9,
                temperature=0.7
            )
            
            response_text = outputs[0]["generated_text"]
            
            # More robust prompt stripping
            if response_text.startswith(prompt):
                response_text = response_text[len(prompt):].strip()
            
            # Additional processing to ensure response quality
            response_text = response_text.split('\n')[0].strip()  # Take first coherent line
            
            # Fallback if response is too short or nonsensical
            if len(response_text) < 20:
                response_text = (
                    "I hear you. Your feelings are valid, and it takes courage to share what you're experiencing. "
                    "Would you like to tell me a bit more about what's on your mind?"
                )
            
            return response_text
        except Exception as e:
            st.error(f"Error generating response: {e}")
            return (
                "I'm experiencing some difficulties right now. "
                "Your feelings are important, and I want to be fully present for you. "
                "Would you be willing to try sharing again?"
            )
        
    def generate_summary(self, conversation_text: str) -> str:
        """Generate a summary of the entire conversation for the user's feeling. Cite from their input how they felt."""
        if not self.summary_model:
            return "Summary model is unavailable at the moment."
        
        try:
            summary_output = self.summary_model(
                conversation_text, 
                max_length=130, 
                min_length=30, 
                do_sample=False
            )
            return summary_output[0]["summary_text"]
        except Exception as e:
            st.error(f"Error generating summary: {e}")
            return "Sorry, I couldn't generate a summary."

def main():
    st.title("🧠 TARS: Therapist Assistance and Response System")
    st.write(
        "A supportive space to share your feelings safely.\n\n"
        "**Disclaimer**: I am not a licensed therapist. If you're in crisis, "
        "please reach out to professional help immediately."
    )
    
    # Note if running on Hugging Face Spaces
    if os.environ.get("SPACE_ID"):
        st.info("Running on Hugging Face Spaces.")
    
    # Instantiate the assistant with phi-1_5 (or your custom 'phi2')
    if "assistant" not in st.session_state:
        st.session_state.assistant = AITherapistAssistant(
            conversation_model_name="microsoft/phi-1_5",  # replace if needed
            summary_model_name="facebook/bart-large-cnn"
        )
    
    # Keep track of conversation
    if "conversation" not in st.session_state:
        st.session_state.conversation = []
    
    # Display existing conversation
    for message in st.session_state.conversation:
        if message["sender"] == "user":
            st.chat_message("user").write(message["text"])
        else:
            st.chat_message("assistant").write(message["text"])
    
    # Collect user input
    if prompt := st.chat_input("How are you feeling today?"):
        # Crisis detection
        if st.session_state.assistant.detect_crisis(prompt):
            st.warning("⚠️ Potential crisis detected.")
            st.markdown("**Immediate Support Resources (Canada):**")
            for org, phone in CRISIS_RESOURCES.items():
                st.markdown(f"- {org}: `{phone}`")
        
        # Display user message
        st.session_state.conversation.append({"sender": "user", "text": prompt})
        st.chat_message("user").write(prompt)
        
        # Generate AI response
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                ai_response = st.session_state.assistant.generate_response(prompt)
                st.write(ai_response)
        st.session_state.conversation.append({"sender": "assistant", "text": ai_response})
    
    # Summarize conversation
    if st.button("Generate Session Summary"):
        if st.session_state.conversation:
            conversation_text = " ".join(msg["text"] for msg in st.session_state.conversation)
            summary = st.session_state.assistant.generate_summary(conversation_text)
            st.subheader("Session Summary")
            st.write(summary)
        else:
            st.info("No conversation to summarize yet.")
    
    # Crisis Support Info in Sidebar
    st.sidebar.title("🆘 Crisis Support")
    st.sidebar.markdown("If you're in crisis, please contact:")
    for org, phone in CRISIS_RESOURCES.items():
        st.sidebar.markdown(f"- **{org}**: `{phone}`")

if __name__ == "__main__":
    main()