app.py

# app.py
import os
import json
import re
import logging
from datetime import datetime
from dotenv import load_dotenv
import gradio as gr
import time
import tempfile
import xml.etree.ElementTree as ET 

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 ---
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:
    # ... (remains the same as v6.4)
    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):
    # ... (remains the same as v6.4)
    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 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 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>`.
ABSOLUTELY NO other text, explanations, or markdown should precede or follow this XML structure.

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):**
... (rest of reflection process - STEP 1, STEP 2, STEP 3 - remains the same) ...

**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
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, 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
        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..."
    # Initialize all potential output components
    updated_rules_text = ""
    updated_mems_json = []
    def_detect_out_md = gr.Markdown(visible=False) 
    def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True) 
    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))
        # Ensure all outputs are provided
        updated_rules_text = ui_refresh_rules_display_fn() # Get current rules
        updated_mems_json = ui_refresh_memories_display_fn() # Get current memories
        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 for chat update
    # Provide current state for rules/memories display even at this early stage
    current_rules_text_init = ui_refresh_rules_display_fn()
    current_mems_json_init = 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_init, current_mems_json_init)
    
    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 generated. Processing learning..." 
                if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
                if updated_gr_hist 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, interactive=True, show_copy_button=True)
                
                if curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
                    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
                        tmpfile.write(curr_bot_disp_msg)
                        temp_dl_file_path = tmpfile.name
                    def_dl_btn = gr.DownloadButton(value=temp_dl_file_path, visible=True, interactive=True) 
                else: 
                    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
                
                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, keeping rules/mems display static for now
            current_rules_text_during = ui_refresh_rules_display_fn()
            current_mems_json_during = 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_during, current_mems_json_during)
            if upd_type == "final_response_and_insights": break 
    
    except Exception as e:
        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)
        
        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)
        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}])
        # ... (history truncation logic remains same) ...
        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):]

        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 
        # ... (error handling logic remains same) ...
    else:
        status_txt = "Processing finished; no valid response or error occurred during main phase."

    # After learning (or if no learning was done), get the fresh data for UI
    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)
    
    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}")

# --- UI Functions for Rules and Memories (ui_refresh_..., ui_download_..., ui_upload_...) ---
# ... (These functions remain THE SAME as v6.4 - no changes needed here)
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

# --- Gradio UI Definition ---
with gr.Blocks(
    theme=gr.themes.Soft(),
    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:
    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.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=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")
            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")

        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
                            )
                            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(equal_height=True):
                        with gr.Column():
                            gr.Markdown("### 📜 Rules Management")
                            rules_disp_ta = gr.TextArea( # This is an output component for handle_gradio_chat_submit
                                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():
                            gr.Markdown("### 📚 Memories Management")
                            mems_disp_json = gr.JSON( # This is an output component for handle_gradio_chat_submit
                                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_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]
    # Add rules_disp_ta and mems_disp_json to the outputs of handle_gradio_chat_submit
    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] 
    
    chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
    
    send_btn.click(**chat_event_args) # Removed .then() calls for UI refresh here
    user_msg_tb.submit(**chat_event_args) # Removed .then() calls for UI refresh here

    # Rules Management events - .then() calls remain for these direct manipulations
    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn)

    def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
        # ... (function body remains the same)
        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."


    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)

    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)

    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)

    # Memories Management events
    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn)

    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)

    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)

    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)

    def app_load_fn():
        initialize_memory_system()
        logger.info("App loaded. Memory system initialized.")
        backend_status = "AI Systems Initialized. Ready."
        # Initial population of all relevant UI components
        rules_on_load = ui_refresh_rules_display_fn()
        mems_on_load = ui_refresh_memories_display_fn()
        return (
            backend_status, 
            rules_on_load, 
            mems_on_load,
            gr.Markdown(visible=False), # detect_out_md
            gr.Textbox(value="*Waiting...*", interactive=True), # fmt_report_tb
            gr.DownloadButton(interactive=False, value=None, visible=False), # dl_report_btn
            # The following are needed if we also want to clear chat on load for some reason
            # but usually, we don't. agent_stat_tb is already covered.
            # "", # user_msg_tb
            # None, # main_chat_disp (can't directly set to None, usually list)
        )
    
    initial_load_outputs = [
        agent_stat_tb, 
        rules_disp_ta, 
        mems_disp_json,
        detect_out_md, 
        fmt_report_tb, 
        dl_report_btn,
        # We need to ensure the outputs here match what app_load_fn returns
        # and what chat_outs expects if they overlap.
        # Since rules_disp_ta and mems_disp_json are already in chat_outs (implicitly now),
        # we don't need extra components here if app_load_fn returns values for them.
    ]
    # Adjust app_load_fn to provide initial values for all components in chat_outs.
    # The last two elements of chat_outs are rules_disp_ta and mems_disp_json.
    # The first 6 are: user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn
    
    # Corrected initial_load_outputs to match the *actual* components being updated by app_load_fn
    # and ensure it provides placeholder/initial values for components also in chat_outs.
    app_load_outputs = [
        agent_stat_tb,          # Updated by app_load_fn
        rules_disp_ta,          # Updated by app_load_fn
        mems_disp_json,         # Updated by app_load_fn
        detect_out_md,          # Initialized by app_load_fn
        fmt_report_tb,          # Initialized by app_load_fn
        dl_report_btn           # Initialized by app_load_fn
    ]
    demo.load(fn=app_load_fn, inputs=None, outputs=app_load_outputs)


if __name__ == "__main__":
    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")
    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.")