Update app.py

app.py

@@ -81,7 +81,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
@@ -133,14 +133,12 @@ def create_story_summary(chat_history):
     }
     return summary_instruction
 
-def format_history_for_api(chat_history):
-    """Format chat history for the API - fixed to ensure proper message format"""
-    formatted_messages = []
-    for user_msg, bot_msg in chat_history:
-        formatted_messages.append({"role": "user", "content": str(user_msg)})
-        formatted_messages.append({"role": "assistant", "content": str(bot_msg)})
-    return formatted_messages
+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) -> Tuple[str, List[Tuple[str, str]]]:
     """Generate a response based on the current message and conversation history."""
     if not message.strip():
@@ -151,23 +149,22 @@ def respond(message: str, chat_history: List[Tuple[str, str]], genre: Optional[s
         api_messages = [{"role": "system", "content": get_enhanced_system_prompt(genre)}]
         logging.debug(f"System Message: {api_messages[0]}")
         
-        # Add chat history - fixed to ensure proper formatting
+        # Add chat history
         if chat_history and use_full_memory:
-            history_messages = format_history_for_api(chat_history[-MEMORY_WINDOW:])
-            api_messages.extend(history_messages)
-            logging.debug(f"Added {len(history_messages)} history messages")
+            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:]}")
         
         # Add current message
         api_messages.append({"role": "user", "content": str(message)})
-        
-        # Log api message structure for debugging
-        logging.debug(f"API messages structure: {[msg['role'] for msg in api_messages]}")
-        logging.debug(f"Final user message: {message}")
+        logging.debug(f"Final Message List: {api_messages}")
 
         # Make API call
         logging.debug("Making API call...")
         response = client.chat_completion(
-            model="HuggingFaceH4/zephyr-7b-beta",  # Explicitly specify model
             messages=api_messages,
             max_tokens=MAX_TOKENS,
             temperature=TEMPERATURE,
@@ -200,27 +197,8 @@ def save_story(chat_history):
     
     return story_text
 
-def test_api_connection():
-    """Test API connection and return result"""
-    try:
-        response = client.chat_completion(
-            model="HuggingFaceH4/zephyr-7b-beta",
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Say hello!"}
-            ],
-            max_tokens=50
-        )
-        return f"✅ API Connection Test: Success! Response: {response.choices[0].message.content[:30]}..."
-    except Exception as e:
-        return f"❌ API Connection Failed: {str(e)}"
-
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🔮 Interactive Story Time")
-    
-    # Add API status indicator
-    api_status = gr.Markdown("Checking API connection...", visible=True)
-    
     gr.Markdown("Create a completely unique literary world, one choice at a time. Dare to explore the unknown.")
     
     with gr.Row():
@@ -230,12 +208,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 height=500,
                 bubble_full_width=True,
                 show_copy_button=True,
-                avatar_images=(None, "🧙"),
+                avatar_images=("👤", "🧙‍♂️"),  # Updated avatars
+                type="chat",
                 container=True,
                 scale=1,
                 min_width=800,
                 value=[],
-                render=True
+                render=True,
+                layout="bubble"
             )
             msg = gr.Textbox(
                 placeholder="Describe what you want to do next in the story...",
@@ -261,10 +241,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 info="When enabled, the AI tries to remember the entire story. If disabled, only the last few exchanges are used."
             )
             
-            # Add debug button
-            debug_btn = gr.Button("Test API Connection", variant="secondary")
-            debug_output = gr.Markdown(visible=True)
-            
             gr.Markdown("## Story Starters")
             
             # Create four placeholder buttons for story starters
@@ -285,34 +261,32 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         results.append("")
                 return tuple(results)
 
-            # Debug button handler
-            debug_btn.click(
-                fn=test_api_connection,
-                inputs=None,
-                outputs=debug_output
-            )
-
-            # Function to handle starter button clicks
+            # New direct handler for starter clicks
             def use_starter(starter_text, history, selected_genre, memory_flag):
-                """Direct handler for starter buttons"""
+                """Handle starter button clicks with proper message formatting"""
                 if not starter_text:
                     return "", history
                 
-                # Simply call the respond function with the starter text
-                empty_msg, new_history = respond(
-                    message=starter_text,
-                    chat_history=history,
-                    genre=selected_genre,
-                    use_full_memory=memory_flag
-                )
-                return empty_msg, new_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 "", history + [(starter_text, error_msg)]
 
-            # Connect each starter button
+            # 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],
+                    outputs=[msg, chatbot],  # Now returning both message and chat history
                     queue=True
                 )
 
@@ -323,30 +297,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 outputs=starter_buttons
             )
 
-    # Handler for user input - now using the fixed respond function
-    def submit_message(message, history, selected_genre, memory_flag):
-        """Handle user message submission"""
-        if not message.strip():
-            return "", history
-        
-        empty_msg, new_history = respond(
-            message=message,
-            chat_history=history,
-            genre=selected_genre,
-            use_full_memory=memory_flag
-        )
-        return empty_msg, new_history
-
+    # Handler for user input
     msg.submit(
-        fn=submit_message,
+        fn=respond,
         inputs=[msg, chatbot, genre, full_memory],
-        outputs=[msg, chatbot]
+        outputs=[msg, chatbot]  # Now returning both message and chat history
     )
-    
     submit.click(
-        fn=submit_message,
+        fn=respond,
         inputs=[msg, chatbot, genre, full_memory],
-        outputs=[msg, chatbot]
+        outputs=[msg, chatbot]  # Now returning both message and chat history
     )
     
     # Clear the chatbot for a new adventure
@@ -380,14 +340,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=starter_buttons,
         queue=False
     )
-    
-    # Also test API connection on load
-    demo.load(
-        fn=test_api_connection,
-        outputs=api_status,
-        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)