Update app.py

app.py

@@ -169,7 +169,7 @@ def process_user_interaction_gradio(user_input: str, provider_name: str, model_d
     logger.debug(f"PUI_GRADIO [{request_id}]: Final LLM User Prompt Start: {final_user_prompt_content_str[:200]}...")
     streamed_response, time_before_llm = "", time.time()
     try:
-        for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_llm_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=2500): # Increased max_tokens for main response
+        for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_llm_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=2500):
             if isinstance(chunk, str) and chunk.startswith("Error:"): streamed_response += f"\n{chunk}\n"; yield "response_chunk", f"\n{chunk}\n"; break
             streamed_response += chunk; yield "response_chunk", chunk
     except Exception as e: streamed_response += f"\n\n(Error: {str(e)[:150]})"; yield "response_chunk", f"\n\n(Error: {str(e)[:150]})"
@@ -184,7 +184,8 @@ def deferred_learning_and_memory_task(user_input: str, bot_response: str, provid
     try:
         metrics = generate_interaction_metrics(user_input, bot_response, provider, model_disp_name, api_key_override)
         logger.info(f"DEFERRED [{task_id}]: Metrics: {metrics}")
-        add_memory_entry(user_input, metrics, bot_response)
+        add_memory_entry(user_input, metrics, bot_response) # This updates in-memory lists quickly
+        
         summary = f"User:\"{user_input}\"\nAI:\"{bot_response}\"\nMetrics(takeaway):{metrics.get('takeaway','N/A')},Success:{metrics.get('response_success_score','N/A')}"
         existing_rules_ctx = "\n".join([f"- \"{r}\"" for r in retrieve_rules_semantic(f"{summary}\n{user_input}", k=10)]) or "No existing rules context."
         
@@ -240,12 +241,10 @@ Combine all findings into a single, valid XML structure as specified in the syst
             if i_d_n: insight_prov, insight_model_disp = i_p, i_d_n
         logger.info(f"DEFERRED [{task_id}]: Generating insights with {insight_prov}/{insight_model_disp} (expecting XML)")
         
-        # Increase max_tokens significantly for XML, as it can be more verbose than JSON
-        raw_ops_xml_full = "".join(list(call_model_stream(provider=insight_prov, model_display_name=insight_model_disp, messages=insight_msgs, api_key_override=api_key_override, temperature=0.0, max_tokens=3500))).strip() # Increased max_tokens
+        raw_ops_xml_full = "".join(list(call_model_stream(provider=insight_prov, model_display_name=insight_model_disp, messages=insight_msgs, api_key_override=api_key_override, temperature=0.0, max_tokens=3500))).strip()
         
         ops_data_list, processed_count = [], 0
         
-        # Try to find XML block, possibly within markdown
         xml_match = re.search(r"```xml\s*(<operations_list>.*</operations_list>)\s*```", raw_ops_xml_full, re.DOTALL | re.IGNORECASE) or \
                     re.search(r"(<operations_list>.*</operations_list>)", raw_ops_xml_full, re.DOTALL | re.IGNORECASE)
 
@@ -267,7 +266,7 @@ Combine all findings into a single, valid XML structure as specified in the syst
                             ops_data_list.append({
                                 "action": action,
                                 "insight": insight_text,
-                                "old_insight_to_replace": old_insight_text # Can be None
+                                "old_insight_to_replace": old_insight_text 
                             })
                         else:
                             logger.warning(f"DEFERRED [{task_id}]: Skipped XML operation due to missing action or insight text. Action: {action}, Insight: {insight_text}")
@@ -287,25 +286,24 @@ Combine all findings into a single, valid XML structure as specified in the syst
                 insight_text = op_data["insight"]
                 old_insight = op_data["old_insight_to_replace"]
                 
-                # Validation of insight_text format
                 if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I|re.DOTALL): 
                     logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to invalid insight_text format from XML: '{insight_text[:100]}...'")
                     continue
 
                 if action == "add":
-                    success, status_msg = add_rule_entry(insight_text)
+                    success, status_msg = add_rule_entry(insight_text) # This updates in-memory lists quickly
                     if success: processed_count +=1
                     else: logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (add from XML): Failed to add rule '{insight_text[:50]}...'. Status: {status_msg}")
                 elif action == "update":
                     if old_insight: 
                         if old_insight != insight_text: 
-                            remove_success = remove_rule_entry(old_insight)
+                            remove_success = remove_rule_entry(old_insight) # This updates in-memory lists quickly
                             if not remove_success:
                                  logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (update from XML): Failed to remove old rule '{old_insight[:50]}...' before adding new.")
                         else: 
                              logger.info(f"DEFERRED [{task_id}]: Op {op_idx} (update from XML): Old insight is identical to new insight. Skipping removal.")
                     
-                    success, status_msg = add_rule_entry(insight_text) 
+                    success, status_msg = add_rule_entry(insight_text) # This updates in-memory lists quickly
                     if success: processed_count +=1
                     else: logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (update from XML): Failed to add/update rule '{insight_text[:50]}...'. Status: {status_msg}")
                 else:
@@ -318,15 +316,10 @@ Combine all findings into a single, valid XML structure as specified in the syst
     except Exception as e: logger.error(f"DEFERRED [{task_id}]: CRITICAL ERROR in deferred task: {e}", exc_info=True)
     logger.info(f"DEFERRED [{task_id}]: END. Total: {time.time() - start_time:.2f}s")
 
-# --- handle_gradio_chat_submit, UI functions, Gradio Layout ... ---
-# (The rest of the file: handle_gradio_chat_submit, all ui_... functions, and the
-# `with gr.Blocks(...) as demo:` section, and `if __name__ == "__main__":`
-# remain THE SAME as the previous fully correct version (v6.1) you provided.
-# No changes are needed in those sections for this XML switch.)
-
 def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_name: str, sel_model_disp_name: str, ui_api_key: str|None, cust_sys_prompt: str):
     global current_chat_session_history
     cleared_input, updated_gr_hist, status_txt = "", list(gr_hist_list), "Initializing..."
+    # Ensure these are initialized with their correct Gradio update types
     def_detect_out_md = gr.Markdown(visible=False) 
     def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True) 
     def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False) 
@@ -402,7 +395,9 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
         if len(current_chat_session_history) > hist_len_check:
             current_chat_session_history = ([current_chat_session_history[0]] if current_chat_session_history[0]["role"] == "system" else []) + current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]
         
-        threading.Thread(target=deferred_learning_and_memory_task, args=(user_msg_txt, final_bot_resp_acc, sel_prov_name, sel_model_disp_name, insights_used_parsed, ui_api_key.strip() if ui_api_key else None), daemon=True).start()
+        # Start the deferred task
+        deferred_task_thread = threading.Thread(target=deferred_learning_and_memory_task, args=(user_msg_txt, final_bot_resp_acc, sel_prov_name, sel_model_disp_name, insights_used_parsed, ui_api_key.strip() if ui_api_key else None), daemon=True)
+        deferred_task_thread.start()
         status_txt = "Response complete. Background learning initiated."
     else: 
         status_txt = "Processing finished; no valid response or error occurred."
@@ -413,8 +408,19 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
             def_fmt_out_txt = gr.Textbox(value=final_bot_resp_acc, interactive=True)
             def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
 
+    # Final yield for the main chat processing
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
     
+    # After the main response is fully yielded, and the deferred task thread is started,
+    # we can add a small delay here before the .then() clauses trigger the UI refresh.
+    # This gives the deferred task's initial, fast in-memory updates (like add_memory_entry)
+    # a better chance to complete before the UI tries to read that data.
+    # The LLM-based rule generation within the deferred task will still take longer.
+    if 'deferred_task_thread' in locals() and deferred_task_thread.is_alive():
+        # Wait a very short period for synchronous parts of deferred task
+        # This is a heuristic. The value might need tuning.
+        time.sleep(0.2) # e.g., 200 milliseconds
+
     if temp_dl_file_path and os.path.exists(temp_dl_file_path):
         try: os.unlink(temp_dl_file_path)
         except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path}: {e_unlink}")
@@ -456,7 +462,6 @@ def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
             line = line.strip()
             if line:
                 try:
-                    # Expect each line to be a JSON string containing the rule text itself
                     rule_text_in_json_string = json.loads(line) 
                     if isinstance(rule_text_in_json_string, str):
                         potential_rules.append(rule_text_in_json_string)
@@ -713,15 +718,28 @@ with gr.Blocks(
     chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
     chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn]
     
+    # Define a dummy function to introduce a delay if needed for UI refresh.
+    # For now, we will try without an explicit Gradio-level delay here,
+    # relying on the small time.sleep in handle_gradio_chat_submit's final part.
+    # def delayed_refresh_trigger():
+    # time.sleep(0.2) # Small delay
+    # return True # Or any value, just to trigger the next .then()
+
+    # chat_event_args is a dictionary containing common arguments for click/submit
     chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
     
+    # Setup events for chat submission
     send_btn_click_event = send_btn.click(**chat_event_args)
     user_msg_submit_event = user_msg_tb.submit(**chat_event_args)
 
+    # Chain UI refreshes to occur after the chat processing logic (including the small sleep)
+    # The `handle_gradio_chat_submit` now includes a small sleep before its final yield
+    # if a deferred task was started. This should be sufficient for the fast in-memory updates.
     for event in [send_btn_click_event, user_msg_submit_event]:
         event.then(fn=ui_refresh_rules_display_fn, inputs=None, outputs=rules_disp_ta, show_progress=False)
         event.then(fn=ui_refresh_memories_display_fn, inputs=None, outputs=mems_disp_json, show_progress=False)
 
+
     # Rules Management events
     dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn)
 
@@ -770,7 +788,7 @@ with gr.Blocks(
     clear_rules_btn.click(
         fn=lambda: ("All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules."),
         outputs=rules_stat_tb,
-        show_progress=False
+        show_progress=False # It's a quick operation
     ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False)
 
     # Memories Management events
@@ -786,7 +804,7 @@ with gr.Blocks(
     clear_mems_btn.click(
         fn=lambda: ("All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories."),
         outputs=mems_stat_tb,
-        show_progress=False
+        show_progress=False # Quick op
     ).then(fn=ui_refresh_memories_display_fn, outputs=mems_disp_json, show_progress=False)
 
     if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
@@ -822,7 +840,7 @@ with gr.Blocks(
 
 
 if __name__ == "__main__":
-    logger.info(f"Starting Gradio AI Research Mega Agent (v6.2 - XML Insights & Max Tokens, Memory: {MEMORY_STORAGE_BACKEND})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v6.3 - Delayed UI Refresh Attempt, Memory: {MEMORY_STORAGE_BACKEND})...")
     app_port = int(os.getenv("GRADIO_PORT", 7860))
     app_server = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
     app_debug = os.getenv("GRADIO_DEBUG", "False").lower() == "true"