Update app.py

app.py

@@ -1320,11 +1320,10 @@ def remove_loading_animations(chat_history):
         "loading-animation" in str(msg.get("content", ""))
     )]
 
-
 def add_user_message(message, chat_history, conversation_state):
     """
     Add user message with proper state management.
-    ✅ FIXED: Returns updated states to Gradio components.
+    ✅ FIXED: Uses Gradio state instead of overwriting with global manager.
     """
     callback_start = log_step("add_user_message")
     
@@ -1332,36 +1331,26 @@ def add_user_message(message, chat_history, conversation_state):
         log_step("add_user_message", callback_start)
         return "", chat_history, conversation_state
     
-    # Get current state from global manager
-    current_state = global_state_manager.get_conversation_state()
-    chat_history = current_state['chat_history']
-    conversation_state = current_state['conversation_state']
-    
     # Add to both states
     conversation_state.append({"role": "user", "content": message})
     chat_history.append({"role": "user", "content": message})
     
-    # Update global state
+    # Update global state for persistence
     global_state_manager.update_conversation_state(chat_history, conversation_state)
     
     log_step("add_user_message", callback_start)
     
-    # ✅ CRITICAL: Return updated states to Gradio
+    # Return updated states to Gradio
     return "", chat_history, conversation_state
 
 
 def add_loading_animation(chat_history, conversation_state):
     """
     Add loading animation with proper state management.
-    ✅ FIXED: Returns updated states to Gradio components.
+    ✅ FIXED: Uses Gradio state instead of overwriting.
     """
     callback_start = log_step("add_loading_animation")
     
-    # Get current state from global manager
-    current_state = global_state_manager.get_conversation_state()
-    chat_history = current_state['chat_history']
-    conversation_state = current_state['conversation_state']
-    
     if not conversation_state:
         log_step("add_loading_animation", callback_start)
         return chat_history, conversation_state
@@ -1378,26 +1367,29 @@ def add_loading_animation(chat_history, conversation_state):
     
     chat_history.append({"role": "assistant", "content": loading_html})
     
-    # Update global state
+    # Update global state for persistence
     global_state_manager.update_conversation_state(chat_history, conversation_state)
     
     log_step("add_loading_animation", callback_start)
     
-    # ✅ CRITICAL: Return updated states to Gradio
+    # Return updated states to Gradio
     return chat_history, conversation_state
 
 
 def generate_response(chat_history, conversation_state):
     """
     Generate response using orchestration with proper streaming.
-    ✅ FIXED: Loading animation stays until first chunk, then streams properly.
+    Uses Gradio state instead of overwriting, only pulls from global on error.
     """
     callback_start = log_step("generate_response")
     
-    # Get fresh state from global manager
-    current_state = global_state_manager.get_conversation_state()
-    chat_history = current_state['chat_history']
-    conversation_state = current_state['conversation_state']
+    # Use the state passed by Gradio
+    # Only pull from global manager if state is empty (error recovery)
+    if not conversation_state:
+        logger.warning("Empty conversation_state in generate_response, attempting recovery from global state")
+        current_state = global_state_manager.get_conversation_state()
+        chat_history = current_state['chat_history']
+        conversation_state = current_state['conversation_state']
     
     if not conversation_state:
         log_step("generate_response", callback_start)
@@ -1415,9 +1407,7 @@ def generate_response(chat_history, conversation_state):
         return chat_history, conversation_state
     
     try:
-        # ✅ DON'T remove loading animation yet - let it show during orchestration
-        
-        # Call orchestration (this takes time)
+        # Call orchestration
         orch_start = log_step("orchestrate_turn call")
         raw_response = orchestrate_turn(last_user_message)
         log_step("orchestrate_turn call", orch_start)
@@ -1425,7 +1415,7 @@ def generate_response(chat_history, conversation_state):
         # Stream the processed response
         first_chunk = True
         for chunk in post_processor.process_and_stream_response(raw_response, last_user_message):
-            # ✅ Remove loading animation on FIRST chunk only
+            # Remove loading animation on FIRST chunk only
             if first_chunk:
                 chat_history = remove_loading_animations(chat_history)
                 first_chunk = False
@@ -1436,7 +1426,7 @@ def generate_response(chat_history, conversation_state):
             else:
                 chat_history.append({"role": "assistant", "content": chunk})
             
-            # ✅ Yield to update UI during streaming
+            # Yield to update UI during streaming
             yield chat_history, conversation_state
         
         # Add final response to conversation state
@@ -1446,7 +1436,7 @@ def generate_response(chat_history, conversation_state):
         # Update global state with final conversation
         global_state_manager.update_conversation_state(chat_history, conversation_state)
         
-        # ✅ Final yield with complete states
+        # Final yield with complete states
         yield chat_history, conversation_state
         
     except Exception as e:
@@ -1487,10 +1477,12 @@ def load_conversation_state():
     current_state = global_state_manager.get_conversation_state()
     log_step("load_conversation_state", callback_start)
     
-    # ✅ Extract and return both states
+    # Extract and return both states
     return current_state['chat_history'], current_state['conversation_state']
 
 
+
+
 # ============================================================================
 # MULTI-PAGE INTERFACE
 # ============================================================================