Update app.py

app.py

@@ -1,225 +1,312 @@
-import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
-import base64
-import io
+# app.py
+
+import os
 import json
 import logging
-import os
-import requests
-import struct
-import numpy as np
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import rsa, padding
-from cryptography.hazmat.primitives.ciphers.aead import AESGCM
-from cryptography.hazmat.primitives import hashes
-
-# --- Configure Logging ---
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+import tempfile
+from dotenv import load_dotenv
+import gradio as gr
+
+load_dotenv()
+
+# --- Keep all your existing constants and setup ---
+MEMORY_STORAGE_TYPE = "HF_DATASET"
+HF_DATASET_MEMORY_REPO = "broadfield-dev/ai-brain"
+HF_DATASET_RULES_REPO = "broadfield-dev/ai-rules"
+
+os.environ['STORAGE_BACKEND'] = MEMORY_STORAGE_TYPE
+if MEMORY_STORAGE_TYPE == "HF_DATASET":
+    os.environ['HF_MEMORY_DATASET_REPO'] = HF_DATASET_MEMORY_REPO
+    os.environ['HF_RULES_DATASET_REPO'] = HF_DATASET_RULES_REPO
+
+from model_logic import get_available_providers, get_model_display_names_for_provider, get_default_model_display_name_for_provider
+from memory_logic import (
+    initialize_memory_system, add_memory_entry, get_all_memories_cached, clear_all_memory_data_backend,
+    add_rule_entry, remove_rule_entry, get_all_rules_cached, clear_all_rules_data_backend,
+    save_faiss_indices_to_disk, STORAGE_BACKEND as MEMORY_STORAGE_BACKEND, SQLITE_DB_PATH as MEMORY_SQLITE_PATH,
+    HF_MEMORY_DATASET_REPO as MEMORY_HF_MEM_REPO, HF_RULES_DATASET_REPO as MEMORY_HF_RULES_REPO
+)
+from tools.orchestrator import orchestrate_and_respond
+from learning import perform_post_interaction_learning
+from utils import load_rules_from_file, load_memories_from_file
+from prompts import DEFAULT_SYSTEM_PROMPT
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(threadName)s - %(message)s')
 logger = logging.getLogger(__name__)
+for lib_name in ["urllib3", "requests", "huggingface_hub", "PIL.PngImagePlugin", "matplotlib", "gradio_client.client", "multipart.multipart", "httpx", "sentence_transformers", "faiss", "datasets"]:
+    if logging.getLogger(lib_name): logging.getLogger(lib_name).setLevel(logging.WARNING)
 
-# ==============================================================================
-#  CONFIGURATION
-# ==============================================================================
-CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
-BASE_HF_URL = "https://huggingface.co/spaces/"
-# Define these here for cleaner linking in the UI
-CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
-SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
-CREATOR_URL = f"{BASE_HF_URL}{CREATOR_SPACE_ID}"
-SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
-
-# ==============================================================================
-#  LOCAL LOGIC (Key and Image Generation)
-# ==============================================================================
-def generate_rsa_keys():
-    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
-    private_pem = private_key.private_bytes(encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption()).decode('utf-8')
-    public_pem = private_key.public_key().public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo).decode('utf-8')
-    return private_pem, public_pem
-
-def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    if not secret_data_str.strip(): raise ValueError("Secret data cannot be empty.")
-    if not public_key_pem.strip(): raise ValueError("Public Key cannot be empty.")
-    data_dict = {}
-    for line in secret_data_str.splitlines():
-        line = line.strip()
-        if not line or line.startswith('#'): continue
-        parts = line.split(':', 1) if ':' in line else line.split('=', 1)
-        if len(parts) == 2: data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
-    if not data_dict: raise ValueError("No valid key-value pairs found.")
-    json_bytes = json.dumps(data_dict).encode('utf-8')
-    public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
-    aes_key, nonce = os.urandom(32), os.urandom(12)
-    ciphertext = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
-    rsa_encrypted_key = public_key.encrypt(aes_key, padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
-    encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
-    img = Image.new('RGB', (800, 600), color=(45, 52, 54))
-    draw = ImageDraw.Draw(img)
-    try: font = ImageFont.truetype("DejaVuSans.ttf", 40)
-    except IOError: font = ImageFont.load_default(size=30)
-    draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
-    pixel_data = np.array(img.convert("RGB")).ravel()
-    binary_payload = ''.join(format(b, '08b') for b in struct.pack('>I', len(encrypted_payload)) + encrypted_payload)
-    if len(binary_payload) > pixel_data.size: raise ValueError("Data too large for image.")
-    pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
-    stego_pixels = pixel_data.reshape((600, 800, 3))
-    return Image.fromarray(stego_pixels, 'RGB')
-
-
-# ==============================================================================
-#  UI HELPER & REMOTE API CALL LOGIC
-# ==============================================================================
-
-def get_server_list():
-    status = f"Fetching server list from remote config..."
-    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
-    try:
-        response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
-        response.raise_for_status()
-        all_entries = response.json()
-        valid_endpoints = []
-        required_keys = ["name", "api_endpoint", "public_key"]
-        for entry in all_entries:
-            if all(key in entry for key in required_keys):
-                valid_endpoints.append(entry)
-            else:
-                logger.warning(f"Skipping invalid entry in configuration file: {entry}")
-        if not valid_endpoints:
-            raise ValueError("No valid server configurations found in the remote file.")
-        endpoint_names = [e['name'] for e in valid_endpoints]
-        status = f"✅ Success! Found {len(endpoint_names)} valid servers."
-        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, valid_endpoints
-    except Exception as e:
-        status = f"❌ Error fetching or parsing configuration: {e}"
-        yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
+MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", 7))
+LOAD_RULES_FILE = os.getenv("LOAD_RULES_FILE")
+LOAD_MEMORIES_FILE = os.getenv("LOAD_MEMORIES_FILE")
+current_chat_session_history = []
+
+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,
+    # New stateful inputs from the UI
+    turn_budget: int,
+    orchestrator_state: dict
+):
+    global current_chat_session_history
+    cleared_input, updated_gr_hist, status_txt = "", list(gr_hist_list), "Initializing..."
+    updated_rules_text = ui_refresh_rules_display_fn()
+    updated_mems_json = ui_refresh_memories_display_fn()
+    def_detect_out_md = gr.Markdown(visible=False)
+    def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True)
+    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
+
+    # Make a mutable copy of the state for this run
+    latest_orchestrator_state = dict(orchestrator_state)
+
+    # Check for invalid input: empty message when no task is running
+    is_ongoing_task = latest_orchestrator_state.get('is_complex_task', False)
+    if not user_msg_txt.strip() and not is_ongoing_task:
+        status_txt = "Error: Please enter a message to start."
+        # Don't add to history, just show status
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json, latest_orchestrator_state)
+        return
+
+    # If user sends empty message during a task, it means "continue"
+    effective_user_msg = user_msg_txt or "(Continuing previous task...)"
+    updated_gr_hist.append((effective_user_msg, "<i>Thinking...</i>"))
+    # Yield initial response to show the user's message in the chat
+    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json, latest_orchestrator_state)
+
+    internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
+    if len(internal_hist) > MAX_HISTORY_TURNS * 2:
+         internal_hist = internal_hist[-(MAX_HISTORY_TURNS * 2):]
+
+    final_bot_resp_acc, insights_used_parsed = "", []
+    temp_dl_file_path = None
 
-def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list):
-    if not service_name: raise gr.Error("Please select a target server.")
-    public_key = next((e['public_key'] for e in available_endpoints if e['name'] == service_name), None)
-    if not public_key: raise gr.Error(f"Could not find public key for '{service_name}'.")
     try:
-        created_image = create_encrypted_image(secret_data, public_key)
-        return created_image, f"✅ Success! Image created for '{service_name}'."
+        processor_gen = orchestrate_and_respond(
+            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,
+            # Pass the new stateful parameters
+            turn_budget=int(turn_budget),
+            orchestrator_state=orchestrator_state
+        )
+        curr_bot_disp_msg = ""
+        for upd_type, upd_data in processor_gen:
+            if upd_type == "state":
+                latest_orchestrator_state = upd_data
+                continue # This is a background update, no need to yield to UI yet
+            
+            if upd_type == "status":
+                status_txt = upd_data
+                if updated_gr_hist[-1][0] == effective_user_msg:
+                    updated_gr_hist[-1] = (effective_user_msg, 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[-1][0] == effective_user_msg:
+                    updated_gr_hist[-1] = (effective_user_msg, 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 generated. Processing learning..."
+                if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
+                if updated_gr_hist[-1][0] == effective_user_msg: updated_gr_hist[-1] = (effective_user_msg, curr_bot_disp_msg or "(No text)")
+                def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg, interactive=True, show_copy_button=True)
+                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)
+                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=bool(insights_used_parsed))
+
+            yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json, latest_orchestrator_state)
+            if upd_type == "final_response_and_insights": break
     except Exception as e:
-        return None, f"❌ Error: {e}"
-
-def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: list):
-    if not service_name: raise gr.Error("Please select a target server.")
-    if image is None: raise gr.Error("Please upload an image to send.")
-    api_endpoint = next((e.get('api_endpoint') for e in available_endpoints if e['name'] == service_name), None)
-    if not api_endpoint: raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
-    status = f"Connecting to endpoint: {api_endpoint}"
-    yield None, status
-    try:
-        with io.BytesIO() as buffer:
-            image.save(buffer, format="PNG")
-            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        payload = {"data": [b64_string]}
-        headers = {"Content-Type": "application/json"}
-        response = requests.post(api_endpoint, headers=headers, json=payload, timeout=45)
-        response_json = response.json()
-        if response.status_code == 200:
-            if "data" in response_json:
-                yield response_json["data"][0], "✅ Success! Data decrypted by remote server."
-            else:
-                raise gr.Error(f"API returned an unexpected success format: {response_json}")
+        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]}"
+        if updated_gr_hist[-1][0] == effective_user_msg: updated_gr_hist[-1] = (effective_user_msg, error_message_for_chat)
+        else: updated_gr_hist.append((effective_user_msg, error_message_for_chat))
+        latest_orchestrator_state = {} # Reset state on error
+        yield (cleared_input, updated_gr_hist, status_txt, gr.Markdown(value="*Error processing request.*", visible=True), gr.Textbox(value=error_message_for_chat, interactive=True), def_dl_btn, ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn(), latest_orchestrator_state)
+        if temp_dl_file_path and os.path.exists(temp_dl_file_path): os.unlink(temp_dl_file_path)
+        return
+
+    # Don't add partial/continue messages to long-term history
+    if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
+        # Add the initial user prompt and the final bot response to history
+        if orchestrator_state.get('is_complex_task'):
+             # For a complex task, the history entry should be the original goal
+            user_prompt_for_history = orchestrator_state.get('original_goal', user_msg_txt)
+            current_chat_session_history.extend([{"role": "user", "content": user_prompt_for_history}, {"role": "assistant", "content": final_bot_resp_acc}])
         else:
-            raise gr.Error(f"API Error (Status {response.status_code}): {response_json.get('error', 'Unknown error')}")
-    except Exception as e:
-        yield None, f"❌ Error calling server API: {e}"
-
-def refresh_and_update_all():
-    for dropdown_update, status_update, state_update in get_server_list():
-        pass
-    return dropdown_update, dropdown_update, status_update, state_update
-
-
-# ==============================================================================
-#  GRADIO DASHBOARD INTERFACE
-# ==============================================================================
-# ==============================================================================
-#  GRADIO DASHBOARD INTERFACE
-# ==============================================================================
-theme = gr.themes.Base(
-    primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
-    font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
-).set(
-    body_background_fill="#F1F5F9", panel_background_fill="white", block_background_fill="white",
-    block_border_width="1px", block_shadow="*shadow_drop_lg",
-    button_primary_background_fill="*primary_600", button_primary_background_fill_hover="*primary_700",
-)
+             current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
 
-with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
-    # This State component will hold the list of dicts fetched from the JSON config
-    endpoints_state = gr.State([])
-
-    gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    gr.Markdown("A centralized dashboard to manage and demonstrate the entire KeyLock ecosystem. Key/Image creation is performed locally, while decryption is handled by a **live, remote API call** to a secure server.")
-
-    with gr.Tabs() as tabs:
-        with gr.TabItem("① Create KeyLock", id=0):
-            gr.Markdown("## Step 1: Create an Encrypted Authentication Image (Local)")
-            gr.Markdown(f"This tool acts as the **Auth Creator**. It fetches a list of available servers from a public [configuration file]({CREATOR_ENDPOINTS_JSON_URL}), then uses the selected server's public key to encrypt your data into a PNG. **This entire creation process happens locally.**")
-            with gr.Row(variant="panel"):
-                with gr.Column(scale=2):
-                    with gr.Row():
-                        creator_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to encrypt data for.")
-                        refresh_button = gr.Button("🔄", scale=0, size="sm")
-                    creator_secret_input = gr.Textbox(lines=8, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
-                    creator_button = gr.Button("✨ Create Auth Image", variant="primary")
-                with gr.Column(scale=1):
-                    creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png")
-
-        with gr.TabItem("② Send KeyLock", id=1):
-            gr.Markdown("## Step 2: Decrypt via Live API Call")
-            gr.Markdown("This tool acts as the **Client**. It sends the encrypted image you created in Step 1 to the live, remote **Decoder Server** you select from the same configuration list. The server uses its securely stored private key to decrypt the data and sends the result back.")
-            with gr.Row(variant="panel"):
-                with gr.Column(scale=1):
-                    gr.Markdown("### Configuration")
-                    send_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to send the image to.")
-                    gr.Markdown("### Image to Send")
-                    client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image Here", sources=["upload", "clipboard"])
-                    client_button = gr.Button("🔓 Decrypt via Remote Server", variant="primary")
-                with gr.Column(scale=1):
-                    gr.Markdown("### Response from Server")
-                    client_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    client_json_output = gr.JSON(label="Decrypted Data")
+        if len(current_chat_session_history) > MAX_HISTORY_TURNS * 2: 
+            current_chat_session_history = current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]
         
-        with gr.TabItem("ℹ️ Info & Key Generation", id=2):
-            gr.Markdown("## Ecosystem Architecture")
-            gr.Markdown(f"This dashboard uses a public [configuration file]({CREATOR_ENDPOINTS_JSON_URL}) to dynamically discover and interact with live services. It demonstrates a secure, decoupled workflow.")
-            with gr.Row():
-                 with gr.Column():
-                    gr.Markdown(f"### 🏭 Auth Creator Service\n- **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})\n- **Role:** Provides an API to encrypt data for various targets defined in its `endpoints.json` file.\n- **Source Code:** [app.py]({CREATOR_APP_PY_URL})")
-                 with gr.Column():
-                    gr.Markdown(f"### 📡 Decoder Server\n- **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})\n- **Role:** The trusted authority. It holds a secret private key and provides a secure API to decrypt images.\n- **Source Code:** [app.py]({SERVER_APP_PY_URL})")
-            
-            with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
-                gr.Markdown("Create a new key pair. In a real scenario, you would add the **Public Key** and the server's **API Endpoint URL** to the `endpoints.json` configuration file, and set the **Private Key** as a secret variable in the corresponding server space.")
-                with gr.Row():
-                    with gr.Column():
-                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
-                    with gr.Column():
-                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key", interactive=False, show_copy_button=True)
-                gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair", variant="secondary")
-
-    # --- Wire up the component logic ---
-    gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
+        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, ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn(), latest_orchestrator_state)
+        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."
+    else: status_txt = final_bot_resp_acc or "Processing finished."
+    
+    current_chat_session_history = [h for h in current_chat_session_history if h['content']] # Clean up empty entries
+
+    updated_rules_text = ui_refresh_rules_display_fn()
+    updated_mems_json = ui_refresh_memories_display_fn()
+    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json, latest_orchestrator_state)
+
+    if temp_dl_file_path and os.path.exists(temp_dl_file_path): os.unlink(temp_dl_file_path)
+
+# --- Keep all your UI helper functions (ui_refresh_rules_display_fn, etc.) unchanged ---
+def ui_refresh_rules_display_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 gr.DownloadButton(value=None, interactive=False, label="No Rules")
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt", encoding='utf-8') as tmpfile:
+        tmpfile.write(rules_content)
+        return tmpfile.name
+def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
+    if not uploaded_file_obj: return "No file provided."
+    added, skipped, errors = load_rules_from_file(uploaded_file_obj.name, progress_callback=lambda p, d: progress(p, desc=d))
+    return f"Rules Upload: Added: {added}, Skipped (duplicates): {skipped}, Errors: {errors}."
+
+def ui_refresh_memories_display_fn(): return get_all_memories_cached() or []
+def ui_download_memories_action_fn():
+    memories = get_all_memories_cached()
+    if not memories:
+        gr.Warning("No memories to download.")
+        return gr.DownloadButton(value=None, interactive=False, label="No Memories")
+    jsonl_content = "\n".join([json.dumps(mem) for mem in memories])
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".jsonl", encoding='utf-8') as tmpfile:
+        tmpfile.write(jsonl_content)
+        return tmpfile.name
+def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
+    if not uploaded_file_obj: return "No file provided."
+    added, format_err, save_err = load_memories_from_file(uploaded_file_obj.name, progress_callback=lambda p, d: progress(p, desc=d))
+    return f"Memories Upload: Added: {added}, Format Errors: {format_err}, Save Errors: {save_err}."
+
+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."
+    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()]
+    unique_rules = sorted(list(set(filter(None, [r.strip() for r in potential_rules]))))
+    if not unique_rules: return "No unique, non-empty rules found."
+    added, skipped, errors, total = 0, 0, 0, len(unique_rules)
+    progress(0, desc=f"Saving {total} unique rules...")
+    for idx, rule_text in enumerate(unique_rules):
+        success, status_msg = add_rule_entry(rule_text)
+        if success: added += 1
+        elif status_msg == "duplicate": skipped += 1
+        else: errors += 1
+        progress((idx + 1) / total, desc=f"Processed {idx+1}/{total} rules...")
+    return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors: {errors} from {total} unique rules."
+
+def app_load_fn():
+    logger.info("App loading. Initializing systems...")
+    initialize_memory_system()
+    rules_added, rules_skipped, rules_errors = load_rules_from_file(LOAD_RULES_FILE)
+    mems_added, mems_format_errors, mems_save_errors = load_memories_from_file(LOAD_MEMORIES_FILE)
+    status = f"Ready. Rules loaded: {rules_added}. Memories loaded: {mems_added}."
+    return (status, ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn(), gr.Markdown(visible=False), gr.Textbox(value="*Waiting...*", interactive=True), gr.DownloadButton(interactive=False, visible=False))
+
+
+with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .status-text { font-size: 0.9em; color: #555; }") as demo:
+    # This component holds the state of the orchestrator between runs
+    orchestrator_state = gr.State({})
+
+    gr.Markdown("# 🤖 AI Research Agent")
+    with gr.Row(variant="compact"):
+        agent_stat_tb = gr.Textbox(label="Agent Status", value="Initializing...", 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)
+            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")
+    with gr.Row():
+        with gr.Sidebar():
+            gr.Markdown("## ⚙️ Configuration")
+            with gr.Group():
+                api_key_tb = gr.Textbox(label="API Key (Override)", type="password", placeholder="Uses .env if blank")
+                available_providers = get_available_providers()
+                default_provider = available_providers[0] if available_providers else None
+                prov_sel_dd = gr.Dropdown(label="AI Provider", choices=available_providers, value=default_provider, interactive=True)
+                model_sel_dd = gr.Dropdown(label="AI Model", choices=get_model_display_names_for_provider(default_provider) if default_provider else [], value=get_default_model_display_name_for_provider(default_provider), interactive=True)
+                # New slider for controlling the agent's loop
+                turn_budget_slider = gr.Slider(label="Max Steps Per Turn (0=Continuous)", minimum=0, maximum=20, step=1, value=5, interactive=True)
+            with gr.Group():
+                sys_prompt_tb = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
+            if MEMORY_STORAGE_BACKEND == "RAM":
+                save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")
+        with gr.Column(scale=3):
+            with gr.Tabs():
+                with gr.TabItem("💬 Chat & Research"):
+                    main_chat_disp = gr.Chatbot(height=400, show_copy_button=True, render_markdown=True)
+                    with gr.Row(variant="compact"):
+                        user_msg_tb = gr.Textbox(show_label=False, placeholder="Ask your research question...", scale=7, lines=1)
+                        send_btn = gr.Button("Send / Continue", variant="primary", scale=1, min_width=120)
+                    with gr.Accordion("📝 Detailed Response & Insights", 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", value=None, 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="Current Rules", lines=10, interactive=True)
+                            save_edited_rules_btn = gr.Button("💾 Save Edited Text", variant="primary")
+                            with gr.Row(variant="compact"):
+                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules")
+                                clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop")
+                            upload_rules_fobj = gr.File(label="Upload Rules File (.txt/.jsonl)", file_types=[".txt", ".jsonl"])
+                            rules_stat_tb = gr.Textbox(label="Rules Status", interactive=False, lines=1)
+                        with gr.Column():
+                            gr.Markdown("### 📚 Memories Management")
+                            mems_disp_json = gr.JSON(label="Current Memories", value=[])
+                            with gr.Row(variant="compact"):
+                                dl_mems_btn = gr.DownloadButton("⬇️ Download Memories")
+                                clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop")
+                            upload_mems_fobj = gr.File(label="Upload Memories File (.json/.jsonl)", file_types=[".json", ".jsonl"])
+                            mems_stat_tb = gr.Textbox(label="Memories Status", interactive=False, lines=1)
+
+    prov_sel_dd.change(lambda p: gr.Dropdown(choices=get_model_display_names_for_provider(p), value=get_default_model_display_name_for_provider(p), interactive=True), prov_sel_dd, model_sel_dd)
+    
+    # Update inputs and outputs for the chat event to include the new stateful components
+    chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb, turn_budget_slider, orchestrator_state]
+    chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn, rules_disp_ta, mems_disp_json, orchestrator_state]
     
-    def refresh_and_update_all():
-        # The generator must be iterated to get its final values
-        # The 'pass' allows the loop to complete, and we use the final yielded values.
-        for dropdown_update, status_update, state_update in get_server_list():
-            pass
-        # Update both dropdowns and the state
-        return dropdown_update, dropdown_update, status_update, state_update
+    chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
+    send_btn.click(**chat_event_args)
+    user_msg_tb.submit(**chat_event_args)
+
+    # Keep all other button clicks the same
+    dl_rules_btn.click(ui_download_rules_action_fn, None, dl_rules_btn)
+    save_edited_rules_btn.click(save_edited_rules_action_fn, [rules_disp_ta], [rules_stat_tb]).then(ui_refresh_rules_display_fn, outputs=rules_disp_ta)
+    upload_rules_fobj.upload(ui_upload_rules_action_fn, [upload_rules_fobj], [rules_stat_tb]).then(ui_refresh_rules_display_fn, outputs=rules_disp_ta)
+    clear_rules_btn.click(lambda: ("Cleared." if clear_all_rules_data_backend() else "Error."), outputs=rules_stat_tb).then(ui_refresh_rules_display_fn, outputs=rules_disp_ta)
+
+    dl_mems_btn.click(ui_download_memories_action_fn, None, dl_mems_btn)
+    upload_mems_fobj.upload(ui_upload_memories_action_fn, [upload_mems_fobj], [mems_stat_tb]).then(ui_refresh_memories_display_fn, outputs=mems_disp_json)
+    clear_mems_btn.click(lambda: ("Cleared." if clear_all_memory_data_backend() else "Error."), outputs=mems_stat_tb).then(ui_refresh_memories_display_fn, outputs=mems_disp_json)
 
-    refresh_button.click(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
-    demo.load(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
+    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
+        save_faiss_sidebar_btn.click(lambda: (gr.Info("Saved FAISS to disk.") if save_faiss_indices_to_disk() is None else gr.Error("Error saving FAISS.")), None, None)
 
-    creator_button.click(fn=create_keylock_wrapper, inputs=[creator_service_dropdown, creator_secret_input, endpoints_state], outputs=[creator_image_output, creator_status])
-    client_button.click(fn=send_keylock_wrapper, inputs=[send_service_dropdown, client_image_input, endpoints_state], outputs=[client_json_output, client_status])
+    app_load_outputs = [agent_stat_tb, rules_disp_ta, mems_disp_json, detect_out_md, fmt_report_tb, dl_report_btn]
+    demo.load(fn=app_load_fn, inputs=None, outputs=app_load_outputs)
 
 if __name__ == "__main__":
-    demo.launch()
+    app_port = int(os.getenv("GRADIO_PORT", 7860))
+    app_server = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
+    logger.info(f"Launching Gradio server: http://{app_server}:{app_port}")
+    demo.queue().launch(server_name=app_server, server_port=app_port)