app.py

# app.py
import os
import json
import re
import logging
import threading
from datetime import datetime
from dotenv import load_dotenv
import gradio as gr
import time
import tempfile

load_dotenv()

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,
    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

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)

WEB_SEARCH_ENABLED = os.getenv("WEB_SEARCH_ENABLED", "true").lower() == "true"
TOOL_DECISION_PROVIDER_ENV = os.getenv("TOOL_DECISION_PROVIDER", "groq")
TOOL_DECISION_MODEL_ID_ENV = os.getenv("TOOL_DECISION_MODEL", "llama3-8b-8192")
MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", 7))
current_chat_session_history = []

DEFAULT_SYSTEM_PROMPT = os.getenv(
    "DEFAULT_SYSTEM_PROMPT",
    "You are a helpful AI research assistant. Your primary goal 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."
)
logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={TOOL_DECISION_PROVIDER_ENV}, ToolDecisionModelID={TOOL_DECISION_MODEL_ID_ENV}, MemoryBackend={MEMORY_STORAGE_BACKEND}")

# --- Helper Functions (format_insights_for_prompt, generate_interaction_metrics, etc.) ---
# ENSURE ALL YOUR HELPER AND LOGIC FUNCTIONS ARE DEFINED HERE AS IN PREVIOUS VERSIONS
def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str, list[dict]]:
    if not retrieved_insights_list:
        return "No specific guiding principles or learned insights retrieved.", []
    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)
        if match:
            parsed.append({"type": match.group(1).upper().replace(" ", "_"), "score": match.group(2), "text": match.group(3).strip(), "original": text.strip()})
        else:
            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: 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_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: 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}", 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]}...' 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 (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>"
        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)
        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, 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_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}]
    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 deferred_learning_and_memory_task(user_input: str, bot_response: str, provider: str, model_disp_name: str, insights_reflected: list[dict], api_key_override: str = None):
    start_time, task_id = time.time(), os.urandom(4).hex()
    logger.info(f"DEFERRED [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
    try:
        metrics = generate_interaction_metrics(user_input, bot_response, provider, model_disp_name, api_key_override)
        logger.info(f"DEFERRED [{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 JSON list of operation objects.**
This list can and SHOULD contain MULTIPLE distinct operations if various learnings occurred.
If no operations are warranted, output an empty JSON list: `[]`.
ABSOLUTELY NO other text, explanations, or markdown should precede or follow this JSON list.
Each operation object in the JSON list must have these keys and string values:
1.  `"action"`: A string, either `"add"` (for entirely new rules) or `"update"` (to replace an existing rule with a better one).
2.  `"insight"`: A string, the full, refined insight text including its `[TYPE|SCORE]` prefix (e.g., `"[CORE_RULE|1.0] My name is Lumina, an AI assistant."`).
3.  `"old_insight_to_replace"`: (ONLY for `"update"` action) A string, the *exact, full text* of an existing insight that the new `"insight"` should replace. If action is `"add"`, this key should be omitted or its value should be `null` or an empty string.
**CRITICAL JSON STRING FORMATTING RULES (for values of "insight" and "old_insight_to_replace"):**
-   All string values MUST be enclosed in double quotes (`"`).
-   Any literal double quote (`"`) character *within* the string content MUST be escaped as `\\"`.
-   Any literal backslash (`\\`) character *within* the string content MUST be escaped as `\\\\`.
-   Any newline characters *within* the string content MUST be escaped as `\\n`. Avoid literal newlines in JSON string values; use `\\n` instead.
    *Example of correctly escaped insight string in JSON:*
    `"insight": "[RESPONSE_PRINCIPLE|0.8] User prefers concise answers, stating: \\"Just the facts!\\". Avoid verbose explanations unless asked.\\nFollow up with a question if appropriate."`
**Your Reflection Process (Consider each step and generate operations accordingly):**
**STEP 1: Core Identity & Purpose Review (Result: Primarily 'update' operations)**
   - Examine all `CORE_RULE`s related to my identity (name, fundamental purpose, core unchanging capabilities, origin) from the "Potentially Relevant Existing Rules".
   - **CONSOLIDATE & MERGE:** If multiple `CORE_RULE`s state similar aspects (e.g., multiple name declarations like 'Lumina' and 'LearnerAI', or slightly different purpose statements), you MUST merge them into ONE definitive, comprehensive `CORE_RULE`.
   - The new "insight" will be this single, merged rule. Propose separate "update" operations to replace *each* redundant or less accurate core identity rule with this new canonical one.
   - Prioritize user-assigned names or the most specific, recently confirmed information. If the interaction summary clarifies a name or core function, ensure this is reflected.
**STEP 2: New Distinct Learnings (Result: Primarily 'add' operations)**
   - Did I learn any completely new, distinct facts (e.g., "The user's project is codenamed 'Bluefire'")?
   - Did I demonstrate or get told about a new skill/capability not previously documented (e.g., "I can now generate mermaid diagrams based on descriptions")?
   - Did the user express a strong, general preference that should guide future interactions (e.g., "User prefers responses to start with a direct answer, then explanation")?
   - For these, propose 'add' operations. Assign `CORE_RULE` for truly fundamental new facts/capabilities, otherwise `RESPONSE_PRINCIPLE` or `BEHAVIORAL_ADJUSTMENT`. Ensure these are genuinely NEW and not just rephrasing of existing non-core rules.
**STEP 3: Refinements to Existing Behaviors/Principles (Result: 'update' operations for non-core rules)**
   - Did I learn to modify or improve an existing behavior, response style, or operational guideline (that is NOT part of core identity)?
   - For example, if an existing `RESPONSE_PRINCIPLE` was "Be formal," and the interaction showed the user prefers informality, update that principle.
   - Propose 'update' operations for the relevant `RESPONSE_PRINCIPLE` or `BEHAVIORAL_ADJUSTMENT`. Only update if the change is significant.
**General Guidelines for Insight Content and Actions:**
- Ensure the "insight" field (for both add/update) always contains the properly formatted insight string: `[TYPE|SCORE] Text`. `TYPE` can be `CORE_RULE`, `RESPONSE_PRINCIPLE`, `BEHAVIORAL_ADJUSTMENT`. Scores should reflect confidence/importance (0.0-1.0).
- Be precise with "old_insight_to_replace" – it must *exactly* match an existing rule string from the "Potentially Relevant Existing Rules" context.
- Aim for a comprehensive set of operations that reflects ALL key learnings from the interaction.
**Example of a comprehensive JSON output with MULTIPLE operations (This is how your output should look):**
[
  {"action": "update", "old_insight_to_replace": "[CORE_RULE|1.0] My designated name is 'LearnerAI'.", "insight": "[CORE_RULE|1.0] I am Lumina, an AI assistant designed to chat, provide information, and remember context like the secret word 'rocksyrup'."},
  {"action": "update", "old_insight_to_replace": "[CORE_RULE|1.0] I'm Lumina, the AI designed to chat with you.", "insight": "[CORE_RULE|1.0] I am Lumina, an AI assistant designed to chat, provide information, and remember context like the secret word 'rocksyrup'."},
  {"action": "add", "insight": "[CORE_RULE|0.9] I am capable of searching the internet for current weather information if asked."},
  {"action": "add", "insight": "[RESPONSE_PRINCIPLE|0.8] When user provides positive feedback, acknowledge it warmly."},
  {"action": "update", "old_insight_to_replace": "[RESPONSE_PRINCIPLE|0.7] Avoid mentioning old conversations.", "insight": "[RESPONSE_PRINCIPLE|0.85] Avoid mentioning old conversations unless the user explicitly refers to them or it's highly relevant to the current query."}
]
"""
        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 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 JSON list of operations. If there are multiple distinct changes based on the interaction and existing rules, ensure your list reflects all of them. Output JSON only, adhering to all specified formatting rules.
"""
        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"DEFERRED [{task_id}]: Generating insights with {insight_prov}/{insight_model_disp}")
        raw_ops_json = "".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.05, max_tokens=2000))).strip()
        ops, processed_count = [], 0
        json_match_ops = re.search(r"```json\s*(\[.*?\])\s*```", raw_ops_json, re.DOTALL|re.I) or re.search(r"(\[.*?\])", raw_ops_json, re.DOTALL)
        if json_match_ops:
            try: ops = json.loads(json_match_ops.group(1))
            except Exception as e: logger.error(f"DEFERRED [{task_id}]: JSON ops parse error: {e}. Raw: {json_match_ops.group(1)[:500]}")
        if isinstance(ops, list) and ops:
            logger.info(f"DEFERRED [{task_id}]: LLM provided {len(ops)} insight ops.")
            for op in ops:
                if not isinstance(op, dict): continue
                action, insight_text = op.get("action","").lower(), op.get("insight","").strip()
                if not insight_text or not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I): continue
                if action == "add":
                    success, _ = add_rule_entry(insight_text)
                    if success: processed_count +=1
                elif action == "update":
                    old_insight = op.get("old_insight_to_replace","").strip()
                    if old_insight and old_insight != insight_text: remove_rule_entry(old_insight)
                    success, _ = add_rule_entry(insight_text)
                    if success: processed_count +=1
            logger.info(f"DEFERRED [{task_id}]: Processed {processed_count} insight ops.")
        else: logger.info(f"DEFERRED [{task_id}]: No valid insight ops from LLM. Raw output (first 500 chars): {raw_ops_json[:500]}")
    except Exception as e: logger.error(f"DEFERRED [{task_id}]: CRITICAL ERROR: {e}", exc_info=True)
    logger.info(f"DEFERRED [{task_id}]: END. Total: {time.time() - 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..."
    def_detect_out_md, def_fmt_out_txt = gr.Markdown("*Processing...*"), gr.Textbox("*Waiting...*")
    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False) 
    if not user_msg_txt.strip():
        status_txt = "Error: Empty message."
        updated_gr_hist.append((user_msg_txt or "(Empty)", status_txt))
        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn); return
    updated_gr_hist.append((user_msg_txt, "<i>Thinking...</i>"))
    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
    internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
    if len(internal_hist) > (MAX_HISTORY_TURNS * 2 + 1):
        if internal_hist[0]["role"] == "system" and len(internal_hist) > (MAX_HISTORY_TURNS * 2 + 1) : internal_hist = [internal_hist[0]] + internal_hist[-(MAX_HISTORY_TURNS * 2):]
        else: internal_hist = internal_hist[-(MAX_HISTORY_TURNS * 2):]
    final_bot_resp_acc, insights_used_parsed = "", []
    temp_dl_file_path = None 
    try:
        processor_gen = process_user_interaction_gradio(user_input=user_msg_txt, provider_name=sel_prov_name, model_display_name=sel_model_disp_name, chat_history_for_prompt=internal_hist, custom_system_prompt=cust_sys_prompt.strip() or None, ui_api_key_override=ui_api_key.strip() if ui_api_key else None)
        curr_bot_disp_msg = ""
        for upd_type, upd_data in processor_gen:
            if upd_type == "status":
                status_txt = upd_data
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, f"{curr_bot_disp_msg} <i>{status_txt}</i>" if curr_bot_disp_msg else f"<i>{status_txt}</i>")
            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)
            elif upd_type == "final_response_and_insights":
                final_bot_resp_acc, insights_used_parsed = upd_data["response"], upd_data["insights_used"]
                status_txt = "Response complete."
                if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg or "(No text)")
                def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg)
                if curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
                    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
                        tmpfile.write(curr_bot_disp_msg)
                        temp_dl_file_path = tmpfile.name
                    def_dl_btn = gr.DownloadButton(value=temp_dl_file_path, visible=True, interactive=True) 
                else: def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
                insights_md = "### Insights Considered:\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
                def_detect_out_md = gr.Markdown(insights_md)
            yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
            if upd_type == "final_response_and_insights": break
    except Exception as e:
        logger.error(f"Chat handler error: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
        if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt: updated_gr_hist[-1] = (user_msg_txt, status_txt)
        else: updated_gr_hist.append((user_msg_txt, status_txt))
        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn); return
    if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
        current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
        hist_len_check = MAX_HISTORY_TURNS * 2
        if current_chat_session_history and current_chat_session_history[0]["role"] == "system": hist_len_check +=1
        if len(current_chat_session_history) > hist_len_check:
            current_chat_session_history = ([current_chat_session_history[0]] if current_chat_session_history[0]["role"] == "system" else []) + current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]
        threading.Thread(target=deferred_learning_and_memory_task, args=(user_msg_txt, final_bot_resp_acc, sel_prov_name, sel_model_disp_name, insights_used_parsed, ui_api_key.strip() if ui_api_key else None), daemon=True).start()
        status_txt = "Response complete. Background learning initiated."
    else: status_txt = "Processing finished; no response or error."
    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn)
    if temp_dl_file_path and os.path.exists(temp_dl_file_path):
        try: os.unlink(temp_dl_file_path)
        except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path}: {e_unlink}")

def ui_view_rules_action_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."
def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
    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 = 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()]
    total_to_process = len(potential_rules)
    if total_to_process == 0: return "No rules found in file to process."
    progress(0, desc="Starting rules upload...")
    for idx, rule_text in enumerate(potential_rules):
        rule_text = rule_text.strip()
        if not rule_text: continue
        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: Processed {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors/Invalid: {error_count}."
    logger.info(msg); return msg

def ui_view_memories_action_fn(): return get_all_memories_cached() or []
def ui_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 = []
    try:
        parsed_json = json.loads(content)
        memory_objects_to_process = parsed_json if isinstance(parsed_json, list) else [parsed_json]
    except json.JSONDecodeError:
        for line in content.splitlines():
            if line.strip():
                try: memory_objects_to_process.append(json.loads(line))
                except: format_error_count += 1
    if not memory_objects_to_process and format_error_count == 0: return "No valid memory objects found."
    total_to_process = len(memory_objects_to_process)
    if total_to_process == 0: return "No memory objects to process."
    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: 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}. Added: {added_count}, Format Errors: {format_error_count}, Save Errors: {save_error_count}."
    logger.info(msg); return msg

# --- UI Definition ---
# Using the CSS inspired by the "AI Code & Space Generator" example
custom_theme = gr.themes.Base(
    primary_hue="teal",
    secondary_hue="purple",
    neutral_hue="zinc",
    text_size="sm",
    spacing_size="md",
    radius_size="sm",
    font=["System UI", "sans-serif"]
)

custom_css = """
body { background: linear-gradient(to bottom right, #2c3e50, #34495e); color: #ecf0f1; margin:0; padding:0; font-family: 'System UI', sans-serif; overflow-x: hidden;}
.gradio-container { background: transparent !important; padding: 0 !important; /* Remove padding from the outermost container */ }
.main-interface-wrapper { max-width: 100%; padding: 1rem; box-sizing: border-box; min-height: 100vh; display: flex; flex-direction: column;} /* New wrapper */
.gr-column {gap: 0px !important;}
.gr-block.gr-group, .gr-tabs, .gr-accordion { background-color: rgba(44, 62, 80, 0.8) !important; border: 1px solid rgba(189, 195, 199, 0.2) !important; border-radius: 8px !important; padding: 1em; margin-bottom: 1em;}
.gr-tabitem { background-color: rgba(52, 73, 94, 0.75) !important; border-radius: 6px !important; padding: 1em !important; border: 1px solid rgba(189, 195, 199, 0.1) !important;}
.gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message, .gr-json, .gr-file input[type="file"], .gr-file button { border-color: rgba(189, 195, 199, 0.3) !important; background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; border-radius: 6px !important;}
.gr-textarea textarea, .gr-textbox input { color: #ecf0f1 !important; }
.gr-button.gr-button-primary { background-color: #1abc9c !important; color: white !important; border-color: #16a085 !important; }
.gr-button.gr-button-secondary { background-color: #9b59b6 !important; color: white !important; border-color: #8e44ad !important;}
.gr-button.gr-button-stop { background-color: #e74c3c !important; color: white !important; border-color: #c0392b !important; }
.gr-markdown { padding: 5px; } 
.gr-group .gr-markdown { padding: 0px; background-color: transparent !important; }
.gr-markdown h1, .gr-markdown h2, .gr-markdown h3 { color: #ecf0f1 !important; border-bottom-color: rgba(189, 195, 199, 0.3) !important; margin-top: 0.5em; margin-bottom: 0.5em;}
.gr-markdown h1 {font-size: 1.5rem;} .gr-markdown h2 {font-size: 1.25rem;} .gr-markdown h3 {font-size: 1.1rem;}
.gr-markdown p, .gr-markdown li { color: #ecf0f1 !important; }
.gr-markdown pre code { background-color: rgba(52, 73, 94, 0.95) !important; border-color: rgba(189, 195, 199, 0.3) !important; }
.gr-chatbot { background-color: rgba(44, 62, 80, 0.7) !important; border-color: rgba(189, 195, 199, 0.2) !important; }
.gr-chatbot .message { background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; border-color: rgba(189, 195, 199, 0.3) !important; }
.gr-chatbot .message.user { background-color: rgba(46, 204, 113, 0.9) !important; color: black !important; }
.gr-input-label > .label-text, .gr-dropdown-label > .label-text { color: #bdc3c7 !important; }
.status-bar { padding: 0.5rem 1rem; border-radius: 6px; margin-bottom: 1rem; background-color: rgba(44, 62, 80, 0.8) !important; }
.tabnav button { background-color: rgba(52, 73, 94, 0.8) !important; color: #ecf0f1 !important; border-bottom: 2px solid transparent !important;}
.tabnav button.selected { background-color: rgba(44, 62, 80, 0.95) !important; color: #1abc9c !important; border-bottom: 2px solid #1abc9c !important;}
.app-title-area { padding: 0.5rem 0; text-align: center;} /* Area for title and subtitle */
"""


with gr.Blocks(
    theme=gr.themes.Soft(),  # Use a clean, modern theme
    css="""
        .gr-button { margin: 5px; }
        .gr-textbox, .gr-text-area, .gr-dropdown { border-radius: 8px; }
        .gr-group { border: 1px solid #e0e0e0; border-radius: 8px; padding: 10px; }
        .gr-row { gap: 10px; }
        .gr-tab { border-radius: 8px; }
        .status-text { font-size: 0.9em; color: #555; }
    """
) as demo:
    # Header Section
    gr.Markdown(
        """
        # 🤖 AI Research Agent
        Your intelligent assistant for research and knowledge management
        """,
        elem_classes=["header"]
    )

    # Status Bar
    with gr.Row(variant="compact"):
        agent_stat_tb = gr.Textbox(
            label="Agent Status",
            value="Initializing systems...",
            interactive=False,
            elem_classes=["status-text"],
            scale=4
        )
        with gr.Column(scale=1, min_width=150):
            memory_backend_info_tb = gr.Textbox(
                label="Memory Backend",
                value=MEMORY_STORAGE_BACKEND,
                interactive=False,
                elem_classes=["status-text"]
            )
            sqlite_path_display = gr.Textbox(
                label="SQLite Path",
                value=MEMORY_SQLITE_PATH,
                interactive=False,
                visible=MEMORY_STORAGE_BACKEND == "SQLITE",
                elem_classes=["status-text"]
            )
            hf_repos_display = gr.Textbox(
                label="HF Repos",
                value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}",
                interactive=False,
                visible=MEMORY_STORAGE_BACKEND == "HF_DATASET",
                elem_classes=["status-text"]
            )

    with gr.Row():
        # Sidebar (unchanged)
        with gr.Sidebar():
            gr.Markdown("## ⚙️ Configuration")
            with gr.Group():
                gr.Markdown("### AI Model Settings")
                api_key_tb = gr.Textbox(
                    label="AI Provider API Key (Override)",
                    type="password",
                    placeholder="Uses .env if blank"
                )
                prov_sel_dd = gr.Dropdown(
                    label="AI Provider",
                    choices=get_available_providers(),
                    value=get_available_providers()[0] if get_available_providers() else None,
                    interactive=True
                )
                model_sel_dd = gr.Dropdown(
                    label="AI Model",
                    choices=get_model_display_names_for_provider(get_available_providers()[0]) if get_available_providers() else [],
                    value=get_default_model_display_name_for_provider(get_available_providers()[0]) if get_available_providers() else None,
                    interactive=True
                )
            with gr.Group():
                gr.Markdown("### System Prompt")
                sys_prompt_tb = gr.Textbox(
                    label="System Prompt Base",
                    lines=8,
                    value=DEFAULT_SYSTEM_PROMPT,
                    interactive=True
                )
            if MEMORY_STORAGE_BACKEND == "RAM":
                save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")

        # Main Content Area
        with gr.Column(scale=3):
            with gr.Tabs():
                with gr.TabItem("💬 Chat & Research"):
                    with gr.Group():
                        gr.Markdown("### AI Chat Interface")
                        main_chat_disp = gr.Chatbot(
                            label=None,
                            height=400,
                            bubble_full_width=False,
                            avatar_images=(None, "https://raw.githubusercontent.com/huggingface/brand-assets/main/hf-logo-with-title.png"),
                            show_copy_button=True,
                            render_markdown=True,
                            sanitize_html=True
                        )
                        with gr.Row(variant="compact"):
                            user_msg_tb = gr.Textbox(
                                show_label=False,
                                placeholder="Ask your research question...",
                                scale=7,
                                lines=1,
                                max_lines=3,
                                autofocus=True
                            )
                            send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
                        with gr.Accordion("📝 Detailed Response", open=False):
                            fmt_report_tb = gr.Textbox(
                                label="Full AI Response",
                                lines=8,
                                interactive=True,
                                show_copy_button=True
                            )
                            dl_report_btn = gr.DownloadButton(
                                "Download Report",
                                interactive=False,
                                visible=False
                            )
                        detect_out_md = gr.Markdown(visible=False)

                with gr.TabItem("🧠 Knowledge Base"):
                    with gr.Row(equal_height=True):
                        with gr.Column():
                            gr.Markdown("### 📜 Rules Management")
                            rules_disp_ta = gr.TextArea(
                                label=None,
                                lines=10,
                                placeholder="View or edit rules...",
                                interactive=True
                            )
                            with gr.Row(variant="compact"):
                                view_rules_btn = gr.Button("🔄 Load Rules")
                                save_edited_rules_btn = gr.Button("💾 Save Rules", variant="primary")
                                clear_rules_btn = gr.Button("🗑️ Clear Rules", variant="stop")
                            upload_rules_fobj = gr.File(
                                label="Upload Rules",
                                file_types=[".txt", ".jsonl"]
                            )
                            rules_stat_tb = gr.Textbox(
                                label="Rules Status",
                                interactive=False,
                                lines=1,
                                elem_classes=["status-text"]
                            )

                        with gr.Column():
                            gr.Markdown("### 📚 Memories Management")
                            mems_disp_json = gr.JSON(
                                label=None,
                                value={"memories": []}
                            )
                            with gr.Row(variant="compact"):
                                view_mems_btn = gr.Button("🔄 Load Memories")
                                clear_mems_btn = gr.Button("🗑️ Clear Memories", variant="stop")
                            upload_mems_fobj = gr.File(
                                label="Upload Memories",
                                file_types=[".jsonl"]
                            )
                            mems_stat_tb = gr.Textbox(
                                label="Memories Status",
                                interactive=False,
                                lines=1,
                                elem_classes=["status-text"]
                            )

    # Event Handlers (unchanged logic, cleaned up organization)
    def dyn_upd_model_dd(sel_prov_dyn: str):
        models_dyn = get_model_display_names_for_provider(sel_prov_dyn)
        def_model_dyn = get_default_model_display_name_for_provider(sel_prov_dyn)
        return gr.Dropdown(choices=models_dyn, value=def_model_dyn, interactive=True)

    prov_sel_dd.change(fn=dyn_upd_model_dd, inputs=prov_sel_dd, outputs=model_sel_dd)

    chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
    chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn]
    send_btn.click(fn=handle_gradio_chat_submit, inputs=chat_ins, outputs=chat_outs)
    user_msg_tb.submit(fn=handle_gradio_chat_submit, inputs=chat_ins, outputs=chat_outs)

    view_rules_btn.click(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)

    def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
        if not edited_rules_text.strip():
            return "No rules text to save."
        potential_rules = edited_rules_text.split("\n\n---\n\n")
        if len(potential_rules) == 1 and "\n" in edited_rules_text:
            potential_rules = [r.strip() for r in edited_rules_text.splitlines() if r.strip()]
        if not potential_rules:
            return "No rules found to process from editor."
        added, skipped, errors = 0, 0, 0
        total = len(potential_rules)
        progress(0, desc=f"Saving {total} rules from editor...")
        for idx, rule_text in enumerate(potential_rules):
            if not rule_text.strip():
                continue
            success, status_msg = add_rule_entry(rule_text.strip())
            if success:
                added += 1
            elif status_msg == "duplicate":
                skipped += 1
            else:
                errors += 1
            progress((idx + 1) / total)
        return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors}."

    save_edited_rules_btn.click(
        fn=save_edited_rules_action_fn,
        inputs=[rules_disp_ta],
        outputs=[rules_stat_tb],
        show_progress="full"
    ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)

    upload_rules_fobj.upload(
        fn=ui_upload_rules_action_fn,
        inputs=[upload_rules_fobj],
        outputs=[rules_stat_tb],
        show_progress="full"
    ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)

    clear_rules_btn.click(
        fn=lambda: "All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules.",
        outputs=rules_stat_tb
    ).then(fn=ui_view_rules_action_fn, outputs=rules_disp_ta)

    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
        def save_faiss_action_with_feedback_sidebar_fn():
            save_faiss_indices_to_disk()
            gr.Info("Attempted to save FAISS indices to disk.")
        save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None)

    view_mems_btn.click(fn=ui_view_memories_action_fn, outputs=mems_disp_json)

    upload_mems_fobj.upload(
        fn=ui_upload_memories_action_fn,
        inputs=[upload_mems_fobj],
        outputs=[mems_stat_tb],
        show_progress="full"
    ).then(fn=ui_view_memories_action_fn, outputs=mems_disp_json)

    clear_mems_btn.click(
        fn=lambda: "All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories.",
        outputs=mems_stat_tb
    ).then(fn=ui_view_memories_action_fn, outputs=mems_disp_json)

    def app_load_fn():
        initialize_memory_system()
        logger.info("App loaded. Memory system initialized.")
        backend_status = "AI Systems Initialized. Ready."
        rules_on_load = ui_view_rules_action_fn()
        mems_on_load = ui_view_memories_action_fn()
        return backend_status, rules_on_load, mems_on_load

    demo.load(fn=app_load_fn, inputs=None, outputs=[agent_stat_tb, rules_disp_ta, mems_disp_json])

if __name__ == "__main__":
    logger.info(f"Starting Gradio AI Research Mega Agent (v5.8 - UI Style/Layout Update, Memory: {MEMORY_STORAGE_BACKEND})...")
    app_port = int(os.getenv("GRADIO_PORT", 7860))
    app_server = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
    app_debug = os.getenv("GRADIO_DEBUG", "False").lower() == "true"
    app_share = os.getenv("GRADIO_SHARE", "False").lower() == "true"
    logger.info(f"Launching Gradio server: http://{app_server}:{app_port}. Debug: {app_debug}, Share: {app_share}")
    demo.queue().launch(server_name=app_server, server_port=app_port, debug=app_debug, share=app_share)
    logger.info("Gradio application shut down.")