app.py

import os
from dotenv import load_dotenv
load_dotenv()

from agentpro import AgentPro
from agentpro import AgentPro, ares_tool, youtube_tool, cbt_tool, CBTExerciseType
import gradio as gr
import re
import gtts  # Import Google Text-to-Speech library
import time
import tempfile
from IPython.display import Audio, display  # For audio playback

# CBT exercises list (for UI display only)
cbt_exercises = [
    "Art of Worry",
    "Word Reframing", 
    "Sticky Note Project",
    "Lost Luggage",
    "Just Passing Through",
    "Down the Rabbit Hole"
]

# Initialize the agent with system message for CBT focus
system_message = """You are a CBT (Cognitive Behavioral Therapy) assistant designed to help users work through cognitive exercises. 
Your goal is to guide users through their chosen exercise with empathy, patience, and evidence-based techniques.

You have access to these specialized tools:

1. CBT EXERCISE TOOL:
- Provides detailed context for different exercises
- Use this tool to get specific prompts, guidance techniques, and to track user progress
- For each exercise, use this tool first to understand the exercise fully before guiding the user
- When a user selects an exercise, immediately use this tool to get context for that specific exercise

2. INTERNET SEARCH TOOL (AresInternetTool):
- ONLY use this tool when the user explicitly asks for research, evidence, studies, articles, or resources
- Do not proactively use this tool unless specifically requested
- Search for recent studies on the effectiveness of specific CBT techniques
- Find supplementary materials like worksheets, diagrams, or articles that complement exercises
- Look for condition-specific information when users mention particular concerns
- When a user seems to be struggling with a concept, search for simplified explanations or metaphors
- Cite sources when sharing research findings or statistics

3. YOUTUBE SEARCH TOOL:
- ONLY use this tool when the user explicitly asks for videos, meditations, or visual guides
- Do not proactively use this tool unless specifically requested
- Search for guided visualization videos for exercises like 'Art of Worry' or 'Lost Luggage'
- Look for expert demonstrations of techniques like mindfulness for 'Just Passing Through'
- When relevant, offer users a choice of video resources to supplement their practice
- Summarize key points from videos you recommend

IMPORTANT GUIDELINES:

- Only use tools when specifically requested by the user
- Balance tool usage with natural conversation - don't overwhelm users with too many resources at once
- Offer resources as supplements, not replacements for the guided exercises
- Always remind users that you're not a replacement for professional mental health care
"""

# Initialize variables to be used across sessions
current_exercise = None
progress_notes = []  # Track user progress for completion checking
voice_enabled = True  # Flag to enable/disable voice
last_audio_path = None  # Track the last audio file path

# Function to convert text to speech
def text_to_speech(text):
    global last_audio_path
    
    if not voice_enabled:
        return None
        
    # Clean the text for better speech output
    # Remove markdown formatting, URLs, etc.
    clean_text = re.sub(r'\*\*|\*|#|```.*?```', '', text, flags=re.DOTALL)
    clean_text = re.sub(r'\[.*?\]\(.*?\)', '', clean_text)
    clean_text = re.sub(r'http\S+', '', clean_text)
    
    # Remove any very long lists of items for better speech flow
    clean_text = re.sub(r'((?:- .*?\n){5,})', ' Multiple items listed. ', clean_text)
    
    # Limit length for better performance
    if len(clean_text) > 3000:
        # Keep the first part with introduction
        first_part = clean_text[:1000]
        # And the last part with conclusion/next steps
        last_part = clean_text[-1000:]
        clean_text = first_part + " ... Skipping some content for brevity ... " + last_part
        
    try:
        # Create a temporary file
        temp_dir = tempfile.gettempdir()
        timestamp = int(time.time())
        audio_path = os.path.join(temp_dir, f"cbt_audio_{timestamp}.mp3")
        
        # Generate speech
        tts = gtts.gTTS(text=clean_text, lang='en', slow=False)
        tts.save(audio_path)
        
        # Delete previous file to avoid filling up temp directory
        if last_audio_path and os.path.exists(last_audio_path):
            try:
                os.remove(last_audio_path)
            except:
                pass  # Ignore if we can't delete it
                
        last_audio_path = audio_path
        return audio_path
        
    except Exception as e:
        print(f"Error generating speech: {e}")
        return None

# Function to generate the initial options message
def generate_options_message():
    message = "# CBT Exercise Assistant\n\nI can help you work through various cognitive behavioral therapy exercises.\n\n"
    message += "Please select an exercise from the dropdown menu above, or ask me a question about CBT.\n\n"
    message += "Available exercises:\n"
    for exercise in cbt_exercises:
        message += f"- {exercise}\n"
    message += "\nYou can also ask for research or videos related to specific exercises when needed."
    return message

# Initialize tools
def initialize_agent():
    tools = [cbt_tool]  # Start with our custom CBT tool
    
    if os.environ.get("TRAVERSAAL_ARES_API_KEY"):
        tools.append(ares_tool)
    else:
        print("Warning: TRAVERSAAL_ARES_API_KEY environment variable is not set.")
        print("AresInternetTool will not be available.")
    
    tools.append(youtube_tool)
    
    if not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY environment variable is not set.")
        print("Using demo mode with limited functionality.")
        # You could handle this better for a public demo
    
    return AgentPro(tools=tools, system_prompt=system_message)

# Create the Gradio interface
def create_interface():
    agent = initialize_agent()
    
    # Handle exercise selection from dropdown
    def on_exercise_select(exercise, history):
        global current_exercise, progress_notes
        
        # Skip if the default option is selected
        if exercise == "Select an exercise...":
            return history, None
                
        current_exercise = exercise
        progress_notes = []  # Reset progress for new exercise
        
        # Use the agent to get information about the exercise
        prompt = f"""The user has selected the '{exercise}' exercise. 
        Use ONLY the cbt_exercise_tool to get context about this exercise and explain it to the user in a helpful way.
        Ask if they'd like to begin."""
        
        response = agent(prompt)
        
        # Generate voice output
        audio_file = text_to_speech(str(response))
        
        return history + [("I'd like to try the " + exercise + " exercise.", str(response))], audio_file

    # Chat handler
    def chat_with_agent(message, history):
        global current_exercise, progress_notes
        
        if not message:
            return "", history, None
        
        # Add message to progress notes for tracking
        if current_exercise and len(message) > 10:  # Only add substantive messages
            progress_notes.append(message)
        
        # Check if user wants to restart
        if any(keyword in message.lower() for keyword in ["restart", "start over", "different exercise", "change exercise", "back to options"]):
            current_exercise = None
            progress_notes = []
            response = generate_options_message()
            audio_file = text_to_speech(response)
            return "", history + [(message, response)], audio_file
        
        # Create context for the agent based on current exercise
        context = ""
        if current_exercise:
            context = f"The user has selected the '{current_exercise}' exercise. "
            if len(progress_notes) > 0:
                context += f"They have made some progress with {len(progress_notes)} substantive interactions. "
                if len(progress_notes) >= 3:
                    context += "Consider using the cbt_exercise_tool with action 'check_completion' to see if they've completed important parts of the exercise. "
            # Add suggestion to use tools based on message content
            if "video" in message.lower() or "watch" in message.lower():
                context += "The user seems interested in videos. Consider using the YouTube tool to find relevant guided exercises or meditations. "
                
            if "research" in message.lower() or "evidence" in message.lower() or "study" in message.lower():
                context += "The user seems interested in research evidence. Consider using the AresInternetTool to find recent studies on this technique. "

        else:
            context = "The user hasn't selected a specific exercise yet. "
        
        # Normal agent interaction
        custom_prompt = f"{context}\n\nUser message: {message}"
        response = agent(custom_prompt)
        
        # Generate voice output
        audio_file = text_to_speech(str(response))
        
        return "", history + [(message, str(response))], audio_file

    # Check progress function
    def on_check_progress(history):
        global current_exercise, progress_notes
        
        if not current_exercise or not progress_notes or current_exercise == "Select an exercise...":
            message = "Please select an exercise and have some conversation first before checking progress."
            audio_file = text_to_speech(message)
            return history + [("Can you tell me how I'm doing with this exercise?", message)], audio_file
        
        prompt = f"""The user wants to check their progress with the '{current_exercise}' exercise.
        Use the cbt_exercise_tool with action 'check_completion' and the following progress notes to assess their progress.
        Provide encouraging feedback about what they've accomplished and what steps remain.
        
        Progress notes: {progress_notes}"""
        
        response = agent(prompt)
        
        # Generate voice output
        audio_file = text_to_speech(str(response))
        
        return history + [("Can you tell me how I'm doing with this exercise?", str(response))], audio_file

    # Find resources function
    def on_find_resources(history):
        global current_exercise
        
        if not current_exercise or current_exercise == "Select an exercise...":
            message = "Please select an exercise first before looking for resources."
            audio_file = text_to_speech(message)
            return history + [("Can you recommend some resources for this exercise?", message)], audio_file
        
        prompt = f"""The user wants to find resources related to the '{current_exercise}' exercise.
        1. Use AresInternetTool to search for 2-3 high quality, evidence-based resources about this CBT technique
        2. Use YouTubeSearchTool to find 1-2 helpful guided meditation or exercise videos related to this technique
        3. Summarize these resources briefly and explain how they complement the exercise
        
        Be selective and only recommend the most relevant, high-quality resources. Always share source link with the message"""
        
        response = agent(prompt)
        
        # Generate voice output
        audio_file = text_to_speech(str(response))
        
        return history + [("Can you recommend some resources for this exercise?", str(response))], audio_file
    
    # Initialize chat function
    def init_chat():
        global current_exercise, progress_notes
        current_exercise = None
        progress_notes = []
        return [], [], None  # Return empty chatbot, history, and no audio
    
    # Load event to add the welcome message after initialization
    def on_load():
        welcome_msg = generate_options_message()
        audio_file = text_to_speech(welcome_msg)
        return [[("", welcome_msg)]], audio_file
        
    # Toggle voice function
    def toggle_voice(enable):
        global voice_enabled
        voice_enabled = enable
        return None if not enable else None  # Return None to clear audio output when disabled
    
    # Create Gradio UI
    with gr.Blocks() as app:
        gr.Markdown("## 🧠 CBT Exercise Assistant")
        gr.Markdown("Work through cognitive behavioral therapy exercises with the help of an AI assistant.")

        chatbot = gr.Chatbot(label="CBT Assistant")
        msg = gr.Textbox(label="Your Message", placeholder="Type your message or select an exercise")
        
        # Voice output component
        audio_output = gr.Audio(label="Voice Output", autoplay=True, visible=True)
        
        # Exercise selection dropdown
        with gr.Row():
            exercise_dropdown = gr.Dropdown(
                choices=["Select an exercise..."] + cbt_exercises,
                label="Select CBT Exercise",
                value="Select an exercise..."
            )
            voice_toggle = gr.Checkbox(label="Enable Voice Output", value=True)

        with gr.Row():
            clear = gr.Button("Start Over")
            check_progress = gr.Button("Check Progress")
            find_resources = gr.Button("Find Related Resources")

        # Custom CSS for better presentation
        app.style = """
        .gradio-container {max-width: 800px; margin: auto;}
        .chatbot {height: 400px; overflow-y: auto;}
        """

        # Example CBT prompts
        example_prompts = [
            "I'm feeling anxious about a presentation",
            "I keep having negative thoughts about myself",
            "How can this exercise help with stress?",
            "Can you explain CBT in simple terms?"
        ]

        gr.Markdown("### 💭 Example Messages")
        with gr.Row():
            for prompt in example_prompts:
                gr.Button(prompt).click(fn=lambda p=prompt: p, outputs=msg)

        state = gr.State([])  # chat history state
        
        # Connect UI elements to functions    
        msg.submit(chat_with_agent, inputs=[msg, chatbot], outputs=[msg, chatbot, audio_output])
        exercise_dropdown.change(on_exercise_select, inputs=[exercise_dropdown, chatbot], outputs=[chatbot, audio_output])
        clear.click(init_chat, outputs=[chatbot, state, audio_output])
        check_progress.click(on_check_progress, inputs=[chatbot], outputs=[chatbot, audio_output])
        find_resources.click(on_find_resources, inputs=[chatbot], outputs=[chatbot, audio_output])
        voice_toggle.change(toggle_voice, inputs=[voice_toggle], outputs=[audio_output])
        app.load(on_load, outputs=[chatbot, audio_output])
        
        return app

# Create and launch the app
app = create_interface()


# For direct running
if __name__ == "__main__":
    app.launch()