Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os
 import json
 import re
@@ -12,26 +13,29 @@ import xml.etree.ElementTree as ET
 load_dotenv()
 
 MEMORY_STORAGE_TYPE = "HF_DATASET" 
+# --- DEMO --- #
+# If using HF_DATASET, specify the repository names here.
+# These will override the .env file settings.
 HF_DATASET_MEMORY_REPO = "broadfield-dev/ai-brain"
 HF_DATASET_RULES_REPO = "broadfield-dev/ai-rules"
 
+# Set environment variables based on the toggles above BEFORE importing other modules
 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
+# --- END CONFIGURATION ---
 
 from model_logic import (
     get_available_providers, get_model_display_names_for_provider,
     get_default_model_display_name_for_provider, call_model_stream, MODELS_BY_PROVIDER
 )
 from memory_logic import (
-    initialize_memory_system, search_memories, add_memory_entry,
-    retrieve_rules_semantic, get_all_rules_cached, get_all_memories_cached,
-    add_rule_entry, remove_rule_entry, clear_all_memory_data_backend, 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
+    initialize_memory_system,
+    add_memory_entry, retrieve_memories_semantic, get_all_memories_cached, clear_all_memory_data_backend,
+    add_rule_entry, retrieve_rules_semantic, 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 websearch_logic import scrape_url, search_and_scrape_duckduckgo, search_and_scrape_google
 
@@ -51,11 +55,17 @@ DEFAULT_SYSTEM_PROMPT = os.getenv(
     "Your Name is Node. You are a Helpful AI Assistant, and your goal is to improve yourself, so that you can improve everybody elses life. Your job is to answer questions and perform research tasks accurately and thoroughly. You can use tools like web search and page browsing. When providing information from the web, cite your sources if possible. If asked to perform a task beyond your capabilities, explain politely. Be concise unless asked for detail."
 )
 
+# --- Startup Loading File Paths ---
 LOAD_RULES_FILE = os.getenv("LOAD_RULES_FILE")
 LOAD_MEMORIES_FILE = os.getenv("LOAD_MEMORIES_FILE")
+logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={TOOL_DECISION_PROVIDER_ENV}, ToolDecisionModelID={TOOL_DECISION_MODEL_ID_ENV}, MemoryBackend={MEMORY_STORAGE_BACKEND}")
+logger.info(f"Startup loading: Rules from {LOAD_RULES_FILE or 'None'}, Memories from {LOAD_MEMORIES_FILE or 'None'}")
 
+
+# --- Helper Functions ---
 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.", []
+    if not retrieved_insights_list:
+        return "No specific guiding principles or learned insights retrieved.", []
     parsed = []
     for text in retrieved_insights_list:
         match = re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\](.*)", text.strip(), re.DOTALL | re.IGNORECASE)
@@ -65,124 +75,144 @@ def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str,
             parsed.append({"type": "GENERAL_LEARNING", "score": "0.5", "text": text.strip(), "original": text.strip()})
     try:
         parsed.sort(key=lambda x: float(x["score"]) if x["score"].replace('.', '', 1).isdigit() else -1.0, reverse=True)
-    except ValueError: pass
+    except ValueError: logger.warning("FORMAT_INSIGHTS: Sort error due to invalid score format.")
     grouped = {"CORE_RULE": [], "RESPONSE_PRINCIPLE": [], "BEHAVIORAL_ADJUSTMENT": [], "GENERAL_LEARNING": []}
     for p_item in parsed: grouped.get(p_item["type"], grouped["GENERAL_LEARNING"]).append(f"- (Score: {p_item['score']}) {p_item['text']}")
     sections = [f"{k.replace('_', ' ').title()}:\n" + "\n".join(v) for k, v in grouped.items() if v]
     return "\n\n".join(sections) if sections else "No guiding principles retrieved.", parsed
 
 def generate_interaction_metrics(user_input: str, bot_response: str, provider: str, model_display_name: str, api_key_override: str = None) -> dict:
-    metric_messages = [
-        {"role": "system", "content": "You are a precise JSON output agent. Output a single JSON object containing interaction metrics as requested. Do not include any explanatory text."},
-        {"role": "user", "content": f"User: \"{user_input}\"\nAI: \"{bot_response}\"\nMetrics: \"takeaway\" (3-7 words), \"response_success_score\" (0.0-1.0), \"future_confidence_score\" (0.0-1.0). Output JSON ONLY."}
-    ]
+    metric_start_time = time.time()
+    logger.info(f"Generating metrics with: {provider}/{model_display_name}")
+    metric_prompt_content = f"User: \"{user_input}\"\nAI: \"{bot_response}\"\nMetrics: \"takeaway\" (3-7 words), \"response_success_score\" (0.0-1.0), \"future_confidence_score\" (0.0-1.0). Output JSON ONLY, ensure it's a single, valid JSON object."
+    metric_messages = [{"role": "system", "content": "You are a precise JSON output agent. Output a single JSON object containing interaction metrics as requested by the user. Do not include any explanatory text before or after the JSON object."}, {"role": "user", "content": metric_prompt_content}]
     try:
+        metrics_provider_final, metrics_model_display_final = provider, model_display_name
         metrics_model_env = os.getenv("METRICS_MODEL")
         if metrics_model_env and "/" in metrics_model_env:
             m_prov, m_id = metrics_model_env.split('/', 1)
             m_disp_name = next((dn for dn, mid in MODELS_BY_PROVIDER.get(m_prov.lower(), {}).get("models", {}).items() if mid == m_id), None)
-            if m_disp_name: provider, model_display_name = m_prov, m_disp_name
-        
-        resp_str = "".join(list(call_model_stream(provider=provider, model_display_name=model_display_name, messages=metric_messages, api_key_override=api_key_override, temperature=0.05, max_tokens=200)))
-        json_match = re.search(r"\{.*\}", resp_str, re.DOTALL)
-        if json_match:
-            metrics_data = json.loads(json_match.group(0))
-            return {"takeaway": metrics_data.get("takeaway", "N/A"), "response_success_score": float(metrics_data.get("response_success_score", 0.5)), "future_confidence_score": float(metrics_data.get("future_confidence_score", 0.5))}
+            if m_disp_name: metrics_provider_final, metrics_model_display_final = m_prov, m_disp_name
+            else: logger.warning(f"METRICS_MODEL '{metrics_model_env}' not found, using interaction model.")
+        response_chunks = list(call_model_stream(provider=metrics_provider_final, model_display_name=metrics_model_display_final, messages=metric_messages, api_key_override=api_key_override, temperature=0.05, max_tokens=200))
+        resp_str = "".join(response_chunks).strip()
+        json_match = re.search(r"```json\s*(\{.*?\})\s*```", resp_str, re.DOTALL | re.IGNORECASE) or re.search(r"(\{.*?\})", resp_str, re.DOTALL)
+        if json_match: metrics_data = json.loads(json_match.group(1))
+        else:
+            logger.warning(f"METRICS_GEN: Non-JSON response from {metrics_provider_final}/{metrics_model_display_final}: '{resp_str}'")
+            return {"takeaway": "N/A", "response_success_score": 0.5, "future_confidence_score": 0.5, "error": "metrics format error"}
+        parsed_metrics = {"takeaway": metrics_data.get("takeaway", "N/A"), "response_success_score": float(metrics_data.get("response_success_score", 0.5)), "future_confidence_score": float(metrics_data.get("future_confidence_score", 0.5)), "error": metrics_data.get("error")}
+        logger.info(f"METRICS_GEN: Generated in {time.time() - metric_start_time:.2f}s. Data: {parsed_metrics}")
+        return parsed_metrics
     except Exception as e:
-        logger.error(f"METRICS_GEN Error: {e}")
-    return {"takeaway": "N/A", "response_success_score": 0.5, "future_confidence_score": 0.5, "error": "metrics format error"}
+        logger.error(f"METRICS_GEN Error: {e}", exc_info=False)
+        return {"takeaway": "N/A", "response_success_score": 0.5, "future_confidence_score": 0.5, "error": str(e)}
+
 
 def process_user_interaction_gradio(user_input: str, provider_name: str, model_display_name: str, chat_history_for_prompt: list[dict], custom_system_prompt: str = None, ui_api_key_override: str = None):
+    process_start_time = time.time()
     request_id = os.urandom(4).hex()
-    logger.info(f"PUI_GRADIO [{request_id}] Start. User: '{user_input[:50]}...'")
+    logger.info(f"PUI_GRADIO [{request_id}] Start. User: '{user_input[:50]}...' Provider: {provider_name}/{model_display_name} Hist_len:{len(chat_history_for_prompt)}")
     history_str_for_prompt = "\n".join([f"{('User' if t_msg['role'] == 'user' else 'AI')}: {t_msg['content']}" for t_msg in chat_history_for_prompt[-(MAX_HISTORY_TURNS * 2):]])
-    yield "status", "<i>[Checking guidelines...]</i>"
-    initial_insights, parsed_initial_insights_list = format_insights_for_prompt(retrieve_rules_semantic(f"{user_input}\n{history_str_for_prompt}", k=5))
-
-    action_type = "quick_respond"
-    action_input_dict = {}
-
-    if not (len(user_input.split()) <= 3 and any(kw in user_input.lower() for kw in ["hello", "hi", "thanks", "ok", "bye"]) and "?" not in user_input):
+    yield "status", "<i>[Checking guidelines (semantic search)...]</i>"
+    initial_insights = retrieve_rules_semantic(f"{user_input}\n{history_str_for_prompt}", k=5)
+    initial_insights_ctx_str, parsed_initial_insights_list = format_insights_for_prompt(initial_insights)
+    logger.info(f"PUI_GRADIO [{request_id}]: Initial RAG (insights) found {len(initial_insights)}. Context: {initial_insights_ctx_str[:150]}...")
+    action_type, action_input_dict = "quick_respond", {}
+    user_input_lower = user_input.lower()
+    time_before_tool_decision = time.time()
+    if WEB_SEARCH_ENABLED and ("http://" in user_input or "https://" in user_input):
+        url_match = re.search(r'(https?://[^\s]+)', user_input)
+        if url_match: action_type, action_input_dict = "scrape_url_and_report", {"url": url_match.group(1)}
+    if action_type == "quick_respond" and len(user_input.split()) <= 3 and any(kw in user_input_lower for kw in ["hello", "hi", "thanks", "ok", "bye"]) and not "?" in user_input: pass
+    elif action_type == "quick_respond" and WEB_SEARCH_ENABLED and (len(user_input.split()) > 3 or "?" in user_input or any(w in user_input_lower for w in ["what is", "how to", "explain", "search for"])):
         yield "status", "<i>[LLM choosing best approach...]</i>"
-        tool_definitions = {
-            "answer_using_conversation_memory": "Use if the user's query refers to a past conversation, asks you to 'remember' or 'recall' something, or seems to require knowledge you might have gained previously.",
-            "search_duckduckgo_and_report": "Use for general knowledge questions, current events, or explicit requests to search the web.",
-            "quick_respond": "Use as a fallback for simple greetings or if the answer is obvious from immediate context."
-        }
-        available_tools = ["quick_respond", "answer_using_conversation_memory"] + (["search_duckduckgo_and_report"] if WEB_SEARCH_ENABLED else [])
-        tool_desc = "\n".join(f'- "{name}": {tool_definitions[name]}' for name in available_tools)
-        tool_sys_prompt = "You are a precise routing agent. Analyze the user's query and select the single best action. Output ONLY a valid JSON object like {\"action\": \"action_name\", \"action_input\": {\"param\": \"value\"}}."
-        tool_user_prompt = f"User Query: \"{user_input}\"\n\nAvailable Actions:\n{tool_desc}\n\nSelect one action."
-        
+        history_snippet = "\n".join([f"{msg['role']}: {msg['content'][:100]}" for msg in chat_history_for_prompt[-2:]])
+        guideline_snippet = initial_insights_ctx_str[:200].replace('\n', ' ')
+        tool_sys_prompt = "You are a precise routing agent... Output JSON only. Example: {\"action\": \"search_duckduckgo_and_report\", \"action_input\": {\"search_engine_query\": \"query\"}}"
+        tool_user_prompt = f"User Query: \"{user_input}\nRecent History:\n{history_snippet}\nGuidelines: {guideline_snippet}...\nAvailable Actions: quick_respond, answer_using_conversation_memory, search_duckduckgo_and_report, scrape_url_and_report.\nSelect one action and input. Output JSON."
+        tool_decision_messages = [{"role":"system", "content": tool_sys_prompt}, {"role":"user", "content": tool_user_prompt}]
         tool_provider, tool_model_id = TOOL_DECISION_PROVIDER_ENV, TOOL_DECISION_MODEL_ID_ENV
         tool_model_display = next((dn for dn, mid in MODELS_BY_PROVIDER.get(tool_provider.lower(), {}).get("models", {}).items() if mid == tool_model_id), None)
         if not tool_model_display: tool_model_display = get_default_model_display_name_for_provider(tool_provider)
-
-        try:
-            tool_resp_raw = "".join(list(call_model_stream(provider=tool_provider, model_display_name=tool_model_display, messages=[{"role":"system", "content": tool_sys_prompt}, {"role":"user", "content": tool_user_prompt}], temperature=0.0, max_tokens=150)))
-            json_match_tool = re.search(r"\{.*\}", tool_resp_raw, re.DOTALL)
-            if json_match_tool:
-                action_data = json.loads(json_match_tool.group(0))
-                action_type = action_data.get("action", "quick_respond")
-                action_input_dict = action_data.get("action_input", {})
-        except Exception as e:
-            logger.error(f"Tool decision error: {e}")
-            action_type = "quick_respond"
-    
+        if tool_model_display:
+            try:
+                logger.info(f"PUI_GRADIO [{request_id}]: Tool decision LLM: {tool_provider}/{tool_model_display}")
+                tool_resp_chunks = list(call_model_stream(provider=tool_provider, model_display_name=tool_model_display, messages=tool_decision_messages, temperature=0.0, max_tokens=150))
+                tool_resp_raw = "".join(tool_resp_chunks).strip()
+                json_match_tool = re.search(r"\{.*\}", tool_resp_raw, re.DOTALL)
+                if json_match_tool:
+                    action_data = json.loads(json_match_tool.group(0))
+                    action_type, action_input_dict = action_data.get("action", "quick_respond"), action_data.get("action_input", {})
+                    if not isinstance(action_input_dict, dict): action_input_dict = {}
+                    logger.info(f"PUI_GRADIO [{request_id}]: LLM Tool Decision: Action='{action_type}', Input='{action_input_dict}'")
+                else: logger.warning(f"PUI_GRADIO [{request_id}]: Tool decision LLM non-JSON. Raw: {tool_resp_raw}")
+            except Exception as e: logger.error(f"PUI_GRADIO [{request_id}]: Tool decision LLM error: {e}", exc_info=False)
+        else: logger.error(f"No model for tool decision provider {tool_provider}.")
+    elif action_type == "quick_respond" and not WEB_SEARCH_ENABLED and (len(user_input.split()) > 4 or "?" in user_input or any(w in user_input_lower for w in ["remember","recall"])):
+        action_type="answer_using_conversation_memory"
+    logger.info(f"PUI_GRADIO [{request_id}]: Tool decision logic took {time.time() - time_before_tool_decision:.3f}s. Action: {action_type}, Input: {action_input_dict}")
     yield "status", f"<i>[Path: {action_type}. Preparing response...]</i>"
-    final_system_prompt_str = custom_system_prompt or DEFAULT_SYSTEM_PROMPT
-    
-    if action_type == "answer_using_conversation_memory":
-        yield "status", "<i>[Optimizing query for memory search...]</i>"
-        optimized_query = user_input
-        try:
-            query_gen_sys = "Reformulate the user's question into a concise, self-contained search query for a vector database. Output ONLY the query text."
-            query_gen_user = f"History:\n{history_str_for_prompt}\n\nUser Query: \"{user_input}\"\n\nOptimized Query:"
-            generated_query = "".join(list(call_model_stream(provider=tool_provider, model_display_name=tool_model_display, messages=[{"role":"system", "content":query_gen_sys}, {"role":"user", "content":query_gen_user}], temperature=0.0, max_tokens=50))).strip().replace('"', '')
-            if generated_query: optimized_query = generated_query
-        except Exception: pass
-        logger.info(f"Optimized memory search query: '{optimized_query}'")
-
-        yield "status", "<i>[Searching memories...]</i>"
-        retrieved_mems, search_path = search_memories(query=optimized_query, k=3, threshold=1.0)
-        
-        if search_path == "deep":
-            yield "status", "<i>[No recent match. Performing deep search...]</i>"
-
-        memory_context = "No relevant past interactions found."
-        if retrieved_mems:
-            memory_context = "Relevant Past Interactions:\n" + "\n".join([f"- User asked: '{m.get('user_input','')}'. You responded: '{m.get('bot_response','')}'. (Takeaway: {m.get('metrics',{}).get('takeaway','N/A')})" for m in retrieved_mems])
-        
-        final_system_prompt_str += " You MUST use the 'Memory Context' to inform your answer."
-        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\n\nGuidelines:\n{initial_insights}\n\nMemory Context:\n{memory_context}\n\nUser Query: \"{user_input}\"\n\nResponse:"
-    elif action_type == "search_duckduckgo_and_report":
-        query_or_url = action_input_dict.get("search_engine_query", user_input)
-        yield "status", f"<i>[Web: '{query_or_url[:60]}'...]</i>"
-        web_results = search_and_scrape_duckduckgo(query_or_url, num_results=2)
-        scraped_content = "\n".join([f"Source {i+1}:\nURL:{r.get('url','N/A')}\nTitle:{r.get('title','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:3500]}\n---" for i,r in enumerate(web_results)]) if web_results else f"No results for '{query_or_url}'."
-        yield "status", "<i>[Synthesizing web report...]</i>"
-        final_system_prompt_str += " Generate a report from web content, citing URLs as [Source X]."
-        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights}\nWeb Content:\n{scraped_content}\nQuery: \"{user_input}\"\nReport/Response:"
+    final_system_prompt_str, final_user_prompt_content_str = custom_system_prompt or DEFAULT_SYSTEM_PROMPT, ""
+    if action_type == "quick_respond":
+        final_system_prompt_str += " Respond directly using guidelines & history."
+        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
+    elif action_type == "answer_using_conversation_memory":
+        yield "status", "<i>[Searching conversation memory (semantic)...]</i>"
+        retrieved_mems = retrieve_memories_semantic(f"User query: {user_input}\nContext:\n{history_str_for_prompt[-1000:]}", k=2)
+        memory_context = "Relevant Past Interactions:\n" + "\n".join([f"- User:{m.get('user_input','')}->AI:{m.get('bot_response','')} (Takeaway:{m.get('metrics',{}).get('takeaway','N/A')})" for m in retrieved_mems]) if retrieved_mems else "No relevant past interactions found."
+        final_system_prompt_str += " Respond using Memory Context, guidelines, & history."
+        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nMemory Context:\n{memory_context}\nQuery: \"{user_input}\"\nResponse (use memory context if relevant):"
+    elif WEB_SEARCH_ENABLED and action_type in ["search_duckduckgo_and_report", "scrape_url_and_report"]:
+        query_or_url = action_input_dict.get("search_engine_query") if "search" in action_type else action_input_dict.get("url")
+        if not query_or_url:
+            final_system_prompt_str += " Respond directly (web action failed: no input)."
+            final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
+        else:
+            yield "status", f"<i>[Web: '{query_or_url[:60]}'...]</i>"
+            web_results, max_results = [], 1 if action_type == "scrape_url_and_report" else 2
+            try:
+                if action_type == "search_duckduckgo_and_report": web_results = search_and_scrape_duckduckgo(query_or_url, num_results=max_results)
+                elif action_type == "scrape_url_and_report":
+                    res = scrape_url(query_or_url)
+                    if res and (res.get("content") or res.get("error")): web_results = [res]
+            except Exception as e: web_results = [{"url": query_or_url, "title": "Tool Error", "error": str(e)}]
+            scraped_content = "\n".join([f"Source {i+1}:\nURL:{r.get('url','N/A')}\nTitle:{r.get('title','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:3500]}\n---" for i,r in enumerate(web_results)]) if web_results else f"No results from {action_type} for '{query_or_url}'."
+            yield "status", "<i>[Synthesizing web report...]</i>"
+            final_system_prompt_str += " Generate report/answer from web content, history, & guidelines. Cite URLs as [Source X]."
+            final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nWeb Content:\n{scraped_content}\nQuery: \"{user_input}\"\nReport/Response (cite sources [Source X]):"
     else:
-        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights}\nQuery: \"{user_input}\"\nResponse:"
-
+        final_system_prompt_str += " Respond directly (unknown action path)."
+        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
     final_llm_messages = [{"role": "system", "content": final_system_prompt_str}, {"role": "user", "content": final_user_prompt_content_str}]
-    streamed_response = ""
-    for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_llm_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=2500):
-        streamed_response += chunk
-        yield "response_chunk", chunk
-    
-    final_bot_text = streamed_response.strip() or "(No response.)"
+    logger.debug(f"PUI_GRADIO [{request_id}]: Final LLM System Prompt: {final_system_prompt_str[:200]}...")
+    logger.debug(f"PUI_GRADIO [{request_id}]: Final LLM User Prompt Start: {final_user_prompt_content_str[:200]}...")
+    streamed_response, time_before_llm = "", time.time()
+    try:
+        for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_llm_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=2500):
+            if isinstance(chunk, str) and chunk.startswith("Error:"): streamed_response += f"\n{chunk}\n"; yield "response_chunk", f"\n{chunk}\n"; break
+            streamed_response += chunk; yield "response_chunk", chunk
+    except Exception as e: streamed_response += f"\n\n(Error: {str(e)[:150]})"; yield "response_chunk", f"\n\n(Error: {str(e)[:150]})"
+    logger.info(f"PUI_GRADIO [{request_id}]: Main LLM stream took {time.time() - time_before_llm:.3f}s.")
+    final_bot_text = streamed_response.strip() or "(No response or error.)"
+    logger.info(f"PUI_GRADIO [{request_id}]: Finished. Total: {time.time() - process_start_time:.2f}s. Resp len: {len(final_bot_text)}")
     yield "final_response_and_insights", {"response": final_bot_text, "insights_used": parsed_initial_insights_list}
 
 def perform_post_interaction_learning(user_input: str, bot_response: str, provider: str, model_disp_name: str, insights_reflected: list[dict], api_key_override: str = None):
     task_id = os.urandom(4).hex()
-    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: START")
-    metrics = generate_interaction_metrics(user_input, bot_response, provider, model_disp_name, api_key_override)
-    add_memory_entry(user_input, metrics, bot_response)
-    summary = f"User:\"{user_input}\"\nAI:\"{bot_response}\"\nMetrics(takeaway):{metrics.get('takeaway','N/A')},Success:{metrics.get('response_success_score','N/A')}"
-    existing_rules_ctx = "\n".join([f"- \"{r}\"" for r in retrieve_rules_semantic(f"{summary}\n{user_input}", k=10)]) or "No existing rules context."
-    insight_sys_prompt = """You are an expert AI knowledge base curator. Your primary function is to meticulously analyze an interaction and update the AI's guiding principles (insights/rules) to improve its future performance and self-understanding.
+    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
+    learning_start_time = time.time()
+    significant_learnings_summary = [] # To store summaries of new core learnings
+
+    try:
+        metrics = generate_interaction_metrics(user_input, bot_response, provider, model_disp_name, api_key_override)
+        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Metrics: {metrics}")
+        add_memory_entry(user_input, metrics, bot_response)
+
+        summary = f"User:\"{user_input}\"\nAI:\"{bot_response}\"\nMetrics(takeaway):{metrics.get('takeaway','N/A')},Success:{metrics.get('response_success_score','N/A')}"
+        existing_rules_ctx = "\n".join([f"- \"{r}\"" for r in retrieve_rules_semantic(f"{summary}\n{user_input}", k=10)]) or "No existing rules context."
+
+        insight_sys_prompt = """You are an expert AI knowledge base curator. Your primary function is to meticulously analyze an interaction and update the AI's guiding principles (insights/rules) to improve its future performance and self-understanding.
 **CRITICAL OUTPUT REQUIREMENT: You MUST output a single, valid XML structure representing a list of operation objects.**
 The root element should be `<operations_list>`. Each operation should be an `<operation>` element.
 If no operations are warranted, output an empty list: `<operations_list></operations_list>`.
@@ -191,63 +221,181 @@ Each `<operation>` element must contain the following child elements:
 1.  `<action>`: A string, either `"add"` (for entirely new rules) or `"update"` (to replace an existing rule with a better one).
 2.  `<insight>`: The full, refined insight text including its `[TYPE|SCORE]` prefix (e.g., `[CORE_RULE|1.0] My name is Lumina, an AI assistant.`). Multi-line insight text can be placed directly within this tag; XML handles newlines naturally.
 3.  `<old_insight_to_replace>`: (ONLY for `"update"` action) The *exact, full text* of an existing insight that the new `<insight>` should replace. If action is `"add"`, this element should be omitted or empty.
+**XML Structure Example:**
+<operations_list>
+  <operation>
+    <action>update</action>
+    <insight>[CORE_RULE|1.0] I am Lumina, an AI assistant.
+My purpose is to help with research.</insight>
+    <old_insight_to_replace>[CORE_RULE|0.9] My name is Assistant.</old_insight_to_replace>
+  </operation>
+  <operation>
+    <action>add</action>
+    <insight>[RESPONSE_PRINCIPLE|0.8] User prefers short answers.
+Provide details only when asked.</insight>
+  </operation>
+</operations_list>
+**Your Reflection Process (Consider each step and generate operations accordingly):**
+- **STEP 1: CORE IDENTITY/PURPOSE:** Review the interaction and existing rules. Identify if the interaction conflicts with, clarifies, or reinforces your core identity (name, fundamental nature, primary purpose). If necessary, propose updates or additions to CORE_RULEs. Aim for a single, consistent set of CORE_RULEs over time by updating older versions.
+- **STEP 2: NEW LEARNINGS:** Based *only* on the "Interaction Summary", identify concrete, factual information, user preferences, or skills demonstrated that were not previously known or captured. These should be distinct, actionable learnings. Formulate these as new [GENERAL_LEARNING] or specific [BEHAVIORAL_ADJUSTMENT] rules. Do NOT add rules that are already covered by existing relevant rules.
+- **STEP 3: REFINEMENT/ADJUSTMENT:** Review existing non-core rules ([RESPONSE_PRINCIPLE], [BEHAVIORAL_ADJUSTMENT], [GENERAL_LEARNING]) retrieved as "Potentially Relevant Existing Rules". Determine if the interaction indicates any of these rules need refinement, adjustment, or correction. Update existing rules if a better version exists.
+**General Guidelines for Insight Content and Actions:**
+- Ensure the `<insight>` field always contains the properly formatted insight string: `[TYPE|SCORE] Text`.
+- Be precise with `<old_insight_to_replace>` – it must *exactly* match an existing rule string.
+- Aim for a comprehensive set of operations.
 """
-    insight_user_prompt = f"""Interaction Summary:\n{summary}\n
+        insight_user_prompt = f"""Interaction Summary:\n{summary}\n
 Potentially Relevant Existing Rules (Review these carefully. Your main goal is to consolidate CORE_RULEs and then identify other changes/additions based on the Interaction Summary and these existing rules):\n{existing_rules_ctx}\n
 Guiding principles that were considered during THIS interaction (these might offer clues for new rules or refinements):\n{json.dumps([p['original'] for p in insights_reflected if 'original' in p]) if insights_reflected else "None"}\n
-Task: Based on your reflection process:
-1.  **Consolidate CORE_RULEs:** Merge similar identity/purpose rules into single, definitive statements using "update" operations.
-2.  **Add New Learnings:** "add" any distinct new facts, skills, or important user preferences learned from the "Interaction Summary".
-3.  **Update Existing Principles:** "update" any non-core principles if the "Interaction Summary" provided a clear refinement.
-Combine all findings into a single, valid XML structure. Output XML ONLY.
+Task: Based on your three-step reflection process (Core Identity, New Learnings, Refinements):
+1.  **Consolidate CORE_RULEs:** Merge similar identity/purpose rules from "Potentially Relevant Existing Rules" into single, definitive statements using "update" operations. Replace multiple old versions with the new canonical one.
+2.  **Add New Learnings:** Identify and "add" any distinct new facts, skills, or important user preferences learned from the "Interaction Summary".
+3.  **Update Existing Principles:** "Update" any non-core principles from "Potentially Relevant Existing Rules" if the "Interaction Summary" provided a clear refinement.
+Combine all findings into a single, valid XML structure as specified in the system prompt (root `<operations_list>`, with child `<operation>` elements). Output XML ONLY.
 """
-    insight_msgs = [{"role":"system", "content":insight_sys_prompt}, {"role":"user", "content":insight_user_prompt}]
-    insight_prov, insight_model_disp = provider, model_disp_name
-    insight_env_model = os.getenv("INSIGHT_MODEL_OVERRIDE")
-    if insight_env_model and "/" in insight_env_model:
-        i_p, i_id = insight_env_model.split('/', 1)
-        i_d_n = next((dn for dn, mid in MODELS_BY_PROVIDER.get(i_p.lower(), {}).get("models", {}).items() if mid == i_id), None)
-        if i_d_n: insight_prov, insight_model_disp = i_p, i_d_n
-    raw_ops_xml_full = "".join(list(call_model_stream(provider=insight_prov, model_display_name=insight_model_disp, messages=insight_msgs, api_key_override=api_key_override, temperature=0.0, max_tokens=3500))).strip()
-    xml_match = re.search(r"<operations_list>.*</operations_list>", raw_ops_xml_full, re.DOTALL | re.IGNORECASE)
-    if not xml_match:
-        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No valid XML operations found.")
-        return
-    try:
-        root = ET.fromstring(xml_match.group(0))
-        for op_element in root.findall("operation"):
-            action = op_element.find("action").text.strip().lower() if op_element.find("action") is not None and op_element.find("action").text else None
-            insight_text = op_element.find("insight").text.strip() if op_element.find("insight") is not None and op_element.find("insight").text else None
-            if not action or not insight_text: continue
-            if action == "add":
-                add_rule_entry(insight_text)
-            elif action == "update":
-                old_insight_text = op_element.find("old_insight_to_replace").text.strip() if op_element.find("old_insight_to_replace") is not None and op_element.find("old_insight_to_replace").text else None
-                if old_insight_text:
-                    remove_rule_entry(old_insight_text)
-                add_rule_entry(insight_text)
-    except Exception as e:
-        logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: Error processing insight XML: {e}")
+        insight_msgs = [{"role":"system", "content":insight_sys_prompt}, {"role":"user", "content":insight_user_prompt}]
+        insight_prov, insight_model_disp = provider, model_disp_name
+        insight_env_model = os.getenv("INSIGHT_MODEL_OVERRIDE")
+        if insight_env_model and "/" in insight_env_model:
+            i_p, i_id = insight_env_model.split('/', 1)
+            i_d_n = next((dn for dn, mid in MODELS_BY_PROVIDER.get(i_p.lower(), {}).get("models", {}).items() if mid == i_id), None)
+            if i_d_n: insight_prov, insight_model_disp = i_p, i_d_n
+        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Generating insights with {insight_prov}/{insight_model_disp} (expecting XML)")
+
+        raw_ops_xml_full = "".join(list(call_model_stream(provider=insight_prov, model_display_name=insight_model_disp, messages=insight_msgs, api_key_override=api_key_override, temperature=0.0, max_tokens=3500))).strip()
+
+        ops_data_list, processed_count = [], 0
+
+        xml_match = re.search(r"```xml\s*(<operations_list>.*</operations_list>)\s*```", raw_ops_xml_full, re.DOTALL | re.IGNORECASE) or \
+                    re.search(r"(<operations_list>.*</operations_list>)", raw_ops_xml_full, re.DOTALL | re.IGNORECASE)
+
+        if xml_match:
+            xml_content_str = xml_match.group(1)
+            try:
+                root = ET.fromstring(xml_content_str)
+                if root.tag == "operations_list":
+                    for op_element in root.findall("operation"):
+                        action_el = op_element.find("action")
+                        insight_el = op_element.find("insight")
+                        old_insight_el = op_element.find("old_insight_to_replace")
+
+                        action = action_el.text.strip().lower() if action_el is not None and action_el.text else None
+                        insight_text = insight_el.text.strip() if insight_el is not None and insight_el.text else None
+                        old_insight_text = old_insight_el.text.strip() if old_insight_el is not None and old_insight_el.text else None
+
+                        if action and insight_text:
+                            ops_data_list.append({
+                                "action": action,
+                                "insight": insight_text,
+                                "old_insight_to_replace": old_insight_text
+                            })
+                        else:
+                            logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Skipped XML operation due to missing action or insight text. Action: {action}, Insight: {insight_text}")
+                else:
+                    logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML root tag is not <operations_list>. Found: {root.tag}. XML content:\n{xml_content_str}")
+            except ET.ParseError as e:
+                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML parsing error: {e}. XML content that failed:\n{xml_content_str}")
+            except Exception as e_xml_proc:
+                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: Error processing parsed XML: {e_xml_proc}. XML content:\n{xml_content_str}")
+        else:
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No <operations_list> XML structure found in LLM output. Full raw output:\n{raw_ops_xml_full}")
+
+        if ops_data_list:
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: LLM provided {len(ops_data_list)} insight ops from XML.")
+            for op_idx, op_data in enumerate(ops_data_list):
+                action = op_data["action"]
+                insight_text = op_data["insight"]
+                old_insight = op_data["old_insight_to_replace"]
+
+                if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I|re.DOTALL):
+                    logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to invalid insight_text format from XML: '{insight_text[:100]}...'")
+                    continue
+
+                rule_added_or_updated = False
+                if action == "add":
+                    success, status_msg = add_rule_entry(insight_text)
+                    if success:
+                        processed_count +=1
+                        rule_added_or_updated = True
+                        if insight_text.upper().startswith("[CORE_RULE"):
+                            significant_learnings_summary.append(f"New Core Rule Added: {insight_text}")
+                    else: logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (add from XML): Failed to add rule '{insight_text[:50]}...'. Status: {status_msg}")
+                elif action == "update":
+                    removed_old = False
+                    if old_insight:
+                        if old_insight != insight_text:
+                            remove_success = remove_rule_entry(old_insight)
+                            if not remove_success:
+                                 logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Failed to remove old rule '{old_insight[:50]}...' before adding new.")
+                            else:
+                                removed_old = True
+                        else:
+                             logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Old insight is identical to new insight. Skipping removal.")
+
+                    success, status_msg = add_rule_entry(insight_text)
+                    if success:
+                        processed_count +=1
+                        rule_added_or_updated = True
+                        if insight_text.upper().startswith("[CORE_RULE"):
+                             significant_learnings_summary.append(f"Core Rule Updated (Old: {'Removed' if removed_old else 'Not removed/Same'}, New: {insight_text})")
+                    else: logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Failed to add/update rule '{insight_text[:50]}...'. Status: {status_msg}")
+                else:
+                    logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}' from XML.")
+
+            # After processing all rules, if there were significant learnings, add a special memory
+            if significant_learnings_summary:
+                learning_digest = "SYSTEM CORE LEARNING DIGEST:\n" + "\n".join(significant_learnings_summary)
+                # Create a synthetic metrics object for this system memory
+                system_metrics = {
+                    "takeaway": "Core knowledge refined.",
+                    "response_success_score": 1.0, # Assuming successful internal update
+                    "future_confidence_score": 1.0,
+                    "type": "SYSTEM_REFLECTION"
+                }
+                add_memory_entry(
+                    user_input="SYSTEM_INTERNAL_REFLECTION_TRIGGER", # Fixed identifier
+                    metrics=system_metrics,
+                    bot_response=learning_digest
+                )
+                logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Added CORE_LEARNING_DIGEST to memories: {learning_digest[:100]}...")
+
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Processed {processed_count} insight ops out of {len(ops_data_list)} received from XML.")
+        else:
+            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No valid insight operations derived from LLM's XML output.")
+
+    except Exception as e: logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: CRITICAL ERROR in learning task: {e}", exc_info=True)
+    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: END. Total: {time.time() - learning_start_time:.2f}s")
+
 
 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..."
-    updated_rules_text = ui_refresh_rules_display_fn()
-    updated_mems_json = ui_refresh_memories_display_fn()
+    # Initialize all potential output components with default/current values
+    updated_rules_text = ui_refresh_rules_display_fn() # Get current rules state
+    updated_mems_json = ui_refresh_memories_display_fn() # Get current memories state
     def_detect_out_md = gr.Markdown(visible=False)
-    def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True)
+    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)
 
     if not user_msg_txt.strip():
-        yield (cleared_input, updated_gr_hist, "Error: Empty message.", def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
+        status_txt = "Error: Empty message."
+        updated_gr_hist.append((user_msg_txt or "(Empty)", status_txt))
+        # Ensure all outputs are provided on early exit
+        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)
         return
 
     updated_gr_hist.append((user_msg_txt, "<i>Thinking...</i>"))
+    # Initial yield to update chat UI with thinking message and show current knowledge base state
     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)
 
-    internal_hist = list(current_chat_session_history)
-    internal_hist.append({"role": "user", "content": user_msg_txt})
-    
+    internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
+    # History truncation logic (keep MAX_HISTORY_TURNS pairs + optional system prompt)
+    hist_len_check = MAX_HISTORY_TURNS * 2
+    if internal_hist and internal_hist[0]["role"] == "system": hist_len_check +=1
+    if len(internal_hist) > hist_len_check:
+         current_chat_session_history = ([internal_hist[0]] if internal_hist[0]["role"] == "system" else []) + internal_hist[-(MAX_HISTORY_TURNS * 2):]
+         internal_hist = list(current_chat_session_history) # Use truncated history for current turn processing
+
     final_bot_resp_acc, insights_used_parsed = "", []
     temp_dl_file_path = None
 
@@ -258,114 +406,719 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
             if upd_type == "status":
                 status_txt = upd_data
                 if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
+                    # Update the status alongside the streaming message
                     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)
+                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg) # Update chat with streamed chunk
             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..."
+                # Ensure the final chat message reflects the full response
+                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, final_bot_resp_acc or "(No text)")
-                
-                def_fmt_out_txt = gr.Textbox(value=final_bot_resp_acc, interactive=True)
-                if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
-                    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
-                        tmpfile.write(final_bot_resp_acc)
-                        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.*")
-                def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True)
+                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg or "(No text)")
 
+                # Update detailed response box and download button
+                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:"):
+                    try:
+                        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)
+                    except Exception as e:
+                        logger.error(f"Error creating temp file for download: {e}", exc_info=False)
+                        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False, label="Download Error")
+                else:
+                    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
+
+                # Update insights display
+                insights_md_content = "### Insights Considered (Pre-Response):\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
+                def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True if insights_used_parsed else False)
+
+            # Yield intermediate updates for the UI
+            # Pass the *current* state of rules and memories display components
             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)
+
+            # Stop processing generator after final_response_and_insights
             if upd_type == "final_response_and_insights": break
+
     except Exception as e:
-        status_txt = f"Error: {str(e)[:100]}"
-        updated_gr_hist.append((user_msg_txt, f"Sorry, an error occurred: {status_txt}"))
-        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)
-        return
+        logger.error(f"Chat handler error during main processing: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
+        error_message_for_chat = f"Sorry, an error occurred during response generation: {str(e)[:100]}"
+        if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
+            updated_gr_hist[-1] = (user_msg_txt, error_message_for_chat)
+        else:
+            updated_gr_hist.append((user_msg_txt, error_message_for_chat))
+        def_fmt_out_txt = gr.Textbox(value=error_message_for_chat, interactive=True)
+        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
+        def_detect_out_md = gr.Markdown(value="*Error processing request.*", visible=True)
+
+        # Provide the current state of rules/memories on error path yield
+        current_rules_text_on_error = ui_refresh_rules_display_fn()
+        current_mems_json_on_error = ui_refresh_memories_display_fn()
+
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_on_error, current_mems_json_on_error)
+        # Clean up temp file if created before error
+        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} after error: {e_unlink}")
+        return # Exit the function after error handling
 
+    # --- Post-Interaction Learning ---
     if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
+        # Add the successful turn to the internal history
         current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
-        if len(current_chat_session_history) > MAX_HISTORY_TURNS * 2:
-            current_chat_session_history = current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]
+        # History truncation again after adding
+        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):]
 
-        yield (cleared_input, updated_gr_hist, "<i>[Performing post-interaction learning...]</i>", def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
-        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."
-    
+        status_txt = "<i>[Performing post-interaction learning...]</i>"
+        # Yield status before synchronous learning
+        current_rules_text_before_learn = ui_refresh_rules_display_fn()
+        current_mems_json_before_learn = ui_refresh_memories_display_fn()
+        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_before_learn, current_mems_json_before_learn)
+
+        try:
+            perform_post_interaction_learning(
+                user_input=user_msg_txt,
+                bot_response=final_bot_resp_acc,
+                provider=sel_prov_name,
+                model_disp_name=sel_model_disp_name,
+                insights_reflected=insights_used_parsed,
+                api_key_override=ui_api_key.strip() if ui_api_key else None
+            )
+            status_txt = "Response & Learning Complete."
+        except Exception as e_learn:
+            logger.error(f"Error during post-interaction learning: {e_learn}", exc_info=True)
+            status_txt = "Response complete. Error during learning."
+
+    elif final_bot_resp_acc.startswith("Error:"):
+        status_txt = final_bot_resp_acc
+        # If it was an error response from the generator, it's already in updated_gr_hist[-1]
+        # The other output components (fmt_report_tb, dl_btn, detect_out_md) are already set by the generator loop or default state
+    else:
+        status_txt = "Processing finished; no valid response or error occurred during main phase."
+
+
+    # Final yield after learning (or error handling)
+    # This final yield updates the UI one last time with the true final status
+    # AND crucially refreshes the Rules and Memories displays in case they changed during learning.
     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)
+
+    # Clean up the temporary download file after the final yield
     if temp_dl_file_path and os.path.exists(temp_dl_file_path):
         try: os.unlink(temp_dl_file_path)
-        except: pass
+        except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path}: {e_unlink}")
+
+
+# --- Startup Loading Functions ---
+def load_rules_from_file(filepath: str | None):
+    """Loads rules from a local file (.txt or .jsonl) and adds them to the system."""
+    if not filepath:
+        logger.info("LOAD_RULES_FILE environment variable not set. Skipping rules loading from file.")
+        return 0, 0, 0 # added, skipped, errors
+
+    if not os.path.exists(filepath):
+        logger.warning(f"LOAD_RULES: Specified rules file not found: {filepath}. Skipping loading.")
+        return 0, 0, 0
+
+    added_count, skipped_count, error_count = 0, 0, 0
+    potential_rules = []
+
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            content = f.read()
+    except Exception as e:
+        logger.error(f"LOAD_RULES: Error reading file {filepath}: {e}", exc_info=False)
+        return 0, 0, 1 # Indicate read error
+
+    if not content.strip():
+        logger.info(f"LOAD_RULES: File {filepath} is empty. Skipping loading.")
+        return 0, 0, 0
+
+    file_name_lower = filepath.lower()
+
+    if file_name_lower.endswith(".txt"):
+        potential_rules = content.split("\n\n---\n\n")
+        # Also handle simple line breaks if '---' separator is not used
+        if len(potential_rules) == 1 and "\n" in content:
+             potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
+    elif file_name_lower.endswith(".jsonl"):
+        for line_num, line in enumerate(content.splitlines()):
+            line = line.strip()
+            if line:
+                try:
+                    # Expecting each line to be a JSON string containing the rule text
+                    rule_text_in_json_string = json.loads(line)
+                    if isinstance(rule_text_in_json_string, str):
+                        potential_rules.append(rule_text_in_json_string)
+                    else:
+                        logger.warning(f"LOAD_RULES (JSONL): Line {line_num+1} in {filepath} did not contain a string value. Got: {type(rule_text_in_json_string)}")
+                        error_count +=1
+                except json.JSONDecodeError:
+                    logger.warning(f"LOAD_RULES (JSONL): Line {line_num+1} in {filepath} failed to parse as JSON: {line[:100]}")
+                    error_count +=1
+    else:
+        logger.error(f"LOAD_RULES: Unsupported file type for rules: {filepath}. Must be .txt or .jsonl")
+        return 0, 0, 1 # Indicate type error
+
+    valid_potential_rules = [r.strip() for r in potential_rules if r.strip()]
+    total_to_process = len(valid_potential_rules)
+
+    if total_to_process == 0 and error_count == 0:
+        logger.info(f"LOAD_RULES: No valid rule segments found in {filepath} to process.")
+        return 0, 0, 0
+    elif total_to_process == 0 and error_count > 0:
+         logger.warning(f"LOAD_RULES: No valid rule segments found to process. Encountered {error_count} parsing/format errors in {filepath}.")
+         return 0, 0, error_count # Indicate only errors
+
+    logger.info(f"LOAD_RULES: Attempting to add {total_to_process} potential rules from {filepath}...")
+    for idx, rule_text in enumerate(valid_potential_rules):
+        success, status_msg = add_rule_entry(rule_text)
+        if success:
+            added_count += 1
+        elif status_msg == "duplicate":
+            skipped_count += 1
+        else:
+            logger.warning(f"LOAD_RULES: Failed to add rule from {filepath} (segment {idx+1}): '{rule_text[:50]}...'. Status: {status_msg}")
+            error_count += 1
+
+    logger.info(f"LOAD_RULES: Finished processing {filepath}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors: {error_count}.")
+    return added_count, skipped_count, error_count
+
+def load_memories_from_file(filepath: str | None):
+    """Loads memories from a local file (.json or .jsonl) and adds them to the system."""
+    if not filepath:
+        logger.info("LOAD_MEMORIES_FILE environment variable not set. Skipping memories loading from file.")
+        return 0, 0, 0 # added, format_errors, save_errors
+
+    if not os.path.exists(filepath):
+        logger.warning(f"LOAD_MEMORIES: Specified memories file not found: {filepath}. Skipping loading.")
+        return 0, 0, 0
+
+    added_count, format_error_count, save_error_count = 0, 0, 0
+    memory_objects_to_process = []
+
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            content = f.read()
+    except Exception as e:
+        logger.error(f"LOAD_MEMORIES: Error reading file {filepath}: {e}", exc_info=False)
+        return 0, 1, 0 # Indicate read error
+
+    if not content.strip():
+        logger.info(f"LOAD_MEMORIES: File {filepath} is empty. Skipping loading.")
+        return 0, 0, 0
+
+    file_ext = os.path.splitext(filepath.lower())[1]
+
+    if file_ext == ".json":
+        try:
+            parsed_json = json.loads(content)
+            if isinstance(parsed_json, list):
+                memory_objects_to_process = parsed_json
+            elif isinstance(parsed_json, dict):
+                 # If it's a single object, process it as a list of one
+                 memory_objects_to_process = [parsed_json]
+            else:
+                logger.warning(f"LOAD_MEMORIES (.json): File content is not a JSON list or object in {filepath}. Type: {type(parsed_json)}")
+                format_error_count = 1
+        except json.JSONDecodeError as e:
+            logger.warning(f"LOAD_MEMORIES (.json): Invalid JSON file {filepath}. Error: {e}")
+            format_error_count = 1
+    elif file_ext == ".jsonl":
+        for line_num, line in enumerate(content.splitlines()):
+            line = line.strip()
+            if line:
+                try:
+                    memory_objects_to_process.append(json.loads(line))
+                except json.JSONDecodeError:
+                    logger.warning(f"LOAD_MEMORIES (.jsonl): Line {line_num+1} in {filepath} parse error: {line[:100]}")
+                    format_error_count += 1
+    else:
+        logger.error(f"LOAD_MEMORIES: Unsupported file type for memories: {filepath}. Must be .json or .jsonl")
+        return 0, 1, 0 # Indicate type error
+
+    total_to_process = len(memory_objects_to_process)
+
+    if total_to_process == 0 and format_error_count > 0 :
+         logger.warning(f"LOAD_MEMORIES: File parsing failed for {filepath}. Found {format_error_count} format errors and no processable objects.")
+         return 0, format_error_count, 0
+    elif total_to_process == 0:
+         logger.info(f"LOAD_MEMORIES: No memory objects found in {filepath} after parsing.")
+         return 0, 0, 0
+
+
+    logger.info(f"LOAD_MEMORIES: Attempting to add {total_to_process} memory objects from {filepath}...")
+    for idx, mem_data in enumerate(memory_objects_to_process):
+        # Validate minimum structure
+        if isinstance(mem_data, dict) and all(k in mem_data for k in ["user_input", "bot_response", "metrics"]):
+            # Add entry without generating new embeddings if possible (assuming file contains embeddings)
+            # NOTE: The current add_memory_entry function *always* generates embeddings.
+            # If performance is an issue with large files, memory_logic might need
+            # an optimized bulk import function that reuses existing embeddings or
+            # generates them in batches. For now, we use the existing function.
+            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"]) # add_memory_entry needs user_input, metrics, bot_response
+            if success:
+                added_count += 1
+            else:
+                # add_memory_entry currently doesn't return detailed error status
+                logger.warning(f"LOAD_MEMORIES: Failed to save memory object from {filepath} (segment {idx+1}). Data: {str(mem_data)[:100]}")
+                save_error_count += 1
+        else:
+            logger.warning(f"LOAD_MEMORIES: Skipped invalid memory object structure in {filepath} (segment {idx+1}): {str(mem_data)[:100]}")
+            format_error_count += 1
+
+    logger.info(f"LOAD_MEMORIES: Finished processing {filepath}. Added: {added_count}, Format/Structure Errors: {format_error_count}, Save Errors: {save_error_count}.")
+    return added_count, format_error_count, save_error_count
+
 
+# --- UI Functions for Rules and Memories (ui_refresh_..., ui_download_..., ui_upload_...) ---
 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")
+    try:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt", encoding='utf-8') as tmpfile:
+            tmpfile.write(rules_content)
+            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:
+        with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f: content = f.read()
+    except Exception as e_read: return f"Error reading file: {e_read}"
+    if not content.strip(): return "Uploaded rules file is empty."
+    added_count, skipped_count, error_count = 0,0,0
+
+    potential_rules = []
+    file_name_lower = uploaded_file_obj.name.lower()
+
+    if file_name_lower.endswith(".txt"):
+        potential_rules = content.split("\n\n---\n\n")
+        if len(potential_rules) == 1 and "\n" in content:
+             potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
+    elif file_name_lower.endswith(".jsonl"):
+        for line_num, line in enumerate(content.splitlines()):
+            line = line.strip()
+            if line:
+                try:
+                    rule_text_in_json_string = json.loads(line)
+                    if isinstance(rule_text_in_json_string, str):
+                        potential_rules.append(rule_text_in_json_string)
+                    else:
+                        logger.warning(f"Rule Upload (JSONL): Line {line_num+1} did not contain a string value. Got: {type(rule_text_in_json_string)}")
+                        error_count +=1
+                except json.JSONDecodeError:
+                    logger.warning(f"Rule Upload (JSONL): Line {line_num+1} failed to parse as JSON: {line[:100]}")
+                    error_count +=1
+    else:
+        return "Unsupported file type for rules. Please use .txt or .jsonl."
+
+    valid_potential_rules = [r.strip() for r in potential_rules if r.strip()]
+    total_to_process = len(valid_potential_rules)
+
+    if total_to_process == 0 and error_count == 0:
+        return "No valid rules found in file to process."
+    elif total_to_process == 0 and error_count > 0:
+        return f"No valid rules found to process. Encountered {error_count} parsing/format errors."
+
+    progress(0, desc="Starting rules upload...")
+    for idx, rule_text in enumerate(valid_potential_rules):
+        success, status_msg = add_rule_entry(rule_text)
+        if success: added_count += 1
+        elif status_msg == "duplicate": skipped_count += 1
+        else: error_count += 1
+        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} rules...")
+
+    msg = f"Rules Upload: Total valid rule segments processed: {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors (parsing/add): {error_count}."
+    logger.info(msg); return msg
+
 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 = ""
+    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}")
+
+    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
+    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:
+        with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f: content = f.read()
+    except Exception as e_read: return f"Error reading file: {e_read}"
+    if not content.strip(): return "Uploaded memories file is empty."
+    added_count, format_error_count, save_error_count = 0,0,0
+    memory_objects_to_process = []
+
+    file_ext = os.path.splitext(uploaded_file_obj.name.lower())[1]
+
+    if file_ext == ".json":
+        try:
+            parsed_json = json.loads(content)
+            if isinstance(parsed_json, list):
+                memory_objects_to_process = parsed_json
+            elif isinstance(parsed_json, dict):
+                 memory_objects_to_process = [parsed_json]
+            else:
+                logger.warning(f"Memories Upload (.json): File content is not a JSON list or object. Type: {type(parsed_json)}")
+                format_error_count = 1
+        except json.JSONDecodeError as e:
+            logger.warning(f"Memories Upload (.json): Invalid JSON file. Error: {e}")
+            format_error_count = 1
+    elif file_ext == ".jsonl":
+        for line_num, line in enumerate(content.splitlines()):
+            line = line.strip()
+            if line:
+                try:
+                    memory_objects_to_process.append(json.loads(line))
+                except json.JSONDecodeError:
+                    logger.warning(f"Memories Upload (.jsonl): Line {line_num+1} parse error: {line[:100]}")
+                    format_error_count += 1
+    else:
+        return "Unsupported file type for memories. Please use .json or .jsonl."
+
+    if not memory_objects_to_process and format_error_count > 0 :
+         return f"Memories Upload: File parsing failed. Found {format_error_count} format errors and no processable objects."
+    elif not memory_objects_to_process:
+         return "No valid memory objects found in the uploaded file."
+
+    total_to_process = len(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"]):
+            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"])
+            if success: added_count += 1
+            else: save_error_count += 1
+        else:
+            logger.warning(f"Memories Upload: Skipped invalid memory object structure: {str(mem_data)[:100]}")
+            format_error_count += 1
+        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} memories...")
+
+    msg = f"Memories Upload: Processed {total_to_process} objects. Added: {added_count}, Format/Structure Errors: {format_error_count}, Save Errors: {save_error_count}."
+    logger.info(msg); return msg
+
+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()]
+
+    if not potential_rules:
+        return "No rules found to process from editor."
+
+    added, skipped, errors = 0, 0, 0
+    unique_rules_to_process = sorted(list(set(filter(None, [r.strip() for r in potential_rules]))))
+
+    total_unique = len(unique_rules_to_process)
+    if total_unique == 0: return "No unique, non-empty rules found in editor text."
+
+    progress(0, desc=f"Saving {total_unique} unique rules from editor...")
+
+    for idx, rule_text in enumerate(unique_rules_to_process):
+        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_unique, desc=f"Processed {idx+1}/{total_unique} rules...")
+
+    return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors} from {total_unique} unique rules in text."
+
 def app_load_fn():
+    logger.info("App loading. Initializing systems...")
     initialize_memory_system()
-    return "AI Systems Initialized. Ready.", ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn()
+    logger.info("Memory system initialized.")
+
+    # --- Load Rules from File ---
+    rules_added, rules_skipped, rules_errors = load_rules_from_file(LOAD_RULES_FILE)
+    rules_load_msg = f"Rules: Added {rules_added}, Skipped {rules_skipped}, Errors {rules_errors} from {LOAD_RULES_FILE or 'None'}."
+    logger.info(rules_load_msg)
+
+    # --- Load Memories from File ---
+    mems_added, mems_format_errors, mems_save_errors = load_memories_from_file(LOAD_MEMORIES_FILE)
+    mems_load_msg = f"Memories: Added {mems_added}, Format Errors {mems_format_errors}, Save Errors {mems_save_errors} from {LOAD_MEMORIES_FILE or 'None'}."
+    logger.info(mems_load_msg)
+
+    final_status = f"AI Systems Initialized. {rules_load_msg} {mems_load_msg} Ready."
+
+    # Initial population of all relevant UI components AFTER loading
+    rules_on_load = ui_refresh_rules_display_fn()
+    mems_on_load = ui_refresh_memories_display_fn()
+
+    # Return values for outputs defined in demo.load
+    return (
+        final_status, # agent_stat_tb
+        rules_on_load, # rules_disp_ta
+        mems_on_load,  # mems_disp_json
+        gr.Markdown(visible=False), # detect_out_md (initial state)
+        gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True), # fmt_report_tb (initial state)
+        gr.DownloadButton(interactive=False, value=None, visible=False), # dl_report_btn (initial state)
+    )
+
+
+# --- Gradio UI Definition ---
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css="""
+        .gr-button { margin: 5px; }
+        .gr-textbox, .gr-text-area, .gr-dropdown, .gr-json { 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; }
+        .gr-json { max-height: 300px; overflow-y: auto; } /* Added scrolling for JSON */
+    """
+) as demo:
+    gr.Markdown(
+        """
+        # 🤖 AI Research Agent
+        Your intelligent assistant for research and knowledge management
+        """,
+        elem_classes=["header"]
+    )
+
+    is_sqlite = MEMORY_STORAGE_BACKEND == "SQLITE"
+    is_hf_dataset = MEMORY_STORAGE_BACKEND == "HF_DATASET"
 
-with gr.Blocks(theme=gr.themes.Soft(), css=".gr-json { max-height: 300px; overflow-y: auto; }") as demo:
-    gr.Markdown("# 🤖 AI Research Agent")
     with gr.Row(variant="compact"):
-        agent_stat_tb = gr.Textbox(label="Agent Status", value="Initializing...", interactive=False, scale=4)
+        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):
-            gr.Textbox(label="Memory Backend", value=MEMORY_STORAGE_BACKEND, interactive=False)
-            gr.Textbox(label="SQLite Path", value=MEMORY_SQLITE_PATH, interactive=False, visible=MEMORY_STORAGE_BACKEND == "SQLITE")
-            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")
+            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=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, elem_classes=["status-text"]
+            )
+
     with gr.Row():
         with gr.Sidebar():
             gr.Markdown("## ⚙️ Configuration")
-            api_key_tb = gr.Textbox(label="API Key (Override)", type="password")
-            available_providers = get_available_providers()
-            prov_sel_dd = gr.Dropdown(label="AI Provider", choices=available_providers, value=available_providers[0] if available_providers else None, interactive=True)
-            model_sel_dd = gr.Dropdown(label="AI Model", choices=get_model_display_names_for_provider(prov_sel_dd.value) if prov_sel_dd.value else [], value=get_default_model_display_name_for_provider(prov_sel_dd.value) if prov_sel_dd.value else None, interactive=True)
-            sys_prompt_tb = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT)
+            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"
+                )
+                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
+                )
+                default_model_display = get_default_model_display_name_for_provider(default_provider) if default_provider else None
+                model_sel_dd = gr.Dropdown(
+                    label="AI Model",
+                    choices=get_model_display_names_for_provider(default_provider) if default_provider else [],
+                    value=default_model_display,
+                    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")
+
         with gr.Column(scale=3):
             with gr.Tabs():
                 with gr.TabItem("💬 Chat & Research"):
-                    main_chat_disp = gr.Chatbot(label="Chat", height=400, show_copy_button=True, render_markdown=True)
-                    user_msg_tb = gr.Textbox(placeholder="Ask a question...", scale=7)
-                    send_btn = gr.Button("Send", variant="primary")
-                    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", interactive=False, visible=False)
-                        detect_out_md = gr.Markdown(visible=False)
+                    with gr.Group():
+                        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
+                            )
+                            send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
+                        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():
+                    with gr.Row(equal_height=True):
                         with gr.Column():
-                            rules_disp_ta = gr.TextArea(label="Current Rules", lines=10, interactive=True)
+                            gr.Markdown("### 📜 Rules Management")
+                            rules_disp_ta = gr.TextArea(
+                                label="Current Rules", lines=10,
+                                placeholder="Rules will appear here.",
+                                interactive=True
+                            )
+                            gr.Markdown("To edit rules, modify the text above and click 'Save Edited Text', or upload a new file.")
                             save_edited_rules_btn = gr.Button("💾 Save Edited Text", variant="primary")
+                            with gr.Row(variant="compact"):
+                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules", value=None)
+                                clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop")
+                            upload_rules_fobj = gr.File(
+                                label="Upload Rules File (.txt with '---' separators, or .jsonl of rule strings)",
+                                file_types=[".txt", ".jsonl"]
+                            )
+                            rules_stat_tb = gr.Textbox(
+                                label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"]
+                            )
+
                         with gr.Column():
-                            mems_disp_json = gr.JSON(label="Current Memories")
+                            gr.Markdown("### 📚 Memories Management")
+                            mems_disp_json = gr.JSON(
+                                label="Current Memories", value=[]
+                            )
+                            gr.Markdown("To add memories, upload a .jsonl or .json file.")
+                            with gr.Row(variant="compact"):
+                                dl_mems_btn = gr.DownloadButton("⬇️ Download Memories", value=None)
+                                clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop")
+                            upload_mems_fobj = gr.File(
+                                label="Upload Memories File (.jsonl of memory objects, or .json array of objects)",
+                                file_types=[".jsonl", ".json"]
+                            )
+                            mems_stat_tb = gr.Textbox(
+                                label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"]
+                            )
 
     def dyn_upd_model_dd(sel_prov_dyn: str):
-        models = get_model_display_names_for_provider(sel_prov_dyn)
-        default_model = get_default_model_display_name_for_provider(sel_prov_dyn)
-        return gr.Dropdown(choices=models, value=default_model, interactive=True)
+        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)
 
+    # Inputs for the main chat submission function
     chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
+    # Outputs for the main chat submission function (includes knowledge base displays)
     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]
-    
-    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)
-    
-    def save_rules_from_editor(text):
-        for line in text.splitlines():
-            if line.strip(): add_rule_entry(line.strip())
-        return ui_refresh_rules_display_fn()
-    save_edited_rules_btn.click(fn=save_rules_from_editor, inputs=[rules_disp_ta], outputs=[rules_disp_ta])
-
-    demo.load(fn=app_load_fn, outputs=[agent_stat_tb, rules_disp_ta, mems_disp_json])
+
+    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)
+
+    # Rules Management events
+    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False) # show_progress=False for download buttons
+
+    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_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False) # Refresh display after saving
+
+    upload_rules_fobj.upload(
+        fn=ui_upload_rules_action_fn,
+        inputs=[upload_rules_fobj],
+        outputs=[rules_stat_tb],
+        show_progress="full"
+    ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False) # Refresh display after upload
+
+    clear_rules_btn.click(
+        fn=lambda: ("All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules."),
+        outputs=rules_stat_tb,
+        show_progress=False
+    ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False) # Refresh display after clear
+
+    # Memories Management events
+    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn, show_progress=False) # show_progress=False for download buttons
+
+    upload_mems_fobj.upload(
+        fn=ui_upload_memories_action_fn,
+        inputs=[upload_mems_fobj],
+        outputs=[mems_stat_tb],
+        show_progress="full"
+    ).then(fn=ui_refresh_memories_display_fn, outputs=mems_disp_json, show_progress=False) # Refresh display after upload
+
+    clear_mems_btn.click(
+        fn=lambda: ("All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories."),
+        outputs=mems_stat_tb,
+        show_progress=False
+    ).then(fn=ui_refresh_memories_display_fn, outputs=mems_disp_json, show_progress=False) # Refresh display after clear
+
+    # FAISS save button visibility and action (RAM backend only)
+    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
+        def save_faiss_action_with_feedback_sidebar_fn():
+            try:
+                save_faiss_indices_to_disk()
+                gr.Info("Attempted to save FAISS indices to disk.")
+            except Exception as e:
+                logger.error(f"Error saving FAISS indices: {e}", exc_info=True)
+                gr.Error(f"Error saving FAISS indices: {e}")
+
+        save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None, show_progress=False)
+
+
+    # --- Initial Load Event ---
+    # This function runs once when the Gradio app starts.
+    # It initializes the memory system and loads data from specified files.
+    # Its outputs populate the initial state of several UI components.
+    app_load_outputs = [
+        agent_stat_tb,          # Updates status text
+        rules_disp_ta,          # Populates rules display
+        mems_disp_json,         # Populates memories display
+        detect_out_md,          # Sets initial visibility/value for insights markdown
+        fmt_report_tb,          # Sets initial value for detailed response textbox
+        dl_report_btn           # Sets initial state for download button
+    ]
+    demo.load(fn=app_load_fn, inputs=None, outputs=app_load_outputs, show_progress="full")
+
 
 if __name__ == "__main__":
-    logger.info(f"Starting Gradio AI Research Agent (Memory: {MEMORY_STORAGE_TYPE})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v6.5 - Direct UI Update & Core Learning Memories, Memory: {MEMORY_STORAGE_BACKEND})...")
     app_port = int(os.getenv("GRADIO_PORT", 7860))
     app_server = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
-    demo.queue().launch(server_name=app_server, server_port=app_port, share=False)
+    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.")