Update app.py

app.py

@@ -10,11 +10,11 @@ client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token)
 
 # Story genres with genre-specific example prompts
 GENRE_EXAMPLES = {
-    "fairy tale": [
-        "I follow the shimmer of fairy dust into a hidden forest",
-        "I meet a talking rabbit who claims to know a secret about the king's lost crown",
-        "A tiny dragon appears at my window, asking for help to find its mother",
-        "I step into a clearing where the trees whisper ancient riddles",
+    "fairy_tale": [
+        "I follow the shimmer of fairy dust into a hidden forest"
+        "I meet a talking rabbit who claims to know a secret about the king’s lost crown"
+        "A tiny dragon appears at my window, asking for help to find its mother"
+        "I step into a clearing where the trees whisper ancient riddles"
         "A friendly witch invites me into her cozy cottage, offering a warm cup of tea"
     ],
     "fantasy": [
@@ -73,7 +73,7 @@ GENRE_EXAMPLES = {
     ]
 }
 
-# Constants
+# 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
@@ -110,81 +110,66 @@ IMPORTANT:
 Keep the story cohesive by referencing previous events and choices."""
     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 user_msg, bot_msg in chat_history:
+        story_text += f"User: {user_msg}\nStory: {bot_msg}\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
+
+def format_history_for_gradio(history_tuples):
+    """Convert chat history to Gradio's message format."""
+    return [(str(user_msg), str(bot_msg)) for user_msg, bot_msg in history_tuples]
+
+# 1. Add type hints for better code maintainability
+# 4. Add input validation
 def respond(
     message: str, 
     chat_history: List[Tuple[str, str]], 
     genre: Optional[str] = None,
     use_full_memory: bool = True
-) -> List[Tuple[str, str]]:
+) -> List[Dict[str, str]]:  # Changed return type
     """Generate a response based on the current message and conversation history."""
     if not message.strip():
-        return chat_history
+        return [{"role": "assistant", "content": "Please provide a message"}]
 
     try:
-        # Create a copy of chat history to avoid modifying the original
-        new_history = list(chat_history)
-        
-        # Construct messages list for API call
-        messages = []
-        
-        # Always start with system message
-        messages.append({
-            "role": "system",
-            "content": get_enhanced_system_prompt(genre)
-        })
+        # Format messages for API
+        api_messages = [{"role": "system", "content": get_enhanced_system_prompt(genre)}]
         
-        # Add relevant chat history if enabled
-        if use_full_memory and new_history:
-            history_to_include = new_history[-MEMORY_WINDOW:] if len(new_history) > MEMORY_WINDOW else new_history
-            for user_msg, bot_msg in history_to_include:
-                if user_msg:  # Only add if user message exists
-                    messages.append({
-                        "role": "user",
-                        "content": str(user_msg)
-                    })
-                if bot_msg:  # Only add if bot message exists
-                    messages.append({
-                        "role": "assistant",
-                        "content": str(bot_msg)
-                    })
+        # Add chat history
+        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)}
+                ])
         
-        # Add current user message
-        messages.append({
-            "role": "user",
-            "content": str(message)
-        })
-        
-        # Make API call with strict dictionary format
-        response = client.post(
-            model="HuggingFaceH4/zephyr-7b-beta",
-            json={
-                "inputs": messages,
-                "parameters": {
-                    "max_new_tokens": MAX_TOKENS,
-                    "temperature": TEMPERATURE,
-                    "top_p": TOP_P
-                }
-            }
+        # Add current message
+        api_messages.append({"role": "user", "content": str(message)})
+
+        # Make API call
+        response = client.chat_completion(
+            messages=api_messages,
+            max_tokens=MAX_TOKENS,
+            temperature=TEMPERATURE,
+            top_p=TOP_P
         )
         
-        # Extract response text
-        if isinstance(response, dict) and "generated_text" in response:
-            bot_message = response["generated_text"]
-        elif isinstance(response, list) and response and "generated_text" in response[0]:
-            bot_message = response[0]["generated_text"]
-        else:
-            # Fallback in case the response structure is different
-            bot_message = str(response)
-        
-        # Add the new exchange to chat history
-        new_history.append((str(message), str(bot_message)))
-        return new_history
+        # Return properly formatted message
+        bot_message = response.choices[0].message.content
+        return [{"role": "assistant", "content": str(bot_message)}]
                         
     except Exception as e:
-        # Print the full error message for debugging
-        print(f"Error: {str(e)}")
-        # Return existing history plus error message
-        return chat_history + [(str(message), f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})")]
+        return [{"role": "assistant", "content": f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})"}]
 
 def save_story(chat_history):
     """Convert chat history to markdown for download"""
@@ -198,7 +183,6 @@ def save_story(chat_history):
     
     return story_text
 
-# Define the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🔮 Interactive Story Time")
     with gr.Row():
@@ -209,15 +193,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=3):
             # Chat window + user input
             chatbot = gr.Chatbot(
-                height=500,
-                bubble_full_width=True,
+                height=500,  # Increased height
+                bubble_full_width=True,  # Allow bubbles to use full width
                 show_copy_button=True,
                 avatar_images=(None, "🧙"),
                 type="messages",
                 container=True,
                 scale=1,
-                min_width=800,
-                value=[]
+                min_width=800,  # Ensure minimum width
+                value=[],  # Initialize with empty list
+                render=True
             )
             msg = gr.Textbox(
                 placeholder="Describe what you want to do next in the story...",
@@ -252,42 +237,50 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             starter_btn4 = gr.Button("Starter 4")
             starter_buttons = [starter_btn1, starter_btn2, starter_btn3, starter_btn4]
 
+            # 1) We'll return a list of 4 dicts, each dict updating 'value' & 'visible'
             def update_starter_buttons(selected_genre):
                 """Update starter buttons with examples for the selected genre."""
                 examples = get_examples_for_genre(selected_genre)
                 results = []
                 for i in range(4):
                     if i < len(examples):
+                        # Return just the string value instead of a dict
                         results.append(examples[i])
                     else:
-                        results.append("")
-                return tuple(results)
+                        results.append("")  # Empty string for hidden buttons
+                return tuple(results)  # Return tuple of strings
+
+            # 2) Initialize them with "fantasy" so they don't stay "Starter X" on page load
+            # We'll just call the function and store the results in a variable, then apply them in a .load() event
+            initial_button_data = update_starter_buttons("fantasy")  # returns 4 dicts
 
-            # Initialize with default genre
-            initial_button_data = update_starter_buttons("fantasy")
+            # 3) We'll define a "pick_starter" function that sets msg to the chosen text
+            def pick_starter(starter_text, chat_history, selected_genre, memory_flag):
+                # Putting 'starter_text' into the msg
+                return starter_text
 
-            # Connect each starter button
+            # 4) Connect each starter button:
             for starter_button in starter_buttons:
-                # First, update chatbot with starter message
                 starter_button.click(
-                    fn=lambda x: [(str(x), "")],
+                    # Format initial message correctly as ChatMessage
+                    fn=lambda x: [{"role": "user", "content": str(x)}],
                     inputs=[starter_button],
                     outputs=[chatbot],
                     queue=False
                 ).success(
-                    # Then process the message with respond function
-                    fn=respond,
-                    inputs=[
-                        starter_button,
-                        chatbot,
-                        genre,
-                        full_memory
-                    ],
+                    # Then process with properly formatted message
+                    fn=lambda x, h, g, m: respond(
+                        message=x["content"],  # Extract content from ChatMessage
+                        chat_history=h if h else [],
+                        genre=g,
+                        use_full_memory=m
+                    ),
+                    inputs=[starter_button, chatbot, genre, full_memory],
                     outputs=chatbot,
                     queue=True
                 )
 
-            # Update buttons when genre changes
+            # 5) Dynamically update the 4 buttons if the user changes the genre
             genre.change(
                 fn=update_starter_buttons,
                 inputs=[genre],
@@ -316,13 +309,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         queue=False
     )
 
-    # Load initial button data
-    demo.load(
-        fn=lambda: initial_button_data,
-        outputs=starter_buttons,
-        queue=False
-    )
+    # 6) Finally, run a "load" event to apply initial_button_data to the 4 button outputs on page load
+    def load_initial_buttons():
+        # Just return our precomputed tuple of 4 dicts
+        return initial_button_data
+
+    demo.load(fn=load_initial_buttons, outputs=starter_buttons, queue=False)
 
 # Run the app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)