Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import json
 import re
 import logging
-import threading
+# import threading # No longer needed for the learning task
 from datetime import datetime
 from dotenv import load_dotenv
 import gradio as gr
@@ -178,13 +178,15 @@ def process_user_interaction_gradio(user_input: str, provider_name: str, model_d
     logger.info(f"PUI_GRADIO [{request_id}]: Finished. Total: {time.time() - process_start_time:.2f}s. Resp len: {len(final_bot_text)}")
     yield "final_response_and_insights", {"response": final_bot_text, "insights_used": parsed_initial_insights_list}
 
-def deferred_learning_and_memory_task(user_input: str, bot_response: str, provider: str, model_disp_name: str, insights_reflected: list[dict], api_key_override: str = None):
-    start_time, task_id = time.time(), os.urandom(4).hex()
-    logger.info(f"DEFERRED [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
+# Renamed from deferred_learning_and_memory_task
+def perform_post_interaction_learning(user_input: str, bot_response: str, provider: str, model_disp_name: str, insights_reflected: list[dict], api_key_override: str = None):
+    task_id = os.urandom(4).hex() # Keep task_id for logging clarity
+    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
+    learning_start_time = time.time()
     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) # This updates in-memory lists quickly
+        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Metrics: {metrics}")
+        add_memory_entry(user_input, metrics, bot_response) 
         
         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."
@@ -239,7 +241,7 @@ Combine all findings into a single, valid XML structure as specified in the syst
             i_p, i_id = insight_env_model.split('/', 1)
             i_d_n = next((dn for dn, mid in MODELS_BY_PROVIDER.get(i_p.lower(), {}).get("models", {}).items() if mid == i_id), None)
             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)")
+        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Generating insights with {insight_prov}/{insight_model_disp} (expecting XML)")
         
         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()
         
@@ -269,57 +271,57 @@ Combine all findings into a single, valid XML structure as specified in the syst
                                 "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}")
+                            logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Skipped XML operation due to missing action or insight text. Action: {action}, Insight: {insight_text}")
                 else:
-                    logger.error(f"DEFERRED [{task_id}]: XML root tag is not <operations_list>. Found: {root.tag}. XML content:\n{xml_content_str}")
+                    logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML root tag is not <operations_list>. Found: {root.tag}. XML content:\n{xml_content_str}")
             except ET.ParseError as e:
-                logger.error(f"DEFERRED [{task_id}]: XML parsing error: {e}. XML content that failed:\n{xml_content_str}")
+                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML parsing error: {e}. XML content that failed:\n{xml_content_str}")
             except Exception as e_xml_proc:
-                logger.error(f"DEFERRED [{task_id}]: Error processing parsed XML: {e_xml_proc}. XML content:\n{xml_content_str}")
+                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: Error processing parsed XML: {e_xml_proc}. XML content:\n{xml_content_str}")
         else:
-            logger.info(f"DEFERRED [{task_id}]: No <operations_list> XML structure found in LLM output. Full raw output:\n{raw_ops_xml_full}")
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No <operations_list> XML structure found in LLM output. Full raw output:\n{raw_ops_xml_full}")
 
         if ops_data_list:
-            logger.info(f"DEFERRED [{task_id}]: LLM provided {len(ops_data_list)} insight ops from XML.")
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: LLM provided {len(ops_data_list)} insight ops from XML.")
             for op_idx, op_data in enumerate(ops_data_list):
                 action = op_data["action"]
                 insight_text = op_data["insight"]
                 old_insight = op_data["old_insight_to_replace"]
                 
                 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]}...'")
+                    logger.warning(f"POST_INTERACTION_LEARNING [{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) # This updates in-memory lists quickly
+                    success, status_msg = add_rule_entry(insight_text)
                     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}")
+                    else: logger.warning(f"POST_INTERACTION_LEARNING [{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) # This updates in-memory lists quickly
+                            remove_success = remove_rule_entry(old_insight)
                             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.")
+                                 logger.warning(f"POST_INTERACTION_LEARNING [{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.")
+                             logger.info(f"POST_INTERACTION_LEARNING [{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) # This updates in-memory lists quickly
+                    success, status_msg = add_rule_entry(insight_text) 
                     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: logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Failed to add/update rule '{insight_text[:50]}...'. Status: {status_msg}")
                 else:
-                    logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}' from XML.")
+                    logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}' from XML.")
 
-            logger.info(f"DEFERRED [{task_id}]: Processed {processed_count} insight ops out of {len(ops_data_list)} received from XML.")
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Processed {processed_count} insight ops out of {len(ops_data_list)} received from XML.")
         else:
-            logger.info(f"DEFERRED [{task_id}]: No valid insight operations derived from LLM's XML output.")
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No valid insight operations derived from LLM's XML output.")
             
-    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")
+    except Exception as e: logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: CRITICAL ERROR in learning task: {e}", exc_info=True)
+    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: END. Total: {time.time() - learning_start_time:.2f}s")
+    # This function now runs synchronously, so its completion means data is updated.
 
 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) 
@@ -355,7 +357,7 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
                     updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg)
             elif upd_type == "final_response_and_insights":
                 final_bot_resp_acc, insights_used_parsed = upd_data["response"], upd_data["insights_used"]
-                status_txt = "Response complete."
+                status_txt = "Response generated. Processing learning..." # Status update
                 if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
                 if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: 
                     updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg or "(No text)")
@@ -370,14 +372,16 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
                 else: 
                     def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
                 
-                insights_md_content = "### Insights Considered:\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
+                insights_md_content = "### Insights Considered (Pre-Response):\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
                 def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True if insights_used_parsed else False)
 
+            # Yield intermediate updates for response streaming
             yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
-            if upd_type == "final_response_and_insights": break
+            if upd_type == "final_response_and_insights": break # Exit loop after main response processing
+    
     except Exception as e:
-        logger.error(f"Chat handler error: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
-        error_message_for_chat = f"Sorry, an error occurred: {str(e)[:100]}"
+        logger.error(f"Chat handler error during main processing: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
+        error_message_for_chat = f"Sorry, an error occurred during response generation: {str(e)[:100]}"
         if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: 
             updated_gr_hist[-1] = (user_msg_txt, error_message_for_chat)
         else: 
@@ -388,6 +392,7 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
         yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
         return
 
+    # --- Post-Interaction Learning (Synchronous within this handler) ---
     if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
         current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
         hist_len_check = MAX_HISTORY_TURNS * 2
@@ -395,32 +400,37 @@ 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):]
         
-        # 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."
-        if final_bot_resp_acc.startswith("Error:"): 
-            status_txt = final_bot_resp_acc 
-            if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
-                 updated_gr_hist[-1] = (user_msg_txt, final_bot_resp_acc)
-            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
+        # Call learning function synchronously
+        status_txt = "<i>[Performing post-interaction learning...]</i>"
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn) # Update status
+
+        try:
+            perform_post_interaction_learning(
+                user_input=user_msg_txt, 
+                bot_response=final_bot_resp_acc, 
+                provider=sel_prov_name, 
+                model_disp_name=sel_model_disp_name, 
+                insights_reflected=insights_used_parsed, 
+                api_key_override=ui_api_key.strip() if ui_api_key else None
+            )
+            status_txt = "Response & Learning Complete."
+        except Exception as e_learn:
+            logger.error(f"Error during post-interaction learning: {e_learn}", exc_info=True)
+            status_txt = "Response complete. Error during learning."
+
+    elif final_bot_resp_acc.startswith("Error:"):
+        status_txt = final_bot_resp_acc 
+        if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
+                updated_gr_hist[-1] = (user_msg_txt, final_bot_resp_acc)
+        def_fmt_out_txt = gr.Textbox(value=final_bot_resp_acc, interactive=True)
+        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
+    else:
+        status_txt = "Processing finished; no valid response or error occurred during main phase."
+
+
+    # Final yield for this handler. The .then() clauses will execute after this.
     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}")
@@ -718,23 +728,14 @@ 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.
+    # The .then() calls will execute after handle_gradio_chat_submit fully completes (all yields).
+    # Since perform_post_interaction_learning is now synchronous within handle_gradio_chat_submit,
+    # the data will be updated before these .then() calls run.
     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)
@@ -788,7 +789,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 # It's a quick operation
+        show_progress=False 
     ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False)
 
     # Memories Management events
@@ -804,7 +805,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 # Quick op
+        show_progress=False 
     ).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():
@@ -840,7 +841,7 @@ with gr.Blocks(
 
 
 if __name__ == "__main__":
-    logger.info(f"Starting Gradio AI Research Mega Agent (v6.3 - Delayed UI Refresh Attempt, Memory: {MEMORY_STORAGE_BACKEND})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v6.4 - Synchronous Learning for UI Refresh, 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"