Update app.py

app.py

@@ -1,10 +1,9 @@
 import gradio as gr
 import os
 from huggingface_hub import InferenceClient, __version__ as hf_version
+import random
 from typing import Generator, Dict, List, Tuple, Optional
 import logging  # Added logging for better debugging
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
 
 # Configure logging with DEBUG level and add version info
 logging.basicConfig(
@@ -13,12 +12,9 @@ logging.basicConfig(
 )
 logging.debug(f"Using huggingface_hub version: {hf_version}")
 
-
 # Get token from environment variable
-hf_token = os.environ.get("QWEN_BOT_TOKEN")
-if not hf_token:
-    raise ValueError("Hugging Face token not found. Please set QWEN_BOT_TOKEN in the Hugging Face Secrets.")
-client = InferenceClient("Qwen/Qwen2.5-72B-Instruct", token=hf_token)
+hf_token = os.environ.get("HF_TOKEN")
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token)
 
 # Story genres with genre-specific example prompts
 GENRE_EXAMPLES = {
@@ -220,45 +216,165 @@ custom_css = """
 }
 """
 
-app = Flask(__name__)
-
-@app.route("/api/respond", methods=["POST"])
-def respond_api():
-    data = request.json
-    message = data.get("message", "")
-    chat_history = data.get("chat_history", [])
-    genre = data.get("genre", "fantasy")
-    use_full_memory = data.get("use_full_memory", True)
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) 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=400,
+                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 three placeholder buttons for story starters
+            starter_btn1 = gr.Button("Starter 1", scale=1, min_width=250, elem_classes="compact-btn")
+            starter_btn2 = gr.Button("Starter 2", scale=1, min_width=250, elem_classes="compact-btn")
+            starter_btn3 = gr.Button("Starter 3", scale=1, min_width=250, elem_classes="compact-btn")
+            starter_buttons = [starter_btn1, starter_btn2, starter_btn3]
 
-    # Generate system prompt
-    system_prompt = f"You are an interactive storyteller creating an immersive {genre} choose-your-own-adventure story."
+            # Simplified update function
+            def update_starter_buttons(selected_genre):
+                examples = get_examples_for_genre(selected_genre)
+                results = []
+                for i in range(3):
+                    if i < len(examples):
+                        results.append(examples[i])
+                    else:
+                        results.append("")
+                return tuple(results)
 
-    # Combine chat history and user message
-    input_text = system_prompt + "\n" + "\n".join(
-        [f"User: {msg}\nBot: {resp}" for msg, resp in chat_history]
-    ) + f"\nUser: {message}\nBot:"
+            # 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)]
 
-    try:
-        # Generate response using Hugging Face Inference API
-        logging.debug("Generating response via Inference API...")
-        response = client.text_generation(input_text, max_new_tokens=1024, temperature=0.7, top_p=0.95)
-        bot_message = response["generated_text"].split("Bot:")[-1].strip()
+            # 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 chat history
-        updated_history = chat_history + [(message, bot_message)]
-        return jsonify({"bot_message": bot_message, "updated_history": updated_history})
-    except Exception as e:
-        logging.error(f"Error generating response: {e}")
-        return jsonify({"error": str(e)}), 500
+            # Update buttons when genre changes
+            genre.change(
+                fn=update_starter_buttons,
+                inputs=[genre],
+                outputs=starter_buttons
+            )
 
-@app.route("/api/genres", methods=["GET"])
-def get_genres():
-    return jsonify(list(GENRE_EXAMPLES.keys()))
+    # 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 improved layout
+    with gr.Row(equal_height=True):  # Force equal height for all children
+        # Use Column for the button to control width
+        with gr.Column(scale=4):
+            save_btn = gr.Button("Download My Story", variant="secondary", size="lg")
+        
+        # Use Column for the file output with matching height
+        with gr.Column(scale=1):
+            story_output = gr.File(
+                label=None,  # Remove the label that adds extra height
+                file_count="single",
+                file_types=[".md"],
+                interactive=False,
+                visible=True,
+                elem_classes="compact-file-output"  # Optional: for custom CSS styling
+            )
+    
+    # Connect the save button to the save_story function
+    save_btn.click(
+        fn=save_story,
+        inputs=[chatbot],
+        outputs=story_output,
+        queue=False  # Process immediately
+    )
 
-@app.route("/api/examples", methods=["GET"])
-def get_examples():
-    genre = request.args.get("genre", "fantasy")
-    return jsonify(get_examples_for_genre(genre))
+    # 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(3)
+    )
+    
+    # Update button text on page load
+    demo.load(
+        fn=lambda: initial_button_data,
+        outputs=starter_buttons,
+        queue=False
+    )
 
+# Run the app
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000)
+    demo.launch(server_name="0.0.0.0", server_port=7860)