Update app.py

app.py

@@ -430,111 +430,271 @@ body { background: linear-gradient(to bottom right, #2c3e50, #34495e); color: #e
 .app-title-area { padding: 0.5rem 0; text-align: center;} /* Area for title and subtitle */
 """
 
-with gr.Blocks(title="AI Research Agent v5.8") as demo:
-    
-    gr.HTML("<h1> 🤖 AI Research Agent </h1>") 
+
+with gr.Blocks(
+    title="AI Research Agent v5.8",
+    theme=gr.themes.Soft(),  # Use a clean, modern theme
+    css="""
+        .gr-button { margin: 5px; }
+        .gr-textbox, .gr-text-area, .gr-dropdown { border-radius: 8px; }
+        .gr-group { border: 1px solid #e0e0e0; border-radius: 8px; padding: 10px; }
+        .gr-row { gap: 10px; }
+        .gr-tab { border-radius: 8px; }
+        .status-text { font-size: 0.9em; color: #555; }
+    """
+) as demo:
+    # Header Section
+    gr.Markdown(
+        """
+        # 🤖 AI Research Agent
+        Your intelligent assistant for research and knowledge management
+        """,
+        elem_classes=["header"]
+    )
+
+    # Status Bar
+    with gr.Row(variant="compact"):
+        agent_stat_tb = gr.Textbox(
+            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,
+                elem_classes=["status-text"]
+            )
+            sqlite_path_display = gr.Textbox(
+                label="SQLite Path",
+                value=MEMORY_SQLITE_PATH,
+                interactive=False,
+                visible=MEMORY_STORAGE_BACKEND == "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=MEMORY_STORAGE_BACKEND == "HF_DATASET",
+                elem_classes=["status-text"]
+            )
+
     with gr.Row():
-        with gr.Column(scale=3): 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)
-        sqlite_path_display = gr.Textbox(label="SQLite Path", value=f"{MEMORY_SQLITE_PATH}", interactive=False, scale=2) # Use visible
-        hf_repos_display = gr.Textbox(label="HF Repos", value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}", interactive=False, scale=2) # Use visible
-
-    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
-    with gr.Sidebar(): # Using Gradio's explicit Sidebar component
-        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")
-            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)
-        with gr.Group():
-            gr.Markdown("### System Prompt")
-            sys_prompt_tb = gr.Textbox(label="System Prompt Base", lines=10, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
-        if MEMORY_STORAGE_BACKEND == "RAM":
-            save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")
-        
-    with gr.Tabs():
-        with gr.TabItem("💬 AI Chat & Research Output"):
-            gr.Markdown("### AI Chat Interface") # Using H3 for section titles within tabs
-            main_chat_disp = gr.Chatbot(label="AI Research Chat", height=500, 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 your research question...", 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 / Details", open=False):
-                fmt_report_tb = gr.Textbox(label="Full AI Response", lines=10, interactive=True, show_copy_button=True)
-                dl_report_btn = gr.DownloadButton("Download Report", interactive=False, visible=False) 
-            detect_out_md = gr.Markdown()
-
-        with gr.TabItem("🧠 Knowledge Base Management"):
-            gr.Markdown("### Manage Stored Knowledge")
-            with gr.Row():
-                with gr.Column():
-                    with gr.Group():
-                        gr.Markdown("#### 📜 Rules (Insights)")
-                        rules_disp_ta = gr.TextArea(label="View/Edit Rules", lines=15, interactive=True, placeholder="Load rules or type new ones...")
-                        with gr.Row(): view_rules_btn = gr.Button("🔄 Load"); save_edited_rules_btn = gr.Button("💾 Save Edits", variant="primary")
-                        upload_rules_fobj = gr.File(label="Upload Rules File", file_types=[".txt", ".jsonl"])
-                        rules_stat_tb = gr.Textbox(label="Rules Status", interactive=False, lines=2)
-                        clear_rules_btn = gr.Button("⚠️ Clear All Rules", variant="stop")
-                with gr.Column():
+        # 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"
+                )
+                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
+                )
+                model_sel_dd = gr.Dropdown(
+                    label="AI Model",
+                    choices=get_model_display_names_for_provider(get_available_providers()[0]) if get_available_providers() else [],
+                    value=get_default_model_display_name_for_provider(get_available_providers()[0]) if get_available_providers() else None,
+                    interactive=True
+                )
+            with gr.Group():
+                gr.Markdown("### System Prompt")
+                sys_prompt_tb = gr.Textbox(
+                    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("#### 📚 Memories")
-                        mems_disp_json = gr.JSON(label="View Memories") 
-                        view_mems_btn = gr.Button("🔄 Load Memories")
-                        upload_mems_fobj = gr.File(label="Upload Memories File", file_types=[".jsonl"])
-                        mems_stat_tb = gr.Textbox(label="Memories Status", interactive=False, lines=2)
-                        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)
+                        gr.Markdown("### AI Chat Interface")
+                        main_chat_disp = gr.Chatbot(
+                            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
+                        )
+                        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,
+                                autofocus=True
+                            )
+                            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
+                            )
+                            dl_report_btn = gr.DownloadButton(
+                                "Download Report",
+                                interactive=False,
+                                visible=False
+                            )
+                        detect_out_md = gr.Markdown(visible=False)
+
+                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
+                            )
+                            with gr.Row(variant="compact"):
+                                view_rules_btn = gr.Button("🔄 Load Rules")
+                                save_edited_rules_btn = gr.Button("💾 Save Rules", variant="primary")
+                                clear_rules_btn = gr.Button("🗑️ Clear Rules", variant="stop")
+                            upload_rules_fobj = gr.File(
+                                label="Upload Rules",
+                                file_types=[".txt", ".jsonl"]
+                            )
+                            rules_stat_tb = gr.Textbox(
+                                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": []}
+                            )
+                            with gr.Row(variant="compact"):
+                                view_mems_btn = gr.Button("🔄 Load Memories")
+                                clear_mems_btn = gr.Button("🗑️ Clear Memories", variant="stop")
+                            upload_mems_fobj = gr.File(
+                                label="Upload Memories",
+                                file_types=[".jsonl"]
+                            )
+                            mems_stat_tb = gr.Textbox(
+                                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)
+        return gr.Dropdown(choices=models_dyn, value=def_model_dyn, interactive=True)
+
     prov_sel_dd.change(fn=dyn_upd_model_dd, inputs=prov_sel_dd, outputs=model_sel_dd)
+
     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)
+    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)
+
     view_rules_btn.click(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
+
     def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
-        if not edited_rules_text.strip(): return "No rules text to save."
+        if not edited_rules_text.strip():
+            return "No rules text to save."
         potential_rules = edited_rules_text.split("\n\n---\n\n")
-        if len(potential_rules) == 1 and "\n" in edited_rules_text: potential_rules = [r.strip() for r in edited_rules_text.splitlines() if r.strip()]
-        if not potential_rules: return "No rules found to process from editor."
-        added, skipped, errors = 0,0,0; total = len(potential_rules)
+        if len(potential_rules) == 1 and "\n" in edited_rules_text:
+            potential_rules = [r.strip() for r in edited_rules_text.splitlines() if r.strip()]
+        if not potential_rules:
+            return "No rules found to process from editor."
+        added, skipped, errors = 0, 0, 0
+        total = len(potential_rules)
         progress(0, desc=f"Saving {total} rules from editor...")
         for idx, rule_text in enumerate(potential_rules):
-            if not rule_text.strip(): continue
+            if not rule_text.strip():
+                continue
             success, status_msg = add_rule_entry(rule_text.strip())
-            if success: added +=1
-            elif status_msg == "duplicate": skipped +=1
-            else: errors +=1
-            progress((idx+1)/total)
+            if success:
+                added += 1
+            elif status_msg == "duplicate":
+                skipped += 1
+            else:
+                errors += 1
+            progress((idx + 1) / total)
         return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors}."
-    save_edited_rules_btn.click(fn=save_edited_rules_action_fn, inputs=[rules_disp_ta], outputs=[rules_stat_tb], show_progress="full").then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
-    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_sidebar_btn' in locals() and save_faiss_sidebar_btn is not None:
-        def save_faiss_action_with_feedback_sidebar_fn(): save_faiss_indices_to_disk(); gr.Info("Attempted to save FAISS indices to disk.")
+
+    save_edited_rules_btn.click(
+        fn=save_edited_rules_action_fn,
+        inputs=[rules_disp_ta],
+        outputs=[rules_stat_tb],
+        show_progress="full"
+    ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)
+
+    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_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)
+
     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)
-    clear_mems_btn.click(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)
-    def app_load_fn(): 
+
+    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)
+
+    clear_mems_btn.click(
+        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)
+
+    def app_load_fn():
         initialize_memory_system()
         logger.info("App loaded. Memory system initialized.")
-        backend_status = f"AI Systems Initialized. Ready."
+        backend_status = "AI Systems Initialized. Ready."
         rules_on_load = ui_view_rules_action_fn()
         mems_on_load = ui_view_memories_action_fn()
-        # Initial visibility for conditional Textboxes based on backend
-        vis_sqlite = MEMORY_STORAGE_BACKEND == "SQLITE"
-        vis_hf = MEMORY_STORAGE_BACKEND == "HF_DATASET"
         return backend_status, rules_on_load, mems_on_load
 
     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})...")
-    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"
+    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"
+    app_share = os.getenv("GRADIO_SHARE", "False").lower() == "true"
     logger.info(f"Launching Gradio server: http://{app_server}:{app_port}. Debug: {app_debug}, Share: {app_share}")
     demo.queue().launch(server_name=app_server, server_port=app_port, debug=app_debug, share=app_share)
     logger.info("Gradio application shut down.")