Update app.py

app.py

@@ -43,7 +43,6 @@ DEFAULT_SYSTEM_PROMPT = os.getenv(
 logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={TOOL_DECISION_PROVIDER_ENV}, ToolDecisionModelID={TOOL_DECISION_MODEL_ID_ENV}, MemoryBackend={MEMORY_STORAGE_BACKEND}")
 
 # --- Helper Functions (format_insights_for_prompt, generate_interaction_metrics, etc.) ---
-# ENSURE ALL YOUR HELPER AND LOGIC FUNCTIONS ARE DEFINED HERE AS IN PREVIOUS VERSIONS
 def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str, list[dict]]:
     if not retrieved_insights_list:
         return "No specific guiding principles or learned insights retrieved.", []
@@ -276,66 +275,121 @@ Combine all findings into a single JSON list of operations. If there are multipl
 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..."
-    def_detect_out_md, def_fmt_out_txt = gr.Markdown("*Processing...*"), gr.Textbox("*Waiting...*")
+    def_detect_out_md = gr.Markdown(visible=False) # Initialize with visible=False
+    def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True) # Ensure interactive
     def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False) 
+    
     if not user_msg_txt.strip():
         status_txt = "Error: Empty message."
         updated_gr_hist.append((user_msg_txt or "(Empty)", status_txt))
-        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn); return
+        # Make sure to yield for all outputs in chat_outs
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
+        return
+
     updated_gr_hist.append((user_msg_txt, "<i>Thinking...</i>"))
+    # Initial yield for chat update
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
+    
     internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
     if len(internal_hist) > (MAX_HISTORY_TURNS * 2 + 1):
         if internal_hist[0]["role"] == "system" and len(internal_hist) > (MAX_HISTORY_TURNS * 2 + 1) : internal_hist = [internal_hist[0]] + internal_hist[-(MAX_HISTORY_TURNS * 2):]
         else: internal_hist = internal_hist[-(MAX_HISTORY_TURNS * 2):]
+    
     final_bot_resp_acc, insights_used_parsed = "", []
     temp_dl_file_path = None 
+    
     try:
         processor_gen = process_user_interaction_gradio(user_input=user_msg_txt, provider_name=sel_prov_name, model_display_name=sel_model_disp_name, chat_history_for_prompt=internal_hist, custom_system_prompt=cust_sys_prompt.strip() or None, ui_api_key_override=ui_api_key.strip() if ui_api_key else None)
         curr_bot_disp_msg = ""
         for upd_type, upd_data in processor_gen:
             if upd_type == "status":
                 status_txt = upd_data
-                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, f"{curr_bot_disp_msg} <i>{status_txt}</i>" if curr_bot_disp_msg else f"<i>{status_txt}</i>")
+                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: 
+                    updated_gr_hist[-1] = (user_msg_txt, f"{curr_bot_disp_msg} <i>{status_txt}</i>" if curr_bot_disp_msg else f"<i>{status_txt}</i>")
             elif upd_type == "response_chunk":
                 curr_bot_disp_msg += upd_data
-                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg)
+                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: 
+                    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."
                 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)")
-                def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg)
+                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)")
+                
+                def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg, interactive=True, show_copy_button=True) # Update
+                
                 if curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
                     with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
                         tmpfile.write(curr_bot_disp_msg)
                         temp_dl_file_path = tmpfile.name
                     def_dl_btn = gr.DownloadButton(value=temp_dl_file_path, visible=True, interactive=True) 
-                else: def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
-                insights_md = "### 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.*")
-                def_detect_out_md = gr.Markdown(insights_md)
+                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.*")
+                def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True) # Update with value and visible
+
             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
     except Exception as e:
         logger.error(f"Chat handler error: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
-        if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, status_txt)
-        else: updated_gr_hist.append((user_msg_txt, status_txt))
-        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn); return
+        error_message_for_chat = f"Sorry, an error occurred: {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: 
+            updated_gr_hist.append((user_msg_txt, error_message_for_chat))
+        def_fmt_out_txt = gr.Textbox(value=error_message_for_chat, interactive=True) # Update
+        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False) # Reset
+        def_detect_out_md = gr.Markdown(value="*Error processing request.*", visible=True) # Update
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
+        return
+
     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
         if current_chat_session_history and current_chat_session_history[0]["role"] == "system": hist_len_check +=1
         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()
         status_txt = "Response complete. Background learning initiated."
-    else: status_txt = "Processing finished; no response or error."
+    else: 
+        status_txt = "Processing finished; no valid response or error occurred."
+        if final_bot_resp_acc.startswith("Error:"): # If there was an error string in final_bot_resp_acc
+            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) # Update chatbot with error
+            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 this handler
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
+    
     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}")
 
+# --- UI Functions for Rules and Memories ---
 def ui_view_rules_action_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."
+
+def ui_download_rules_action_fn():
+    rules_content = "\n\n---\n\n".join(get_all_rules_cached())
+    if not rules_content.strip():
+        gr.Warning("No rules to download.")
+        # Return an update to clear any previous file and disable button
+        return gr.DownloadButton(value=None, interactive=False, label="No Rules") 
+    
+    try:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt", encoding='utf-8') as tmpfile:
+            tmpfile.write(rules_content)
+            # Gradio's DownloadButton click handler should return the filepath string
+            return tmpfile.name 
+    except Exception as e:
+        logger.error(f"Error creating rules download file: {e}")
+        gr.Error(f"Failed to prepare rules for download: {e}")
+        return gr.DownloadButton(value=None, interactive=False, label="Error")
+
 def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
     if not uploaded_file_obj: return "No file provided for rules upload."
     try:
@@ -359,7 +413,35 @@ def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
     msg = f"Rules Upload: Processed {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors/Invalid: {error_count}."
     logger.info(msg); return msg
 
-def ui_view_memories_action_fn(): return get_all_memories_cached() or []
+def ui_view_memories_action_fn(): return get_all_memories_cached() or [] # Returns list of dicts
+
+def ui_download_memories_action_fn():
+    memories = get_all_memories_cached() # list of dicts
+    if not memories:
+        gr.Warning("No memories to download.")
+        return gr.DownloadButton(value=None, interactive=False, label="No Memories")
+
+    jsonl_content = ""
+    for mem_dict in memories:
+        try:
+            jsonl_content += json.dumps(mem_dict) + "\n"
+        except Exception as e:
+            logger.error(f"Error serializing memory for download: {mem_dict}, Error: {e}")
+            # Skip problematic memory for download
+    
+    if not jsonl_content.strip():
+        gr.Warning("No valid memories to serialize for download.")
+        return gr.DownloadButton(value=None, interactive=False, label="No Data")
+
+    try:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".jsonl", encoding='utf-8') as tmpfile:
+            tmpfile.write(jsonl_content)
+            return tmpfile.name # Return filepath string
+    except Exception as e:
+        logger.error(f"Error creating memories download file: {e}")
+        gr.Error(f"Failed to prepare memories for download: {e}")
+        return gr.DownloadButton(value=None, interactive=False, label="Error")
+
 def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
     if not uploaded_file_obj: return "No file provided for memories upload."
     try:
@@ -378,7 +460,7 @@ def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
                 except: format_error_count += 1
     if not memory_objects_to_process and format_error_count == 0: return "No valid memory objects found."
     total_to_process = len(memory_objects_to_process)
-    if total_to_process == 0: return "No memory objects to process."
+    if total_to_process == 0: return "No memory objects to process (after parsing)."
     progress(0, desc="Starting memories upload...")
     for idx, mem_data in enumerate(memory_objects_to_process):
         if isinstance(mem_data, dict) and all(k in mem_data for k in ["user_input", "bot_response", "metrics"]):
@@ -390,49 +472,9 @@ def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
     msg = f"Memories Upload: Processed {total_to_process}. Added: {added_count}, Format Errors: {format_error_count}, Save Errors: {save_error_count}."
     logger.info(msg); return msg
 
-# --- UI Definition ---
-# Using the CSS inspired by the "AI Code & Space Generator" example
-custom_theme = gr.themes.Base(
-    primary_hue="teal",
-    secondary_hue="purple",
-    neutral_hue="zinc",
-    text_size="sm",
-    spacing_size="md",
-    radius_size="sm",
-    font=["System UI", "sans-serif"]
-)
-
-custom_css = """
-body { background: linear-gradient(to bottom right, #2c3e50, #34495e); color: #ecf0f1; margin:0; padding:0; font-family: 'System UI', sans-serif; overflow-x: hidden;}
-.gradio-container { background: transparent !important; padding: 0 !important; /* Remove padding from the outermost container */ }
-.main-interface-wrapper { max-width: 100%; padding: 1rem; box-sizing: border-box; min-height: 100vh; display: flex; flex-direction: column;} /* New wrapper */
-.gr-column {gap: 0px !important;}
-.gr-block.gr-group, .gr-tabs, .gr-accordion { background-color: rgba(44, 62, 80, 0.8) !important; border: 1px solid rgba(189, 195, 199, 0.2) !important; border-radius: 8px !important; padding: 1em; margin-bottom: 1em;}
-.gr-tabitem { background-color: rgba(52, 73, 94, 0.75) !important; border-radius: 6px !important; padding: 1em !important; border: 1px solid rgba(189, 195, 199, 0.1) !important;}
-.gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message, .gr-json, .gr-file input[type="file"], .gr-file button { border-color: rgba(189, 195, 199, 0.3) !important; background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; border-radius: 6px !important;}
-.gr-textarea textarea, .gr-textbox input { color: #ecf0f1 !important; }
-.gr-button.gr-button-primary { background-color: #1abc9c !important; color: white !important; border-color: #16a085 !important; }
-.gr-button.gr-button-secondary { background-color: #9b59b6 !important; color: white !important; border-color: #8e44ad !important;}
-.gr-button.gr-button-stop { background-color: #e74c3c !important; color: white !important; border-color: #c0392b !important; }
-.gr-markdown { padding: 5px; } 
-.gr-group .gr-markdown { padding: 0px; background-color: transparent !important; }
-.gr-markdown h1, .gr-markdown h2, .gr-markdown h3 { color: #ecf0f1 !important; border-bottom-color: rgba(189, 195, 199, 0.3) !important; margin-top: 0.5em; margin-bottom: 0.5em;}
-.gr-markdown h1 {font-size: 1.5rem;} .gr-markdown h2 {font-size: 1.25rem;} .gr-markdown h3 {font-size: 1.1rem;}
-.gr-markdown p, .gr-markdown li { color: #ecf0f1 !important; }
-.gr-markdown pre code { background-color: rgba(52, 73, 94, 0.95) !important; border-color: rgba(189, 195, 199, 0.3) !important; }
-.gr-chatbot { background-color: rgba(44, 62, 80, 0.7) !important; border-color: rgba(189, 195, 199, 0.2) !important; }
-.gr-chatbot .message { background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; border-color: rgba(189, 195, 199, 0.3) !important; }
-.gr-chatbot .message.user { background-color: rgba(46, 204, 113, 0.9) !important; color: black !important; }
-.gr-input-label > .label-text, .gr-dropdown-label > .label-text { color: #bdc3c7 !important; }
-.status-bar { padding: 0.5rem 1rem; border-radius: 6px; margin-bottom: 1rem; background-color: rgba(44, 62, 80, 0.8) !important; }
-.tabnav button { background-color: rgba(52, 73, 94, 0.8) !important; color: #ecf0f1 !important; border-bottom: 2px solid transparent !important;}
-.tabnav button.selected { background-color: rgba(44, 62, 80, 0.95) !important; color: #1abc9c !important; border-bottom: 2px solid #1abc9c !important;}
-.app-title-area { padding: 0.5rem 0; text-align: center;} /* Area for title and subtitle */
-"""
-
 
 with gr.Blocks(
-    theme=gr.themes.Soft(),  # Use a clean, modern theme
+    theme=gr.themes.Soft(),
     css="""
         .gr-button { margin: 5px; }
         .gr-textbox, .gr-text-area, .gr-dropdown { border-radius: 8px; }
@@ -442,7 +484,6 @@ with gr.Blocks(
         .status-text { font-size: 0.9em; color: #555; }
     """
 ) as demo:
-    # Header Section
     gr.Markdown(
         """
         # 🤖 AI Research Agent
@@ -451,56 +492,39 @@ with gr.Blocks(
         elem_classes=["header"]
     )
 
-    # Status Bar
     is_sqlite = MEMORY_STORAGE_BACKEND == "SQLITE"
     is_hf_dataset = MEMORY_STORAGE_BACKEND == "HF_DATASET"
     
     with gr.Row(variant="compact"):
         agent_stat_tb = gr.Textbox(
-            label="Agent Status",
-            value="Initializing systems...",
-            interactive=False,
-            elem_classes=["status-text"],
-            scale=4
+            label="Agent Status", value="Initializing systems...", interactive=False,
+            elem_classes=["status-text"], scale=4
         )
         with gr.Column(scale=1, min_width=150):
             memory_backend_info_tb = gr.Textbox(
-                label="Memory Backend",
-                value=MEMORY_STORAGE_BACKEND,
-                interactive=False,
+                label="Memory Backend", value=MEMORY_STORAGE_BACKEND, interactive=False,
                 elem_classes=["status-text"]
             )
             sqlite_path_display = gr.Textbox(
-                label="SQLite Path",
-                value=MEMORY_SQLITE_PATH,
-                interactive=False,
-                visible=is_sqlite,  # Use precomputed boolean
-                elem_classes=["status-text"]
+                label="SQLite Path", value=MEMORY_SQLITE_PATH, interactive=False,
+                visible=is_sqlite, elem_classes=["status-text"]
             )
             hf_repos_display = gr.Textbox(
-                label="HF Repos",
-                value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}",
-                interactive=False,
-                visible=is_hf_dataset,  # Use precomputed boolean
-                elem_classes=["status-text"]
+                label="HF Repos", value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}",
+                interactive=False, visible=is_hf_dataset, elem_classes=["status-text"]
             )
 
     with gr.Row():
-        # Sidebar (unchanged)
         with gr.Sidebar():
             gr.Markdown("## ⚙️ Configuration")
             with gr.Group():
                 gr.Markdown("### AI Model Settings")
                 api_key_tb = gr.Textbox(
-                    label="AI Provider API Key (Override)",
-                    type="password",
-                    placeholder="Uses .env if blank"
+                    label="AI Provider API Key (Override)", type="password", placeholder="Uses .env if blank"
                 )
                 prov_sel_dd = gr.Dropdown(
-                    label="AI Provider",
-                    choices=get_available_providers(),
-                    value=get_available_providers()[0] if get_available_providers() else None,
-                    interactive=True
+                    label="AI Provider", choices=get_available_providers(),
+                    value=get_available_providers()[0] if get_available_providers() else None, interactive=True
                 )
                 model_sel_dd = gr.Dropdown(
                     label="AI Model",
@@ -511,98 +535,75 @@ with gr.Blocks(
             with gr.Group():
                 gr.Markdown("### System Prompt")
                 sys_prompt_tb = gr.Textbox(
-                    label="System Prompt Base",
-                    lines=8,
-                    value=DEFAULT_SYSTEM_PROMPT,
-                    interactive=True
+                    label="System Prompt Base", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True
                 )
             if MEMORY_STORAGE_BACKEND == "RAM":
                 save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")
 
-        # Main Content Area
         with gr.Column(scale=3):
             with gr.Tabs():
                 with gr.TabItem("💬 Chat & Research"):
                     with gr.Group():
                         gr.Markdown("### AI Chat Interface")
                         main_chat_disp = gr.Chatbot(
-                            label=None,
-                            height=400,
-                            bubble_full_width=False,
+                            label=None, height=400, bubble_full_width=False,
                             avatar_images=(None, "https://raw.githubusercontent.com/huggingface/brand-assets/main/hf-logo-with-title.png"),
-                            show_copy_button=True,
-                            render_markdown=True,
-                            sanitize_html=True
+                            show_copy_button=True, render_markdown=True, sanitize_html=True
                         )
                         with gr.Row(variant="compact"):
                             user_msg_tb = gr.Textbox(
-                                show_label=False,
-                                placeholder="Ask your research question...",
-                                scale=7,
-                                lines=1,
-                                max_lines=3
+                                show_label=False, placeholder="Ask your research question...",
+                                scale=7, lines=1, max_lines=3
                             )
                             send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
                         with gr.Accordion("📝 Detailed Response", open=False):
                             fmt_report_tb = gr.Textbox(
-                                label="Full AI Response",
-                                lines=8,
-                                interactive=True,
-                                show_copy_button=True
+                                label="Full AI Response", lines=8, interactive=True, show_copy_button=True
                             )
                             dl_report_btn = gr.DownloadButton(
-                                "Download Report",
-                                interactive=False,
-                                visible=False
+                                "Download Report", value=None, interactive=False, visible=False # Ensure value is None initially
                             )
-                        detect_out_md = gr.Markdown(visible=False)
+                        detect_out_md = gr.Markdown(visible=False) # Ensure visible is False initially
 
                 with gr.TabItem("🧠 Knowledge Base"):
                     with gr.Row(equal_height=True):
                         with gr.Column():
                             gr.Markdown("### 📜 Rules Management")
                             rules_disp_ta = gr.TextArea(
-                                label=None,
-                                lines=10,
-                                placeholder="View or edit rules...",
-                                interactive=True
+                                label=None, lines=10, placeholder="View or edit rules...", interactive=True
                             )
                             with gr.Row(variant="compact"):
                                 view_rules_btn = gr.Button("🔄 Load Rules")
                                 save_edited_rules_btn = gr.Button("💾 Save Rules", variant="primary")
+                            with gr.Row(variant="compact"):
+                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules", value=None) # Add download button
                                 clear_rules_btn = gr.Button("🗑️ Clear Rules", variant="stop")
                             upload_rules_fobj = gr.File(
-                                label="Upload Rules",
+                                label="Upload Rules File (.txt, .jsonl)",
                                 file_types=[".txt", ".jsonl"]
                             )
                             rules_stat_tb = gr.Textbox(
-                                label="Rules Status",
-                                interactive=False,
-                                lines=1,
-                                elem_classes=["status-text"]
+                                label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"]
                             )
 
                         with gr.Column():
                             gr.Markdown("### 📚 Memories Management")
                             mems_disp_json = gr.JSON(
-                                label=None,
-                                value={"memories": []}
+                                label=None, value=[] # Initialize with empty list
                             )
                             with gr.Row(variant="compact"):
                                 view_mems_btn = gr.Button("🔄 Load Memories")
+                                dl_mems_btn = gr.DownloadButton("⬇️ Download Memories", value=None) # Add download button
+                            with gr.Row(variant="compact"):
                                 clear_mems_btn = gr.Button("🗑️ Clear Memories", variant="stop")
                             upload_mems_fobj = gr.File(
-                                label="Upload Memories",
-                                file_types=[".jsonl"]
+                                label="Upload Memories File (.jsonl)",
+                                file_types=[".jsonl", ".json"] # Allow .json for single list
                             )
                             mems_stat_tb = gr.Textbox(
-                                label="Memories Status",
-                                interactive=False,
-                                lines=1,
-                                elem_classes=["status-text"]
+                                label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"]
                             )
 
-    # Event Handlers (unchanged logic, cleaned up organization)
     def dyn_upd_model_dd(sel_prov_dyn: str):
         models_dyn = get_model_display_names_for_provider(sel_prov_dyn)
         def_model_dyn = get_default_model_display_name_for_provider(sel_prov_dyn)
@@ -612,10 +613,23 @@ 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]
-    send_btn.click(fn=handle_gradio_chat_submit, inputs=chat_ins, outputs=chat_outs)
-    user_msg_tb.submit(fn=handle_gradio_chat_submit, inputs=chat_ins, outputs=chat_outs)
+    
+    # Chat submission events
+    chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
+    
+    send_btn_click_event = send_btn.click(**chat_event_args)
+    user_msg_submit_event = user_msg_tb.submit(**chat_event_args)
+
+    # Chain UI refreshes for rules and memories after chat interaction
+    for event in [send_btn_click_event, user_msg_submit_event]:
+        event.then(fn=ui_view_rules_action_fn, inputs=None, outputs=rules_disp_ta)
+        event.then(fn=ui_view_memories_action_fn, inputs=None, outputs=mems_disp_json)
+
 
+    # Rules Management events
     view_rules_btn.click(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
+    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn)
+
 
     def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
         if not edited_rules_text.strip():
@@ -638,7 +652,7 @@ with gr.Blocks(
                 skipped += 1
             else:
                 errors += 1
-            progress((idx + 1) / total)
+            progress((idx + 1) / total, desc=f"Processed {idx+1}/{total} rules...")
         return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors}."
 
     save_edited_rules_btn.click(
@@ -656,17 +670,14 @@ with gr.Blocks(
     ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
 
     clear_rules_btn.click(
-        fn=lambda: "All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules.",
+        fn=lambda: ("All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules."),
         outputs=rules_stat_tb
     ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
 
-    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
-        def save_faiss_action_with_feedback_sidebar_fn():
-            save_faiss_indices_to_disk()
-            gr.Info("Attempted to save FAISS indices to disk.")
-        save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None)
-
+    # Memories Management events
     view_mems_btn.click(fn=ui_view_memories_action_fn, outputs=mems_disp_json)
+    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn)
+
 
     upload_mems_fobj.upload(
         fn=ui_upload_memories_action_fn,
@@ -676,10 +687,17 @@ with gr.Blocks(
     ).then(fn=ui_view_memories_action_fn, outputs=mems_disp_json)
 
     clear_mems_btn.click(
-        fn=lambda: "All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories.",
+        fn=lambda: ("All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories."),
         outputs=mems_stat_tb
     ).then(fn=ui_view_memories_action_fn, outputs=mems_disp_json)
 
+    # Save FAISS for RAM backend
+    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
+        def save_faiss_action_with_feedback_sidebar_fn():
+            save_faiss_indices_to_disk()
+            gr.Info("Attempted to save FAISS indices to disk.")
+        save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None)
+
     def app_load_fn():
         initialize_memory_system()
         logger.info("App loaded. Memory system initialized.")
@@ -691,7 +709,7 @@ with gr.Blocks(
     demo.load(fn=app_load_fn, inputs=None, outputs=[agent_stat_tb, rules_disp_ta, mems_disp_json])
 
 if __name__ == "__main__":
-    logger.info(f"Starting Gradio AI Research Mega Agent (v5.8 - UI Style/Layout Update, Memory: {MEMORY_STORAGE_BACKEND})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v5.9 - UI Refresh & Download/Upload Enhancements, 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"