app.py

import gradio as gr
import os
import asyncio
import numpy as np
import random
from openai import AsyncOpenAI
import json
import time

# Initialize OpenAI client with better error handling
def get_openai_client():
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key or api_key == "":
        return None
    try:
        return AsyncOpenAI(api_key=api_key)
    except Exception as e:
        print(f"Failed to initialize OpenAI client: {e}")
        return None

openai_client = get_openai_client()

class ImprovPartner:
    def __init__(self):
        self.scene_contexts = [
            "a coffee shop where the coffee machine is broken",
            "a spaceship that's running out of fuel", 
            "a doctor's office where the doctor is allergic to patients",
            "a library where all the books are singing",
            "a pet store that only sells invisible animals",
            "a cooking show where everything is made of marshmallows",
            "a gym where everyone is moving in slow motion",
            "a bank that only trades in compliments",
            "a haunted house where the ghosts are afraid of people",
            "a classroom where the teacher is a student"
        ]
        self.current_scene = random.choice(self.scene_contexts)
        self.persona_traits = [
            "overly enthusiastic",
            "mysteriously quiet", 
            "dramatic and theatrical",
            "constantly confused",
            "extremely literal",
            "sings everything",
            "speaks only in questions",
            "uses pirate slang",
            "acts like a robot",
            "thinks they're a superhero"
        ]
        self.current_persona = random.choice(self.persona_traits)
        self.conversation_history = []
        
    async def generate_improv_response(self, user_input: str) -> str:
        """Generate an improv response based on user input."""
        global openai_client
        
        # Try to get client if not available
        if not openai_client:
            openai_client = get_openai_client()
        
        if not openai_client:
            # Fallback responses if no OpenAI client
            fallback_responses = [
                f"Yes, and that reminds me of something crazy that happened in {self.current_scene}!",
                "That's brilliant! What if we also...",
                "I love that! And what if suddenly...",
                "Perfect! And at that exact moment...",
                f"Absolutely! And then out of nowhere, in {self.current_scene}..."
            ]
            return random.choice(fallback_responses)
        
        try:
            system_prompt = f"""
            You are an improv partner in a scene. Be creative, supportive, and move the scene forward.
            
            Current scene: {self.current_scene}
            Your character traits: {self.current_persona}
            
            Rules:
            - Always say "Yes, and..." or equivalent to accept and build on what the partner says
            - Keep responses concise (1-2 sentences max)
            - Stay in character and scene
            - Be playful and supportive
            - End with an open-ended suggestion or question to keep the scene moving
            - Never break character or mention you're an AI
            
            Recent dialogue:
            User: "{user_input}"
            """
            
            response = await openai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_input}
                ],
                max_tokens=100,
                temperature=0.9
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            print(f"OpenAI API error: {e}")
            # Fallback responses if API fails
            fallback_responses = [
                f"Yes, and that reminds me of something crazy that happened in {self.current_scene}!",
                "That's brilliant! What if we also...",
                "I love that! And what if suddenly...",
                "Perfect! And at that exact moment...",
                "Absolutely! And then out of nowhere..."
            ]
            return random.choice(fallback_responses)

# Global improv partner instance
improv_partner = ImprovPartner()

def update_scene_info():
    """Return current scene information."""
    return improv_partner.current_scene, improv_partner.current_persona

def change_scene(option):
    """Change scene or persona based on user selection."""
    if option == "Change Scene":
        improv_partner.current_scene = random.choice(improv_partner.scene_contexts)
        improv_partner.conversation_history = []
    elif option == "Change Persona":
        improv_partner.current_persona = random.choice(improv_partner.persona_traits)
    elif option == "Reset Scene":
        improv_partner.current_scene = random.choice(improv_partner.scene_contexts)
        improv_partner.current_persona = random.choice(improv_partner.persona_traits)
        improv_partner.conversation_history = []
    
    return update_scene_info()

def chat_response(user_input, history):
    """Generate chat response."""
    if not user_input or user_input.strip() == "":
        return history, "Please say something to start the scene!"
    
    try:
        # Add user input to history
        history.append([user_input, ""])
        
        # Generate AI response (synchronously for Gradio)
        import asyncio
        try:
            loop = asyncio.get_event_loop()
        except:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        
        response = loop.run_until_complete(improv_partner.generate_improv_response(user_input))
        
        # Update history with AI response
        history[-1][1] = response
        
        return history, ""
        
    except Exception as e:
        print(f"Error in chat_response: {e}")
        fallback = f"Yes, and that's interesting! What happens next in {improv_partner.current_scene}?"
        history.append([user_input, fallback])
        return history, ""

def generate_prompt_suggestions():
    """Generate prompt suggestions for the user."""
    suggestions = [
        "I just got here, what's going on?",
        "Oh no! Did you hear that noise?",
        "This is the strangest place I've ever been!",
        "Can you help me with something?",
        "I have something important to tell you...",
        "Wait a minute, I think I know you!",
        "Something incredible just happened!",
        "We need to get out of here, now!",
        "I have a brilliant idea!",
        "You won't believe what I just found!"
    ]
    return random.sample(suggestions, 3)

def check_api_status():
    """Check if OpenAI API is configured."""
    global openai_client
    if not openai_client:
        openai_client = get_openai_client()
    
    if openai_client:
        return "✅ OpenAI API Connected", gr.Button(visible=False), gr.Textbox(visible=False)
    else:
        return "❌ OpenAI API Not Configured", gr.Button(visible=True), gr.Textbox(visible=True)

def set_api_key(api_key):
    """Set OpenAI API key from user input."""
    if api_key and api_key.strip():
        # Set environment variable for this session
        os.environ["OPENAI_API_KEY"] = api_key.strip()
        global openai_client
        try:
            openai_client = AsyncOpenAI(api_key=api_key.strip())
            return "✅ API Key Set Successfully!", gr.Button(visible=False), gr.Textbox(visible=False)
        except Exception as e:
            return f"❌ Error setting API key: {e}", gr.Button(visible=True), gr.Textbox(visible=True)
    else:
        return "❌ Please enter a valid API key", gr.Button(visible=True), gr.Textbox(visible=True)

# Build the Gradio interface
with gr.Blocks(title="Improv Partner Chatbot", theme=gr.themes.Soft(), css="""
    .scene-info {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 15px;
        border-radius: 10px;
        color: white;
        margin: 10px 0;
    }
    .chat-container {
        height: 400px;
        overflow-y: auto;
    }
    .api-status {
        background: #f0f0f0;
        padding: 10px;
        border-radius: 5px;
        margin: 10px 0;
        text-align: center;
    }
""") as demo:
    
    gr.Markdown("""
    # 🎭 Real-Time Improv Partner
    ### Your AI companion for spontaneous theater!
    
    **Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)**
    
    ---
    """)
    
    # API Status Section
    with gr.Row():
        api_status = gr.Textbox(label="API Status", value="Checking...", interactive=False)
    
    with gr.Row(visible=False) as api_setup_row:
        api_key_input = gr.Textbox(
            label="Enter your OpenAI API Key:",
            placeholder="sk-...",
            type="password",
            scale=3
        )
        set_api_btn = gr.Button("Set API Key", variant="primary", scale=1)
    
    # Check API status on load
    demo.load(
        check_api_status,
        outputs=[api_status, api_setup_row, api_key_input]
    )
    
    # Set API key handler
    set_api_btn.click(
        set_api_key,
        inputs=[api_key_input],
        outputs=[api_status, api_setup_row, api_key_input]
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            # Scene information display
            with gr.Row():
                with gr.Column():
                    gr.Markdown('<div class="scene-info">📍 <strong>Current Scene:</strong></div>', elem_classes="scene-info")
                    scene_info = gr.Textbox(
                        value=improv_partner.current_scene,
                        interactive=False,
                        container=False,
                        show_label=False
                    )
                
                with gr.Column():
                    gr.Markdown('<div class="scene-info">🎭 <strong>Current Persona:</strong></div>', elem_classes="scene-info")
                    persona_info = gr.Textbox(
                        value=improv_partner.current_persona,
                        interactive=False,
                        container=False,
                        show_label=False
                    )
            
            # Chat interface
            gr.Markdown("### 💬 Scene Dialogue")
            chatbot = gr.Chatbot(
                height=400,
                show_copy_button=True,
                bubble_full_width=False,
                elem_classes="chat-container"
            )
            
            # Input area
            with gr.Row():
                msg = gr.Textbox(
                    label="Your line:",
                    placeholder="Say your line... (Remember: 'Yes, and...')",
                    scale=4
                )
                submit_btn = gr.Button("🎤 Speak", variant="primary", scale=1)
            
            # Prompt suggestions
            with gr.Row():
                gr.Markdown("💡 **Need inspiration? Try one of these:**")
            with gr.Row():
                suggestion1 = gr.Button(generate_prompt_suggestions()[0], size="sm")
                suggestion2 = gr.Button(generate_prompt_suggestions()[1], size="sm") 
                suggestion3 = gr.Button(generate_prompt_suggestions()[2], size="sm")
            
            # Clear button
            clear_btn = gr.Button("🗑️ Clear Scene", variant="secondary")
        
        with gr.Column(scale=1):
            gr.Markdown("""
            ### 🎛️ Scene Controls
            """)
            
            scene_options = gr.Dropdown(
                choices=["Change Scene", "Change Persona", "Reset Scene"],
                label="Scene Options",
                value="Reset Scene"
            )
            
            apply_btn = gr.Button("✨ Apply Changes", variant="primary")
            
            gr.Markdown("""
            ### 🎯 How to Play:
            
            1. **Accept & Build**: Always say "Yes, and..." to accept your partner's ideas
            2. **Stay in Character**: Keep the scene going with your persona
            3. **Be Creative**: Add unexpected but relevant details
            4. **Listen Carefully**: Build on what your partner says
            5. **Have Fun**: There are no wrong answers in improv!
            
            ### 📝 Tips:
            - Keep responses brief
            - Ask questions to keep dialogue flowing
            - Use physical actions and emotions
            - Don't block or negate ideas
            - Embrace the unexpected!
            """)
            
            # API Setup Info
            gr.Markdown("""
            ### 🔑 API Setup
            
            To get AI responses, you need an OpenAI API key:
            1. Go to [platform.openai.com](https://platform.openai.com/api-keys)
            2. Create an API key
            3. Enter it above when prompted
            4. The app will save it for your session
            
            Without an API key, you'll get generic responses.
            """)
            
            # Stats
            gr.Markdown("### 📊 Scene Stats")
            scene_count = gr.Number(label="Scenes Played", value=0, interactive=False)
            
            # Refresh suggestions button
            refresh_suggestions = gr.Button("🔄 New Suggestions", variant="secondary", size="sm")
    
    # Event handlers
    def handle_submit(user_input, history):
        if not user_input.strip():
            return history, ""
        
        new_history = history.copy()
        new_history, _ = chat_response(user_input, new_history)
        return new_history, ""
    
    # Main chat interaction
    msg.submit(handle_submit, [msg, chatbot], [chatbot, msg])
    submit_btn.click(handle_submit, [msg, chatbot], [chatbot, msg])
    
    # Suggestion buttons
    def use_suggestion(suggestion_text, current_history):
        return handle_submit(suggestion_text, current_history)
    
    suggestion1.click(use_suggestion, [suggestion1, chatbot], [chatbot, msg])
    suggestion2.click(use_suggestion, [suggestion2, chatbot], [chatbot, msg])
    suggestion3.click(use_suggestion, [suggestion3, chatbot], [chatbot, msg])
    
    # Scene controls
    def apply_scene_changes(option, current_count):
        new_scene, new_persona = change_scene(option)
        new_count = current_count + 1 if option == "Reset Scene" else current_count
        return new_scene, new_persona, new_count
    
    apply_btn.click(
        apply_scene_changes,
        [scene_options, scene_count],
        [scene_info, persona_info, scene_count]
    )
    
    # Clear conversation
    def clear_conversation():
        return [], ""
    
    clear_btn.click(clear_conversation, outputs=[chatbot, msg])
    
    # Refresh suggestions
    def refresh_suggestions_fn():
        return random.sample([
            "I just got here, what's going on?",
            "Oh no! Did you hear that noise?",
            "This is the strangest place I've ever been!",
            "Can you help me with something?",
            "I have something important to tell you...",
            "Wait a minute, I think I know you!",
            "Something incredible just happened!",
            "We need to get out of here, now!",
            "I have a brilliant idea!",
            "You won't believe what I just found!"
        ], 3)
    
    refresh_suggestions.click(
        refresh_suggestions_fn,
        outputs=[suggestion1, suggestion2, suggestion3]
    )
    
    # Welcome message
    demo.load(
        lambda: [[None, f"Welcome! I'm your improv partner. Our scene is {improv_partner.current_scene}, and I'm feeling {improv_partner.current_persona}. What's your opening line?"]],
        outputs=[chatbot]
    )

if __name__ == "__main__":
    demo.launch()