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.) ---
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 repair_json_string_newslines(json_like_string: str) -> str:
    """
    Attempts to fix unescaped literal newlines within string literals in a JSON-like string.
    This is a heuristic and focuses on the most common LLM error for this task.
    It converts literal \n, \r, \r\n characters found *inside* what it
    determines to be string literals into the two-character sequence "\\n".
    """
    output = []
    i = 0
    n = len(json_like_string)
    in_string_literal = False
    
    while i < n:
        char = json_like_string[i]
        
        if char == '"':
            # Check if this quote is escaped (preceded by an odd number of backslashes)
            num_backslashes = 0
            k = i - 1
            while k >= 0 and json_like_string[k] == '\\':
                num_backslashes += 1
                k -= 1
            
            if num_backslashes % 2 == 0: # This quote is not escaped, it's a delimiter
                in_string_literal = not in_string_literal
            output.append(char)
            i += 1
            continue

        if in_string_literal:
            if char == '\n': # Literal newline
                output.append('\\\\n') # Append literal backslash then 'n'
            elif char == '\r':
                if i + 1 < n and json_like_string[i+1] == '\n': # CRLF
                    output.append('\\\\n')
                    i += 1 # Consume the \n as well, it's part of the CRLF pair
                else: # Standalone CR
                    output.append('\\\\n')
            # We are not attempting to fix other escape issues here,
            # as newlines are the primary cause of "Unterminated string".
            # If we see a backslash, we assume it's either a correct escape
            # or part of content that doesn't break basic parsing like a raw newline does.
            else:
                output.append(char)
        else: # Not in string literal
            output.append(char)
        i += 1
            
    return "".join(output)

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.
Your output will be directly parsed by Python's `json.loads()` function.

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.

**ULTRA-CRITICAL JSON STRING FORMATTING RULES (ESPECIALLY FOR THE "insight" FIELD):**
-   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 `\\\\`.
-   **NEWLINES ARE THE #1 CAUSE OF ERRORS. READ CAREFULLY:**
    -   If the *content* of your `insight` text (or any other string value) spans multiple lines OR contains newline characters for formatting, EACH such newline character *inside the string's content* MUST be represented as the two-character sequence: a backslash followed by an 'n' (i.e., `\\n`).
    -   **DO NOT use a literal newline character (pressing Enter) inside a JSON string value.** This will break the JSON parser and result in an "Unterminated string" error.
    -   **Example of BAD JSON (literal newline):**
        `"insight": "This is line one.
        This is line two."`
        (This is INVALID because of the actual newline between "one." and "This". The parser stops reading the string at "one.")
    -   **Example of GOOD JSON (escaped newline using `\\n`):**
        `"insight": "This is line one.\\nThis is line two."`
        (This is VALID. The `\\n` sequence correctly tells the parser there's a newline *within* the string's content.)
-   The entire JSON response (the list `[...]`) can be pretty-printed with newlines *between* JSON elements (like after a comma or a brace) for readability, but *within* any quoted string value, the newline rule above is absolute.

**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:**
... (General Guidelines remain the same) ...

**Example of a comprehensive JSON output (Pay close attention to `\\n` for newlines within insight text if the insight content spans multiple lines):**
[
  {"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'.\\nMy purpose is to assist the user with research and tasks."},
  {"action": "add", "insight": "[CORE_RULE|0.9] I am capable of searching the internet for current weather information if asked.\\nThis capability was confirmed on [date] based on user query regarding weather."},
  {"action": "add", "insight": "[RESPONSE_PRINCIPLE|0.8] When user provides positive feedback like 'good job', acknowledge it warmly with phrases like 'Thank you!' or 'Glad I could help!'."},
  {"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.\\nThis rule was updated due to user preference for more contextual continuity when appropriate."}
]
"""
        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.
**ULTRA-IMPORTANT FINAL REMINDER: YOUR ENTIRE RESPONSE MUST BE A SINGLE, VALID JSON LIST. DOUBLE-CHECK ALL STRING VALUES, ESPECIALLY THE 'insight' TEXT, FOR CORRECTLY ESCAPED NEWLINES (using the two characters `\\` and `n`, i.e., `\\n`) AND QUOTES (using `\\"`). AN UNESCAPED LITERAL NEWLINE CHARACTER (ASCII 10) WITHIN AN 'insight' STRING VALUE WILL BREAK THE JSON AND CAUSE THE UPDATE TO FAIL. BE METICULOUS. VALIDATE YOUR JSON STRUCTURE AND STRING CONTENTS BEFORE FINALIZING YOUR OUTPUT.**
"""
        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_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=2000))).strip()
        ops, processed_count = [], 0
        
        json_match_ops = re.search(r"```json\s*(\[.*?\])\s*```", raw_ops_json_full, re.DOTALL|re.I) or \
                         re.search(r"(\[.*?\])", raw_ops_json_full, re.DOTALL)

        ops_json_str = None
        if json_match_ops:
            ops_json_str = json_match_ops.group(1)
            try: 
                ops = json.loads(ops_json_str)
            except json.JSONDecodeError as e: 
                error_char_pos = e.pos
                context_window = 40
                start_context = max(0, error_char_pos - context_window)
                end_context = min(len(ops_json_str), error_char_pos + context_window)
                problem_context = ops_json_str[start_context:end_context].replace("\n", "\\n") # Show newlines escaped in log
                
                logger.warning(
                    f"DEFERRED [{task_id}]: Initial JSON ops parse error: {e}. "
                    f"Error at char {error_char_pos}. Context around error: '...{problem_context}...'. "
                    f"Attempting to repair newlines in the JSON string."
                )
                # logger.debug(f"DEFERRED [{task_id}]: Problematic JSON string before repair:\n>>>>>>>>>>\n{ops_json_str}\n<<<<<<<<<<")

                repaired_json_str = repair_json_string_newslines(ops_json_str)
                # logger.debug(f"DEFERRED [{task_id}]: JSON string after repair attempt:\n>>>>>>>>>>\n{repaired_json_str}\n<<<<<<<<<<")
                
                try:
                    ops = json.loads(repaired_json_str)
                    logger.info(f"DEFERRED [{task_id}]: JSON successfully parsed after repair attempt.")
                except json.JSONDecodeError as e2:
                    logger.error(
                        f"DEFERRED [{task_id}]: JSON ops parse error EVEN AFTER REPAIR: {e2}. "
                        f"The repair attempt was not sufficient. Skipping insight operations for this turn."
                    )
                    # Log the string that failed after repair for further debugging
                    # logger.error(f"DEFERRED [{task_id}]: Repaired JSON string that still failed:\n>>>>>>>>>>\n{repaired_json_str}\n<<<<<<<<<<")
                    ops = [] 
        else:
             logger.info(f"DEFERRED [{task_id}]: No JSON list structure (e.g., starting with '[') found in LLM output. Full raw output:\n>>>>>>>>>>\n{raw_ops_json_full}\n<<<<<<<<<<")
        
        if isinstance(ops, list) and ops:
            logger.info(f"DEFERRED [{task_id}]: LLM provided {len(ops)} insight ops to process.")
            for op_idx, op in enumerate(ops):
                if not isinstance(op, dict): 
                    logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped non-dict item in ops list: {op}")
                    continue
                
                action = op.get("action","").lower()
                insight_text_raw = op.get("insight") 
                old_insight_raw = op.get("old_insight_to_replace")

                if not isinstance(insight_text_raw, str):
                    logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to non-string insight_text (type: {type(insight_text_raw)}): {op}")
                    continue
                insight_text = insight_text_raw.strip()

                old_insight = None
                if old_insight_raw is not None:
                    if not isinstance(old_insight_raw, str):
                        logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to non-string/non-null old_insight_to_replace (type: {type(old_insight_raw)}): {op}")
                        continue
                    old_insight = old_insight_raw.strip()
                
                if not insight_text or not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I|re.DOTALL): 
                    logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to invalid/empty insight_text format: '{insight_text[:100]}...' from op: {op}")
                    continue

                if action == "add":
                    success, status_msg = add_rule_entry(insight_text)
                    if success: processed_count +=1
                    else: logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (add): Failed to add rule '{insight_text[:50]}...'. Status: {status_msg}")
                elif action == "update":
                    if old_insight: 
                        if old_insight != insight_text: 
                            remove_success = remove_rule_entry(old_insight)
                            if not remove_success:
                                 logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (update): Failed to remove old rule '{old_insight[:50]}...' before adding new. This might lead to duplicates if the new rule is different.")
                        else: 
                             logger.info(f"DEFERRED [{task_id}]: Op {op_idx} (update): Old insight is identical to new insight. Skipping removal and effectively treating as 'add if not present'.")
                    
                    success, status_msg = add_rule_entry(insight_text) 
                    if success: processed_count +=1
                    else: logger.warning(f"DEFERRED [{task_id}]: Op {op_idx} (update): Failed to add/update rule '{insight_text[:50]}...'. Status: {status_msg}")
                else:
                    logger.warning(f"DEFERRED [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}': {op}")

            logger.info(f"DEFERRED [{task_id}]: Processed {processed_count} insight ops out of {len(ops)} received.")
        elif not json_match_ops : 
            pass 
        else: 
            logger.info(f"DEFERRED [{task_id}]: No valid list of insight ops from LLM after parsing was attempted. Raw match (ops_json_str - first 500 chars):\n>>>>>>>>>>\n{ops_json_str[:500] if ops_json_str else 'N/A'}\n<<<<<<<<<<")
            
    except Exception as e: logger.error(f"DEFERRED [{task_id}]: CRITICAL ERROR in deferred task: {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 = 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))
        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, 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:\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 (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]}"
        error_message_for_chat = f"Sorry, an error occurred: {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)
        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 valid response or error occurred."
        if final_bot_resp_acc.startswith("Error:"): 
            status_txt = final_bot_resp_acc 
            if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                 updated_gr_hist[-1] = (user_msg_txt, final_bot_resp_acc)
            def_fmt_out_txt = gr.Textbox(value=final_bot_resp_acc, interactive=True)
            def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)

    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}")

# --- UI Functions for Rules and Memories ---
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
    # For .txt, split by '---'
    if uploaded_file_obj.name.lower().endswith(".txt"):
        potential_rules = content.split("\n\n---\n\n")
        if len(potential_rules) == 1 and "\n" in content: # Fallback for simple newline separation in .txt
             potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
    elif uploaded_file_obj.name.lower().endswith(".jsonl"):
        potential_rules = []
        for line in content.splitlines():
            if line.strip():
                try:
                    # Expect each line to be a JSON string containing the rule text
                    rule_text_in_json_string = json.loads(line)
                    if isinstance(rule_text_in_json_string, str):
                        potential_rules.append(rule_text_in_json_string)
                    else:
                        logger.warning(f"Rule Upload: Skipped non-string rule from JSONL: {rule_text_in_json_string}")
                        error_count +=1
                except json.JSONDecodeError:
                    logger.warning(f"Rule Upload: Failed to parse JSONL line for rule: {line}")
                    error_count +=1
    else:
        return "Unsupported file type for rules. Please use .txt or .jsonl."


    total_to_process = len(potential_rules)
    if total_to_process == 0 and error_count == 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 # Increment error count for add_rule_entry failures too
        progress((idx + 1) / total_to_process if total_to_process > 0 else 1, desc=f"Processed {idx+1}/{total_to_process} rules...")
    
    msg = f"Rules Upload: Total lines/segments processed: {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors/Invalid: {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): # Single object
                memory_objects_to_process = [parsed_json]
            else:
                format_error_count = 1 # Not a list or object
        except json.JSONDecodeError:
            format_error_count = 1 # Invalid JSON
    elif file_ext == ".jsonl":
        for line_num, line in enumerate(content.splitlines()):
            if line.strip():
                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."
    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


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(
                                label="Current Rules (Read-only, Edit via Upload/Save)", lines=10, 
                                placeholder="Rules will appear here. Use 'Save Edited Text' or 'Upload File' to modify.", 
                                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(
                                label="Current Memories (Read-only)", 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]
    chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn]
    
    chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
    
    send_btn_click_event = send_btn.click(**chat_event_args)
    user_msg_submit_event = user_msg_tb.submit(**chat_event_args)

    for event in [send_btn_click_event, user_msg_submit_event]:
        event.then(fn=ui_refresh_rules_display_fn, inputs=None, outputs=rules_disp_ta, show_progress=False)
        event.then(fn=ui_refresh_memories_display_fn, inputs=None, outputs=mems_disp_json, show_progress=False)

    # Rules Management events
    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()):
        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."
        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), 
            gr.Textbox(value="*Waiting...*", interactive=True), 
            gr.DownloadButton(interactive=False, value=None, visible=False)
        )
    
    initial_load_outputs = [
        agent_stat_tb, 
        rules_disp_ta, 
        mems_disp_json,
        detect_out_md, 
        fmt_report_tb, 
        dl_report_btn  
    ]
    demo.load(fn=app_load_fn, inputs=None, outputs=initial_load_outputs)


if __name__ == "__main__":
    logger.info(f"Starting Gradio AI Research Mega Agent (v6.1 - JSON Repair & UI Polish, 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.")