Update app.py

app.py

@@ -984,46 +984,88 @@ def warmup_agent():
     except Exception as e:
         logger.error(f"Warmup failed: {e}")
 
-def respond_and_update(message, history):
-    """Handle user input and generate streaming response"""
+def add_user_message(message, chat_history, conversation_state):
+    """Add user message to state and display immediately"""
     if not message.strip():
-        return history, ""
+        return "", chat_history, conversation_state
     
-    # Add user message to history
-    history.append({"role": "user", "content": message})
+    # Add to conversation state
+    conversation_state.append({"role": "user", "content": message})
     
-    # yield to show user message
-    yield history, ""
+    # Update chat display
+    chat_history.append({"role": "user", "content": message})
     
-    # Add empty assistant message that will be updated
-    history.append({"role": "assistant", "content": ""})
+    return "", chat_history, conversation_state
+
+def add_thinking_indicator(chat_history, conversation_state):
+    """Add thinking indicator to chat display"""
+    if not conversation_state:
+        return chat_history, conversation_state
+    
+    # Get the last user message for context-aware thinking indicator
+    last_message = conversation_state[-1]["content"] if conversation_state else ""
+    thinking_html = get_thinking_indicator_by_query(last_message)
+    
+    # Add thinking indicator to chat display (not permanent state)
+    chat_history.append({"role": "assistant", "content": thinking_html})
+    
+    return chat_history, conversation_state
+
+def generate_response(chat_history, conversation_state):
+    """Generate streaming response from the agent"""
+    if not conversation_state:
+        return chat_history, conversation_state
+    
+    # Get the last user message
+    last_user_message = ""
+    for msg in reversed(conversation_state):
+        if msg["role"] == "user":
+            last_user_message = msg["content"]
+            break
     
-    #  Yield again to show empty assistant bubble
-    thinking_html = get_thinking_indicator_by_query(message)
-    history.append({"role": "assistant", "content": thinking_html})
-    yield history, ""
+    if not last_user_message:
+        return chat_history, conversation_state
     
     try:
-        # Generate streaming response
+        # Stream the response
         full_response = ""
-        for chunk in agent.stream_query(message):
+        
+        for chunk in agent.stream_query(last_user_message):
             full_response = chunk
-            # Update the last message in history
-            history[-1]["content"] = full_response
-            yield history, ""
+            
+            # Update the last message in chat display (replace thinking indicator)
+            if chat_history and chat_history[-1]["role"] == "assistant":
+                chat_history[-1]["content"] = full_response
+            else:
+                chat_history.append({"role": "assistant", "content": full_response})
+            
+            yield chat_history, conversation_state
+        
+        # Add final response to permanent conversation state
+        conversation_state.append({"role": "assistant", "content": full_response})
+        yield chat_history, conversation_state
         
     except Exception as e:
-        logger.error(f"Error in respond_and_update: {e}")
-        history[-1]["content"] = f"I encountered an error: {str(e)}"
-        yield history, ""
+        logger.error(f"Error in generate_response: {e}")
+        error_msg = f"I encountered an error: {str(e)}"
+        
+        # Update display
+        if chat_history and chat_history[-1]["role"] == "assistant":
+            chat_history[-1]["content"] = error_msg
+        else:
+            chat_history.append({"role": "assistant", "content": error_msg})
+        
+        # Add to permanent state
+        conversation_state.append({"role": "assistant", "content": error_msg})
+        yield chat_history, conversation_state
 
-def clear_chat():
-    """Clear the chat history"""
-    return [], ""
+def reset_conversation():
+    """Reset both chat display and conversation state"""
+    return [], []
     
 # --- UI: Interface Creation ---
 def create_interface():
-    """Creates and configures the complete Gradio interface."""
+    """Creates and configures the complete Gradio interface with proper state management."""
     start_create_interface_time = time.perf_counter()
     current_time = datetime.now()
     
@@ -1046,7 +1088,9 @@ def create_interface():
         # Add head content
         gr.HTML(html_head_content)
         gr.HTML(force_light_mode)
-
+        
+        # State management - this is the key addition
+        conversation_state = gr.State([])  # Persistent conversation memory
         
         with gr.Column(elem_classes=["main-container"]):
             # Title Section
@@ -1064,9 +1108,10 @@ def create_interface():
                     elem_id="main-chatbot",
                     scale=1,
                     height="70vh",
+                    value=[],  # Initialize with empty list
                     latex_delimiters=[
-                        {"left": "$$", "right": "$$", "display": True},  # For centered display math
-                        {"left": "$", "right": "$", "display": False},    # For inline math
+                        {"left": "$$", "right": "$$", "display": True},
+                        {"left": "$", "right": "$", "display": False},
                     ]
                 )
             
@@ -1085,12 +1130,50 @@ def create_interface():
                     send = gr.Button("Send", elem_classes=["send-button"], size="sm")
                     clear = gr.Button("Clear", elem_classes=["clear-button"], size="sm")
             
-            # Event handlers
-            msg.submit(respond_and_update, [msg, chatbot], [chatbot, msg])
-            send.click(respond_and_update, [msg, chatbot], [chatbot, msg])
-            clear.click(clear_chat, outputs=[chatbot, msg])
+            # event chaining with state management
+            submit_event = msg.submit(
+                add_user_message,
+                inputs=[msg, chatbot, conversation_state],
+                outputs=[msg, chatbot, conversation_state],
+                show_progress="hidden"
+            ).then(
+                add_thinking_indicator,
+                inputs=[chatbot, conversation_state],
+                outputs=[chatbot, conversation_state],
+                show_progress="hidden"
+            ).then(
+                generate_response,
+                inputs=[chatbot, conversation_state],
+                outputs=[chatbot, conversation_state],
+                show_progress="hidden"
+            )
+            
+            send_event = send.click(
+                add_user_message,
+                inputs=[msg, chatbot, conversation_state],
+                outputs=[msg, chatbot, conversation_state],
+                show_progress="hidden"
+            ).then(
+                add_thinking_indicator,
+                inputs=[chatbot, conversation_state],
+                outputs=[chatbot, conversation_state],
+                show_progress="hidden"
+            ).then(
+                generate_response,
+                inputs=[chatbot, conversation_state],
+                outputs=[chatbot, conversation_state],
+                show_progress="hidden"
+            )
+            
+            # Clear button
+            clear.click(
+                reset_conversation,
+                inputs=None,
+                outputs=[chatbot, conversation_state],
+                show_progress="hidden"
+            )
             
-            # Apply CSS at the very end
+            # Apply CSS
             gr.HTML(f'<style>{custom_css}</style>')
     
     end_create_interface_time = time.perf_counter()