Update app.py

app.py

@@ -269,10 +269,7 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
     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...*")
-    # For DownloadButton, create a placeholder value that's a valid file path (even if empty initially)
-    # Or simply update it to gr.DownloadButton(interactive=False, value=None, visible=False)
     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))
@@ -302,12 +299,10 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
                 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 curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
-                    # Create a temporary file for the download
                     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 # This path will be passed to DownloadButton
-                    # Update DownloadButton with the path to the temp file
-                    def_dl_btn = gr.DownloadButton(label="Download Report (.md)", value=temp_dl_file_path, visible=True, interactive=True)
+                        temp_dl_file_path = tmpfile.name
+                    def_dl_btn = gr.DownloadButton(value=temp_dl_file_path, visible=True, interactive=True) # `label` is optional here, Gradio uses filename
                 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)
@@ -328,14 +323,9 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
         status_txt = "Response complete. Background learning initiated."
     else: status_txt = "Processing finished; no response or error."
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
-    # Clean up the temporary file after the download button has been yielded (and potentially used)
-    # This cleanup is tricky with Gradio's yield. The file needs to exist when Gradio processes the yield.
-    # A more robust way is to have Gradio manage the temp file or use a cleanup mechanism that runs after the HTTP response.
-    # For simplicity, we'll leave it, but in a prod app, temp files should be managed.
-    # 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}")
-
+    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}")
 
 def ui_view_rules_action_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."
 def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
@@ -392,71 +382,92 @@ 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
 
-custom_theme = gr.themes.Base(primary_hue="teal", secondary_hue="purple", neutral_hue="zinc", text_size="sm", spacing_size="sm", radius_size="sm", font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"])
-custom_css = """ body { font-family: 'Inter', sans-serif; } .gradio-container { max-width: 98% !important; margin: auto !important; padding-top: 0.5rem !important; } footer { display: none !important; } .gr-button { white-space: nowrap; } .gr-input, .gr-textarea textarea, .gr-dropdown input { border-radius: 8px !important; } .gr-chatbot .message { border-radius: 10px !important; box-shadow: 0 2px 5px rgba(0,0,0,0.08) !important; } .prose { h1 { font-size: 1.8rem; margin-bottom: 0.6em; margin-top: 0.8em; } h2 { font-size: 1.4rem; margin-bottom: 0.5em; margin-top: 0.7em; } h3 { font-size: 1.15rem; margin-bottom: 0.4em; margin-top: 0.6em; } p { margin-bottom: 0.8em; line-height: 1.65; } ul, ol { margin-left: 1.5em; margin-bottom: 0.8em; } code { background-color: #f1f5f9; padding: 0.2em 0.45em; border-radius: 4px; font-size: 0.9em; } pre > code { display: block; padding: 0.8em; overflow-x: auto; background-color: #f8fafc; border: 1px solid #e2e8f0; border-radius: 6px;}} .compact-group .gr-input-label, .compact-group .gr-dropdown-label { font-size: 0.8rem !important; padding-bottom: 2px !important;} .status-bar { padding: 0.5em; margin-bottom: 0.5em; border-radius: 6px; background-color: #eef2ff;} """ # Added status-bar style
+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; min-height: 100vh; }
+.gradio-container { background: transparent !important; max-width: 100% !important; padding: 10px !important;}
+.gr-box, .gr-panel, .gr-tabs { background-color: rgba(44, 62, 80, 0.85) !important; border-color: rgba(189, 195, 199, 0.2) !important; border-radius: 8px !important;}
+.gr-tabitem { background-color: rgba(52, 73, 94, 0.7) !important; border-radius: 6px !important; margin-bottom: 5px !important;}
+.gr-accordion { background-color: rgba(52, 73, 94, 0.75) !important; border-radius: 6px !important; }
+.gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message, .gr-json { 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 { color: #ecf0f1 !important; } /* Ensure textarea text is light */
+.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 { background-color: transparent !important; padding: 0px; border-radius: 0px; } /* Make markdown background transparent within styled boxes */
+.gr-markdown h1, .gr-markdown h2, .gr-markdown h3 { color: #1abc9c !important; border-bottom: 1px solid rgba(189, 195, 199, 0.2) !important; padding-bottom: 0.3em; margin-top:1em; }
+.gr-markdown p, .gr-markdown li { color: #ecf0f1 !important; }
+.gr-markdown pre code { background-color: rgba(30, 40, 50, 0.95) !important; border-color: rgba(189, 195, 199, 0.3) !important; color: #ecf0f1; border-radius: 4px; padding: 0.8em; }
+.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: 1px solid rgba(189, 195, 199, 0.2) !important; box-shadow: 0 1px 2px rgba(0,0,0,0.1) !important; }
+.gr-chatbot .message.user { background-color: rgba(46, 204, 113, 0.8) !important; color: #2c3e50 !important; }
+.gr-input-label > .label-text, .gr-dropdown-label > .label-text { color: #bdc3c7 !important; font-size: 0.9rem !important; } /* Lighter labels */
+.gr-info { color: #bdc3c7 !important; font-size: 0.85rem !important; }
+#status_bar_info { background-color: rgba(44, 62, 80, 0.8) !important; padding: 8px; border-radius: 6px; margin-bottom: 10px; }
+"""
 
-with gr.Blocks(theme=custom_theme, css=custom_css, title="AI Research Mega Agent v5") as demo:
-    gr.Markdown("# 🚀 AI Research Mega Agent", elem_classes="prose")
+with gr.Blocks(theme=custom_theme, css=custom_css, title="AI Research Mega Agent v5.1") as demo:
+    gr.Markdown("# 🚀 AI Research Mega Agent Dashboard")
     
     avail_provs = get_available_providers()
     def_prov = avail_provs[0] if avail_provs else None
     def_models = get_model_display_names_for_provider(def_prov) if def_prov else []
     def_model_disp = get_default_model_display_name_for_provider(def_prov) if def_prov else None
+    save_faiss_ram_btn_kb = None # For KB tab
 
-    with gr.Row(elem_classes="status-bar"):
-        agent_stat_tb = gr.Textbox(label="Agent Status", interactive=False, lines=1, value="Initializing systems...", scale=3)
-        memory_backend_info_tb = gr.Textbox(label="Memory Backend", value=f"{MEMORY_STORAGE_BACKEND}", interactive=False, scale=1)
+    with gr.Row(elem_id="status_bar_info"):
+        with gr.Column(scale=2): agent_stat_tb = gr.Textbox(label="Agent Status", interactive=False, lines=1, value="Initializing systems...")
+        with gr.Column(scale=1): memory_backend_info_tb = gr.Textbox(label="Memory Backend", value=f"{MEMORY_STORAGE_BACKEND}", interactive=False)
         if MEMORY_STORAGE_BACKEND == "SQLITE":
-            gr.Textbox(label="SQLite Path", value=f"{MEMORY_SQLITE_PATH}", interactive=False, scale=2)
+            with gr.Column(scale=2): gr.Textbox(label="SQLite Path", value=f"{MEMORY_SQLITE_PATH}", interactive=False)
         elif MEMORY_STORAGE_BACKEND == "HF_DATASET":
-            gr.Textbox(label="HF Repos", value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}", interactive=False, scale=2)
-
-
+            with gr.Column(scale=2): gr.Textbox(label="HF Repos", value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}", interactive=False)
+    
     with gr.Tabs() as main_tabs:
-        with gr.TabItem("💬 Dashboard & Chat", id=0):
+        with gr.TabItem("⚙️ Agent Configuration & Chat", id=0):
             with gr.Row():
-                with gr.Column(scale=1, min_width=320):
-                    gr.Markdown("### ⚙️ LLM Configuration", elem_classes="prose")
-                    with gr.Group(elem_classes="compact-group"):
-                        prov_sel_dd = gr.Dropdown(label="Provider", choices=avail_provs, value=def_prov, interactive=True)
-                        model_sel_dd = gr.Dropdown(label="Model", choices=def_models, value=def_model_disp, interactive=True)
-                        api_key_tb = gr.Textbox(label="API Key Override (Optional)", type="password", placeholder="Enter API key if not in .env")
-                    with gr.Group(elem_classes="compact-group"):
+                with gr.Column(scale=1, min_width=350):
+                    gr.Markdown("## ⚙️ LLM & System Configuration")
+                    with gr.Group():
+                        prov_sel_dd = gr.Dropdown(label="AI Provider", choices=avail_provs, value=def_prov, interactive=True)
+                        model_sel_dd = gr.Dropdown(label="AI Model", choices=def_models, value=def_model_disp, interactive=True)
+                        api_key_tb = gr.Textbox(label="API Key Override (Optional)", type="password", placeholder="Paste key for selected provider")
+                    with gr.Group():
                         sys_prompt_tb = gr.Textbox(label="System Prompt Base", lines=10, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
                 
                 with gr.Column(scale=2):
-                    main_chat_disp = gr.Chatbot(label="AI Research Chat", height=600, 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)
+                    gr.Markdown("## 💬 AI Chat Interface")
+                    main_chat_disp = gr.Chatbot(label="AI Research Chat", height=550, 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)
                     with gr.Row():
-                        user_msg_tb = gr.Textbox(show_label=False, placeholder="Ask a question or give an instruction...", scale=7, lines=1, max_lines=5, autofocus=True)
+                        user_msg_tb = gr.Textbox(show_label=False, placeholder="Ask your research question or give an instruction...", scale=7, lines=1, max_lines=5, autofocus=True)
                         send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
             
-            with gr.Accordion("📝 Full Response / Output", open=True):
-                fmt_report_tb = gr.Textbox(label="Current Research Output", lines=10, interactive=True, show_copy_button=True, value="*AI responses will appear here...*")
+            with gr.Accordion("📝 Full Response / Output Details", open=False):
+                fmt_report_tb = gr.Textbox(label="Full AI Response", lines=10, interactive=True, show_copy_button=True, value="*AI responses will appear here...*")
                 dl_report_btn = gr.DownloadButton("Download Report", interactive=False, visible=False) 
-            detect_out_md = gr.Markdown("*Insights used or other intermediate details will show here...*")
+                detect_out_md = gr.Markdown("*Insights used or other intermediate details will show here...*")
 
         with gr.TabItem("🧠 Knowledge Base Management", id=1):
-            gr.Markdown("## Manage Stored Knowledge (Backend: " + MEMORY_STORAGE_BACKEND + ")", elem_classes="prose")
+            gr.Markdown("## Manage Stored Knowledge")
             with gr.Row(equal_height=False):
                 with gr.Column(scale=1):
-                    gr.Markdown("### Rules (Learned Insights)", elem_classes="prose")
+                    gr.Markdown("### 📜 Rules (Learned Insights)")
                     rules_disp_ta = gr.TextArea(label="View/Edit Rules (one per line or '---' separated)", lines=15, interactive=True, placeholder="Load or type rules here...")
                     with gr.Row():
-                        view_rules_btn = gr.Button("🔄 Load/Refresh Rules View")
-                        save_edited_rules_btn = gr.Button("💾 Save Edited Rules", variant="primary")
+                        view_rules_btn = gr.Button("🔄 Load/Refresh Rules"); save_edited_rules_btn = gr.Button("💾 Save Edited Rules", variant="primary")
                     upload_rules_fobj = gr.File(label="Upload Rules File (.txt/.jsonl)", file_types=[".txt", ".jsonl"])
-                    rules_stat_tb = gr.Textbox(label="Rules Action Status", interactive=False, lines=2, placeholder="Status of rule operations...")
+                    rules_stat_tb = gr.Textbox(label="Rules Op Status", interactive=False, lines=2, placeholder="Status...")
                     with gr.Row():
                         clear_rules_btn = gr.Button("⚠️ Clear All Rules", variant="stop")
-                        save_faiss_ram_btn = gr.Button("Save FAISS Indices", visible=(MEMORY_STORAGE_BACKEND == "RAM"))
+                        if MEMORY_STORAGE_BACKEND == "RAM": 
+                            save_faiss_ram_btn_kb = gr.Button("Save FAISS Indices")
                 
                 with gr.Column(scale=1):
-                    gr.Markdown("### Memories (Past Interactions)", elem_classes="prose")
-                    mems_disp_json = gr.JSON(label="View Memories (JSON format)", scale=2) # scale makes it taller
-                    with gr.Row():
-                        view_mems_btn = gr.Button("🔄 Load/Refresh Memories View")
+                    gr.Markdown("### 📚 Memories (Past Interactions)")
+                    mems_disp_json = gr.JSON(label="View Memories (JSON format)", scale=2)
+                    with gr.Row(): view_mems_btn = gr.Button("🔄 Load/Refresh Memories")
                     upload_mems_fobj = gr.File(label="Upload Memories File (.jsonl)", file_types=[".jsonl"])
-                    mems_stat_tb = gr.Textbox(label="Memories Action Status", interactive=False, lines=2, placeholder="Status of memory operations...")
+                    mems_stat_tb = gr.Textbox(label="Memories Op Status", interactive=False, lines=2, placeholder="Status...")
                     clear_mems_btn = gr.Button("⚠️ Clear All Memories", variant="stop")
     
     def dyn_upd_model_dd(sel_prov_dyn:str): models_dyn, def_model_dyn = get_model_display_names_for_provider(sel_prov_dyn), get_default_model_display_name_for_provider(sel_prov_dyn); return gr.Dropdown(choices=models_dyn, value=def_model_dyn, interactive=True)
@@ -486,9 +497,9 @@ with gr.Blocks(theme=custom_theme, css=custom_css, title="AI Research Mega Agent
     upload_rules_fobj.upload(fn=ui_upload_rules_action_fn, inputs=[upload_rules_fobj], outputs=[rules_stat_tb], show_progress="full").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.", outputs=rules_stat_tb).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
     
-    if MEMORY_STORAGE_BACKEND == "RAM" and save_faiss_ram_btn is not None:
-        def save_faiss_action_with_feedback_fn(): save_faiss_indices_to_disk(); gr.Info("Attempted to save FAISS indices to disk.")
-        save_faiss_ram_btn.click(fn=save_faiss_action_with_feedback_fn, inputs=None, outputs=None)
+    if MEMORY_STORAGE_BACKEND == "RAM" and save_faiss_ram_btn_kb is not None:
+        def save_faiss_action_with_feedback_fn_kb(): save_faiss_indices_to_disk(); gr.Info("Attempted to save FAISS indices to disk.")
+        save_faiss_ram_btn_kb.click(fn=save_faiss_action_with_feedback_fn_kb, inputs=None, outputs=None)
     
     view_mems_btn.click(fn=ui_view_memories_action_fn, outputs=mems_disp_json)
     upload_mems_fobj.upload(fn=ui_upload_memories_action_fn, inputs=[upload_mems_fobj], outputs=[mems_stat_tb], show_progress="full").then(fn=ui_view_memories_action_fn, outputs=mems_disp_json)
@@ -497,18 +508,27 @@ with gr.Blocks(theme=custom_theme, css=custom_css, title="AI Research Mega Agent
     def app_load_fn(): 
         initialize_memory_system()
         logger.info("App loaded. Memory system initialized.")
-        backend_status = f"AI Systems Initialized (Backend: {MEMORY_STORAGE_BACKEND}). Ready."
-        # Initial load for KB views
+        backend_status = f"AI Systems Initialized. Ready." # Agent status, backend info is separate
         rules_on_load = ui_view_rules_action_fn()
         mems_on_load = ui_view_memories_action_fn()
+        # Update memory backend info text box too
+        mem_backend_text = f"{MEMORY_STORAGE_BACKEND}"
+        if MEMORY_STORAGE_BACKEND == "SQLITE": mem_backend_path_text = f"{MEMORY_SQLITE_PATH}"
+        elif MEMORY_STORAGE_BACKEND == "HF_DATASET": mem_backend_path_text = f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}"
+        else: mem_backend_path_text = "N/A" # For RAM
+
+        # The number of outputs from app_load_fn must match the number of components in demo.load's outputs list
+        # Make sure memory_backend_info_tb and its potential path display are updated if they are separate components
+        # The current UI has one tb for backend and one for path (conditional)
+        # For simplicity here, I'll assume they are part of the static display or updated differently if needed.
+        # The agent_stat_tb is the primary one to update here.
         return backend_status, rules_on_load, mems_on_load
 
     # Update outputs of demo.load to populate the KB views on startup
     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 with Dashboard UI & Advanced Memory: {MEMORY_STORAGE_BACKEND})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v5.1 with New UI & Advanced Memory: {MEMORY_STORAGE_BACKEND})...")
     app_port, app_server = int(os.getenv("GRADIO_PORT", 7860)), os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
     app_debug, app_share = os.getenv("GRADIO_DEBUG", "False").lower()=="true", os.getenv("GRADIO_SHARE", "False").lower()=="true"
     logger.info(f"Launching Gradio server: http://{app_server}:{app_port}. Debug: {app_debug}, Share: {app_share}")