Update app.py

app.py

@@ -1,97 +1,356 @@
 import gradio as gr
 import os
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, __version__ as hf_version
 import random
-from typing import List, Tuple
+from typing import Generator, Dict, List, Tuple, Optional
+import logging  # Added logging for better debugging
+
+# Configure logging with DEBUG level and add version info
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logging.debug(f"Using huggingface_hub version: {hf_version}")
 
 # Get token from environment variable
 hf_token = os.environ.get("HF_TOKEN")
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token)
 
-# Story genres and starter prompts
+# Story genres with genre-specific example prompts
 GENRE_EXAMPLES = {
-    "Fantasy": [
-        "You enter the ancient forest seeking the wizard's tower.",
-        "Approaching the dragon cautiously, you raise your shield."
+    "fairy tale": [
+        "I follow the shimmer of fairy dust into a hidden forest",
+        "A tiny dragon appears at my window, asking for help to find its mother",
+        "A friendly witch invites me into her cozy cottage, offering a warm cup of tea"
+    ],
+    "fantasy": [
+        "I enter the ancient forest seeking the wizard's tower",
+        "I approach the dragon cautiously with my shield raised",
+        "I try to bargain with the elven council for safe passage"
+    ],
+    "sci-fi": [
+        "I investigate the strange signal coming from the abandoned planet",
+        "I negotiate with the alien ambassador about the peace treaty",
+        "I try to repair my damaged spacecraft before oxygen runs out"
+    ],
+    "mystery": [
+        "I examine the crime scene for overlooked evidence",
+        "I question the nervous butler about the night of the murder",
+        "I follow the suspicious figure through the foggy streets"
     ],
-    "Sci-Fi": [
-        "Hacking into the space station's mainframe, you uncover a secret.",
-        "Investigating the strange signal from the abandoned planet, you find more than you expected."
+    "horror": [
+        "I read the forbidden text while the candles flicker",
+        "I hide under the bed as footsteps approach",
+        "I investigate the strange noises coming from the attic"
     ],
-    "Mystery": [
-        "Examining the crime scene, you notice an overlooked clue.",
-        "Following the suspicious figure through foggy streets leads you to a discovery."
+    "western": [
+        "I challenge the outlaw to a duel at high noon",
+        "I track the bandits through the desert canyon",
+        "I enter the saloon looking for information"
     ],
-    "Horror": [
-        "You slowly open the creaking door to the basement.",
-        "Reading the forbidden text, the candles flicker ominously."
+    "cyberpunk": [
+        "I jack into the corporate mainframe to steal data",
+        "I hide in the neon-lit alleyway from corporate security",
+        "I meet my mysterious client in the underground bar"
+    ],
+    "historical": [
+        "I join the resistance against the occupying forces",
+        "I navigate the dangerous politics of the royal court",
+        "I set sail on a voyage to discover new lands"
+    ],
+    "post-apocalyptic": [
+        "I scavenge the abandoned shopping mall for supplies",
+        "I navigate through the radioactive zone using my old map",
+        "I hide from the approaching group of raiders"
+    ],
+    "steampunk": [
+        "I pilot my airship through the lightning storm",
+        "I present my new invention to the Royal Academy",
+        "I sneak aboard the emperor's armored train"
     ]
 }
 
-# System prompt for generating stories
-def get_story_prompt(genre: str) -> str:
-    return f"""You are an AI storyteller crafting an interactive {genre} adventure. 
-Write engaging scenes with vivid details and always end each response with three numbered choices.
-Follow the format:
-
-1. [Complete action-based choice]
-2. [Complete action-based choice]
-3. [Complete action-based choice]"""
-
-# Generate initial story
-def generate_story_intro(genre: str) -> str:
-    prompt = get_story_prompt(genre)
-    example = random.choice(GENRE_EXAMPLES[genre])
-    response = client.text_generation(prompt=f"{prompt}\n\nUser: {example}\nAssistant:", max_new_tokens=300)
-    return response[0]['generated_text'].strip()
-
-# Continue the story
-def continue_story(user_choice: str, history: List[Tuple[str, str]], genre: str) -> str:
-    prompt = get_story_prompt(genre)
-    conversation = "\n".join([f"User: {turn[0]}\nAssistant: {turn[1]}" for turn in history[-5:]])
-    full_prompt = f"{prompt}\n\n{conversation}\nUser: {user_choice}\nAssistant:"
-    response = client.text_generation(prompt=full_prompt, max_new_tokens=300)
-    return response[0]['generated_text'].strip()
-
-# Reset the conversation
-def reset_story() -> Tuple[List[Tuple[str, str]], str, str]:
-    return [], "", ""
-
-# Gradio UI Setup
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # 🌟 AI Story Studio
-        **Create Your Own Adventure with AI!**  
-        Choose a genre, start your journey, and guide the story with your choices.
+# 2. Add constants at the top for magic numbers
+MAX_HISTORY_LENGTH = 20
+MEMORY_WINDOW = 5  # Reduced from 10 to limit context
+MAX_TOKENS = 1024  # Reduced from 2048 for faster responses
+TEMPERATURE = 0.7  # Slightly reduced for faster convergence
+TOP_P = 0.95
+MIN_RESPONSE_LENGTH = 100  # Reduced from 200 for quicker display
+
+def get_examples_for_genre(genre):
+    """Get example prompts specific to the selected genre"""
+    return GENRE_EXAMPLES.get(genre, GENRE_EXAMPLES["fantasy"])
+
+def get_enhanced_system_prompt(genre=None):
+    """Generate a detailed system prompt with optional genre specification"""
+    selected_genre = genre or "fantasy"
+    system_message = f"""You are an interactive storyteller creating an immersive {selected_genre} choose-your-own-adventure story.
+For each response you MUST:
+1. Write 100-200 words describing the scene, using vivid sensory details
+2. Always use second-person perspective ("you", "your") to maintain reader immersion
+3. Include dialogue or your character's thoughts that reveal personality and motivations
+4. Create a strong sense of atmosphere appropriate for {selected_genre}
+5. End with EXACTLY THREE numbered choices and NOTHING ELSE AFTER THEM:
+1. [Complete sentence in second-person starting with a verb]
+2. [Complete sentence in second-person starting with a verb]
+3. [Complete sentence in second-person starting with a verb]
+CRITICAL RULES:
+- Provide only ONE set of three choices at the very end of your response
+- Never continue the story after giving choices
+- Never provide additional choices
+- Keep all narrative before the choices
+- End every response with exactly three numbered options
+- Each choice must start with "You" followed by a verb
+Remember: The story continues ONLY when the player makes a choice."""
+    return system_message
+
+def create_story_summary(chat_history):
+    """Create a concise summary of the story so far if the history gets too long"""
+    if len(chat_history) <= 2:
+        return None
+    
+    story_text = ""
+    for i in range(0, len(chat_history), 2):
+        if i+1 < len(chat_history):
+            story_text += f"User: {chat_history[i]}\nStory: {chat_history[i+1]}\n\n"
+    
+    summary_instruction = {
+        "role": "system",
+        "content": "The conversation history is getting long. Please create a brief summary of the key plot points and character development so far to help maintain context without exceeding token limits."
+    }
+    return summary_instruction
+
+# Modified function for proper Gradio format (lists)
+def respond(message: str, chat_history: List[Tuple[str, str]], genre: Optional[str] = None, use_full_memory: bool = True) -> Tuple[str, List[Tuple[str, str]]]:
+    """Generate a response based on the current message and conversation history."""
+    if not message.strip():
+        return "", chat_history
+    
+    try:
+        # Start with system prompt
+        api_messages = [{"role": "system", "content": get_enhanced_system_prompt(genre)}]
+        logging.debug(f"System Message: {api_messages[0]}")
         
-        ## 🕹️ How to Play:
-        1. **Pick a genre** from the dropdown.
-        2. **Start with a story prompt** or enter your own beginning.
-        3. **Choose from the AI's options** or type your own response to shape the adventure!
+        # Add chat history - convert from tuples to API format
+        if chat_history and use_full_memory:
+            for user_msg, bot_msg in chat_history[-MEMORY_WINDOW:]:
+                api_messages.extend([
+                    {"role": "user", "content": str(user_msg)},
+                    {"role": "assistant", "content": str(bot_msg)}
+                ])
+            logging.debug(f"Chat History Messages: {api_messages[1:]}")
         
-        ✨ *Tip: Your choices affect the story's outcome!*
-        """
-    )
+        # Add current message
+        api_messages.append({"role": "user", "content": str(message)})
+        logging.debug(f"Final Message List: {api_messages}")
 
-    with gr.Row():
-        genre_dropdown = gr.Dropdown(
-            choices=list(GENRE_EXAMPLES.keys()), label="Select Genre", value="Fantasy"
+        # Make API call without timeout
+        logging.debug("Making API call...")
+        response = client.chat_completion(
+            messages=api_messages,
+            max_tokens=MAX_TOKENS,
+            temperature=TEMPERATURE,
+            top_p=TOP_P
         )
-        start_story_btn = gr.Button("Start New Story")
+        logging.debug("API call completed")
+        
+        # Extract response
+        bot_message = response.choices[0].message.content
+        logging.debug(f"Bot Response: {bot_message[:100]}...")
+        
+        # Update history using tuple format [(user_msg, bot_msg), ...]
+        updated_history = list(chat_history)  # Create a copy
+        updated_history.append((message, bot_message))  # Add as tuple
+        return "", updated_history
+    
+    except Exception as e:
+        logging.error("Error in respond function", exc_info=True)
+        error_msg = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"
+        return "", list(chat_history) + [(message, error_msg)]
+
+def save_story(chat_history):
+    """Convert chat history to markdown and return as downloadable file"""
+    if not chat_history:
+        return gr.File.update(value=None)
+    
+    try:
+        story_text = "# My Adventure\n\n"
+        for user_msg, bot_msg in chat_history:
+            story_text += f"**Player:** {user_msg}\n\n"
+            story_text += f"**Story:** {bot_msg}\n\n---\n\n"
+        
+        # Create temporary file
+        temp_file = "my_story.md"
+        with open(temp_file, "w", encoding="utf-8") as f:
+            f.write(story_text)
+        
+        return temp_file
+        
+    except Exception as e:
+        logging.error(f"Error saving story: {e}")
+        return gr.File.update(value=None)
+
+# Add this function to get a custom avatar image URL
+def get_storyteller_avatar_url():
+    """Get a URL for the storyteller avatar from a free image service"""
+    # Using an external wizard avatar image
+    return "https://api.dicebear.com/7.x/bottts/svg?seed=wizard&backgroundColor=b6e3f4&eyes=bulging"
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    # Header section with improved instructions
+    gr.Markdown("""
+    # 🔮 AI Story Studio
+**Collaborate with AI to craft your own adventure, one scene at a time.**  
+Pick a genre, start with a prompt or write your own, and guide the story with your choices.  
+
+    > **Tip:** The more detail you provide, the deeper the story becomes.
+    """)
+    
+    wizard_avatar = get_storyteller_avatar_url()
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            # Chat window + user input - USING LIST FORMAT
+            chatbot = gr.Chatbot(
+                height=500,
+                bubble_full_width=True,
+                show_copy_button=True,
+                avatar_images=(None, wizard_avatar),
+                container=True,
+                scale=1,
+                min_width=800,
+                value=[],  # Empty list for messages
+                render=True
+            )
+            msg = gr.Textbox(
+                placeholder="Describe your next move...",
+                container=False,
+                scale=4,
+            )
+            
+            with gr.Row():
+                submit = gr.Button("Continue Story", variant="primary")
+                clear = gr.Button("Start New Adventure")
+                
+        with gr.Column(scale=1):
+            gr.Markdown("## Adventure Settings")
+            genre = gr.Dropdown(
+                choices=list(GENRE_EXAMPLES.keys()),
+                label="Story Genre",
+                info="Choose the theme of your next adventure",
+                value="fantasy"
+            )
+            full_memory = gr.Checkbox(
+                label="Full Story Memory",
+                value=True,
+                info="When enabled, the AI tries to remember the entire story. If disabled, only the last few exchanges are used."
+            )
+            
+            gr.Markdown("## Story Starters")
+            
+            # Create four placeholder buttons for story starters
+            starter_btn1 = gr.Button("Starter 1")
+            starter_btn2 = gr.Button("Starter 2")
+            starter_btn3 = gr.Button("Starter 3")
+            starter_btn4 = gr.Button("Starter 4")
+            starter_buttons = [starter_btn1, starter_btn2, starter_btn3, starter_btn4]
+
+            # Simplified update function
+            def update_starter_buttons(selected_genre):
+                examples = get_examples_for_genre(selected_genre)
+                results = []
+                for i in range(4):
+                    if i < len(examples):
+                        results.append(examples[i])
+                    else:
+                        results.append("")
+                return tuple(results)
 
-    chat_history = gr.Chatbot(height=400)
-    user_input = gr.Textbox(placeholder="Type your next move...")
+            # New direct handler for starter clicks
+            def use_starter(starter_text: str, history: List[Tuple[str, str]], selected_genre: str, memory_flag: bool) -> Tuple[str, List[Tuple[str, str]]]:
+                """Handle starter button clicks with proper message formatting"""
+                if not starter_text:
+                    return "", history
+                
+                try:
+                    # Use the respond function for consistent handling
+                    _, updated_history = respond(
+                        message=starter_text,
+                        chat_history=history,
+                        genre=selected_genre,
+                        use_full_memory=memory_flag
+                    )
+                    return "", updated_history
+                                
+                except Exception as e:
+                    error_msg = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"
+                    return "", list(history) + [(starter_text, error_msg)]
 
+            # Simplified button connections
+            for starter_button in starter_buttons:
+                starter_button.click(
+                    fn=use_starter,
+                    inputs=[starter_button, chatbot, genre, full_memory],
+                    outputs=[msg, chatbot],
+                    queue=True
+                )
+
+            # Update buttons when genre changes
+            genre.change(
+                fn=update_starter_buttons,
+                inputs=[genre],
+                outputs=starter_buttons
+            )
+
+    # Handler for user input
+    msg.submit(
+        fn=respond,
+        inputs=[msg, chatbot, genre, full_memory],
+        outputs=[msg, chatbot]
+    )
+    submit.click(
+        fn=respond,
+        inputs=[msg, chatbot, genre, full_memory],
+        outputs=[msg, chatbot]
+    )
+    
+    # Clear the chatbot for a new adventure
+    clear.click(lambda: [], None, chatbot, queue=False)
+    clear.click(lambda: "", None, msg, queue=False)
+    
+    # "Download My Story" row
     with gr.Row():
-        submit_btn = gr.Button("Continue Story", variant="primary")
-        clear_btn = gr.Button("Reset")
+        save_btn = gr.Button("Download My Story", variant="secondary")
+        story_output = gr.File(
+            label="Download your story",
+            file_count="single",
+            file_types=[".md"],
+            interactive=False,
+            visible=True
+        )
+    
+    save_btn.click(
+        fn=save_story,
+        inputs=[chatbot],
+        outputs=story_output,
+        queue=False  # Process immediately
+    )
 
-    # Function connections
-    start_story_btn.click(fn=generate_story_intro, inputs=[genre_dropdown], outputs=chat_history)
-    submit_btn.click(fn=continue_story, inputs=[user_input, chat_history, genre_dropdown], outputs=chat_history)
-    clear_btn.click(fn=reset_story, inputs=[], outputs=[chat_history, user_input])
+    # Initialize buttons with default fantasy genre examples
+    initial_examples = get_examples_for_genre("fantasy")
+    initial_button_data = tuple(
+        initial_examples[i] if i < len(initial_examples) else ""
+        for i in range(4)
+    )
+    
+    # Update button text on page load
+    demo.load(
+        fn=lambda: initial_button_data,
+        outputs=starter_buttons,
+        queue=False
+    )
 
-# Launch the app
+# Run the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)