Update app.py

app.py

@@ -31,13 +31,23 @@ import time
 import tempfile
 import xml.etree.ElementTree as ET
 
+# --- New Imports for Image KV Functionality ---
+import io
+import struct
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from cryptography.hazmat.primitives import hashes
+from cryptography.exceptions import InvalidTag
+# --- End New Imports ---
+
 load_dotenv() # Load .env file, but our settings above will take precedence if set.
 
 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,
@@ -71,6 +81,240 @@ logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={
 logger.info(f"Startup loading: Rules from {LOAD_RULES_FILE or 'None'}, Memories from {LOAD_MEMORIES_FILE or 'None'}")
 
 
+# --- KV to Image Functions (Constants and Implementation) ---
+KEY_SIZE = 32
+SALT_SIZE = 16
+NONCE_SIZE = 12
+TAG_SIZE = 16
+PBKDF2_ITERATIONS = 480000
+LENGTH_HEADER_SIZE = 4  # struct.pack('>I') uses 4 bytes
+PREFERRED_FONTS = ["Arial", "Helvetica", "DejaVu Sans", "Verdana", "Calibri", "sans-serif"]
+MAX_KEYS_TO_DISPLAY_OVERLAY = 15
+
+def _get_font(preferred_fonts, base_size):
+    fp = None
+    safe_base_size = int(base_size)
+    if safe_base_size <= 0: safe_base_size = 10
+    for n in preferred_fonts:
+        try: ImageFont.truetype(n.lower()+".ttf",10); fp=n.lower()+".ttf"; break
+        except IOError:
+            try: ImageFont.truetype(n,10); fp=n; break
+            except IOError: continue
+    if fp:
+        try: return ImageFont.truetype(fp, safe_base_size)
+        except IOError: logger.warning(f"Font '{fp}' load failed with size {safe_base_size}. Defaulting.")
+    try: return ImageFont.load_default(size=safe_base_size)
+    except TypeError: return ImageFont.load_default()
+
+def set_pil_image_format_to_png(image:Image.Image)->Image.Image:
+    buf=io.BytesIO(); image.save(buf,format='PNG'); buf.seek(0)
+    reloaded=Image.open(buf); reloaded.format="PNG"; return reloaded
+
+def _derive_key(pw:str,salt:bytes)->bytes:
+    kdf=PBKDF2HMAC(algorithm=hashes.SHA256(),length=KEY_SIZE,salt=salt,iterations=PBKDF2_ITERATIONS)
+    return kdf.derive(pw.encode('utf-8'))
+
+def encrypt_data(data:bytes,pw:str)->bytes:
+    s=os.urandom(SALT_SIZE);k=_derive_key(pw,s);a=AESGCM(k);n=os.urandom(NONCE_SIZE)
+    ct=a.encrypt(n,data,None); return s+n+ct
+
+def decrypt_data(payload:bytes,pw:str)->bytes:
+    ml=SALT_SIZE+NONCE_SIZE+TAG_SIZE;
+    if len(payload)<ml: raise ValueError("Payload too short.")
+    s,n,ct_tag=payload[:SALT_SIZE],payload[SALT_SIZE:SALT_SIZE+NONCE_SIZE],payload[SALT_SIZE+NONCE_SIZE:]
+    k=_derive_key(pw,s);a=AESGCM(k)
+    try: return a.decrypt(n,ct_tag,None)
+    except InvalidTag: raise ValueError("Decryption failed: Invalid password/corrupted data.")
+    except Exception as e: logger.error(f"Decrypt error: {e}",exc_info=True); raise
+
+def _d2b(d:bytes)->str: return ''.join(format(b,'08b') for b in d)
+def _b2B(b:str)->bytes:
+    if len(b)%8!=0: raise ValueError("Bits not multiple of 8.")
+    return bytes(int(b[i:i+8],2) for i in range(0,len(b),8))
+
+def embed_data_in_image(img_obj:Image.Image,data:bytes)->Image.Image:
+    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
+    lb=struct.pack('>I',len(data));fp=lb+data;db=_d2b(fp);nb=len(db)
+    if nb>len(fpx): raise ValueError(f"Data too large: {nb} bits needed, {len(fpx)} available.")
+    for i in range(nb): fpx[i]=(fpx[i]&0xFE)|int(db[i])
+    spx=fpx.reshape(px.shape); return Image.fromarray(spx.astype(np.uint8),'RGB')
+
+def extract_data_from_image(img_obj:Image.Image)->bytes:
+    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
+    hbc=LENGTH_HEADER_SIZE*8
+    if len(fpx)<hbc: raise ValueError("Image too small for header.")
+    lb="".join(str(fpx[i]&1) for i in range(hbc))
+    try: pl=struct.unpack('>I',_b2B(lb))[0]
+    except Exception as e: raise ValueError(f"Header decode error: {e}")
+    if pl==0: return b""
+    if pl>(len(fpx)-hbc)/8: raise ValueError("Header len corrupted or > capacity.")
+    tpb=pl*8; so=hbc; eo=so+tpb
+    if len(fpx)<eo: raise ValueError("Image truncated or header corrupted.")
+    pb="".join(str(fpx[i]&1) for i in range(so,eo)); return _b2B(pb)
+
+def parse_kv_string_to_dict(kv_str:str)->dict:
+    if not kv_str or not kv_str.strip(): return {}
+    dd={};
+    for ln,ol in enumerate(kv_str.splitlines(),1):
+        l=ol.strip()
+        if not l or l.startswith('#'): continue
+        lc=l.split('#',1)[0].strip();
+        if not lc: continue
+        p=lc.split('=',1) if '=' in lc else lc.split(':',1) if ':' in lc else []
+        if len(p)!=2: raise ValueError(f"L{ln}: Invalid format '{ol}'.")
+        k,v=p[0].strip(),p[1].strip()
+        if not k: raise ValueError(f"L{ln}: Empty key in '{ol}'.")
+        dd[k]=v
+    return dd
+
+def generate_brain_carrier_image(w=800, h=800, msg="iLearn Knowledge Base") -> Image.Image:
+    """
+    Generates a carrier image with a radial gradient, a central text-based icon, and a message.
+    """
+    # --- Color Palette ---
+    bg_center_color = (210, 220, 255)    # Light blue center
+    bg_outer_color = (0, 53, 139)       # Deep blue edges
+    icon_color = (180, 220, 255, 200) # Light, slightly transparent color for the icon text
+    text_color = (230, 235, 245)      # Bright text color for the main message
+    shadow_color = (0, 0, 0, 128)     # Text shadow
+    border_color = (255, 255, 255, 50) # Subtle white border
+
+    # --- Font and Icon Selection ---
+    PREFERRED_MONO_FONTS = ["Courier New", "Consolas", "Menlo", "Monaco", "Courier", "monospace"]
+    
+    # Use the single line "thinking" bubble as the main icon.
+    ascii_art_icon = ". o O ( hmm... )"
+
+    # --- Setup ---
+    img = Image.new("RGBA", (w, h), bg_outer_color)
+    draw = ImageDraw.Draw(img)
+
+    # --- 1. Draw Radial Gradient Background ---
+    center_x, center_y = w / 2, h / 2
+    max_radius = int((center_x**2 + center_y**2)**0.5)
+
+    for r in range(max_radius, 0, -3):
+        ratio = 1 - (r / max_radius)
+        inter_color = tuple(int(bg_outer_color[i] + (bg_center_color[i] - bg_center_color[i]) * ratio) for i in range(3))
+        box = [center_x - r, center_y - r, center_x + r, center_y + r]
+        draw.ellipse(box, fill=inter_color)
+
+    # --- 2. Draw Text-based Icon ---
+    # Make the icon much larger, based on image width.
+    icon_font_size = max(24, int(w / 15))
+    icon_font = _get_font(PREFERRED_MONO_FONTS, icon_font_size)
+
+    # Position the icon lower, closer to the bottom text.
+    icon_cx = w / 2
+    icon_cy = h * 0.58 
+    
+    # Draw the single-line text icon, centered.
+    draw.text(
+        (icon_cx, icon_cy), 
+        ascii_art_icon,
+        font=icon_font, 
+        fill=icon_color, 
+        anchor="mm"
+    )
+
+    # --- 3. Draw Text Message at the bottom ---
+    text_font_size = max(18, int(w / 30))
+    text_font = _get_font(PREFERRED_FONTS, text_font_size)
+    text_y_pos = h * 0.80 # Moved text slightly lower to give the new icon more space
+    shadow_offset = max(1, int(text_font_size / 20))
+
+    draw.text((center_x, text_y_pos), msg, font=text_font, fill=shadow_color, anchor="ms", stroke_width=shadow_offset*2, stroke_fill=shadow_color)
+    draw.text((center_x, text_y_pos), msg, font=text_font, fill=text_color, anchor="ms")
+
+    # --- 4. Draw Border ---
+    border_width = max(2, int(min(w,h) / 150))
+    draw.rectangle([(0,0), (w-1, h-1)], outline=border_color, width=border_width)
+
+    # --- Finalize ---
+    final_image_rgb = Image.new("RGB", img.size, (0, 0, 0))
+    final_image_rgb.paste(img, (0, 0), img) 
+
+    return final_image_rgb
+
+def _get_text_measurement(draw_obj, text_str, font_obj):
+    if hasattr(draw_obj, 'textbbox'):
+        try:
+            bbox = draw_obj.textbbox((0, 0), text_str, font=font_obj)
+            width = bbox[2] - bbox[0]
+            height = bbox[3] - bbox[1]
+            return width, height
+        except Exception: pass
+    try:
+        if hasattr(font_obj, 'getsize'): return font_obj.getsize(text_str)
+        width, height = draw_obj.textsize(text_str, font=font_obj)
+        return width, height
+    except AttributeError:
+        try:
+            char_width_approx = font_obj.size * 0.6
+            char_height_approx = font_obj.size
+            return int(len(text_str) * char_width_approx), int(char_height_approx)
+        except: return len(text_str) * 8, 10
+
+def draw_key_list_dropdown_overlay(image: Image.Image, keys: list[str] = None, title: str = "Data Embedded") -> Image.Image:
+    if not title and (keys is None or not keys):
+        return set_pil_image_format_to_png(image.copy())
+    img_overlayed = image.copy(); draw = ImageDraw.Draw(img_overlayed)
+    margin = 10; padding = {'title_x':10,'title_y':6,'key_x':10,'key_y':5}; line_spacing = 4
+    title_bg_color=(60,60,60); title_text_color=(230,230,90)
+    key_list_bg_color=(50,50,50); key_text_color=(210,210,210); ellipsis_color=(170,170,170)
+    OVERLAY_TARGET_WIDTH_RATIO = 0.30; MIN_OVERLAY_WIDTH_PX = 180; MAX_OVERLAY_WIDTH_PX = 500
+    final_overlay_box_width = min(max(int(image.width*OVERLAY_TARGET_WIDTH_RATIO),MIN_OVERLAY_WIDTH_PX),MAX_OVERLAY_WIDTH_PX)
+    final_overlay_box_width = min(final_overlay_box_width, image.width - 2 * margin)
+    TITLE_FONT_HEIGHT_RATIO=0.030; TITLE_FONT_OVERLAY_WIDTH_RATIO=0.08; MIN_TITLE_FONT_SIZE=14; MAX_TITLE_FONT_SIZE=28
+    title_font_size = min(max(min(int(image.height*TITLE_FONT_HEIGHT_RATIO),int(final_overlay_box_width*TITLE_FONT_OVERLAY_WIDTH_RATIO)),MIN_TITLE_FONT_SIZE),MAX_TITLE_FONT_SIZE)
+    title_font = _get_font(PREFERRED_FONTS, title_font_size)
+    KEY_FONT_HEIGHT_RATIO=0.025; KEY_FONT_OVERLAY_WIDTH_RATIO=0.07; MIN_KEY_FONT_SIZE=12; MAX_KEY_FONT_SIZE=22
+    key_font_size = min(max(min(int(image.height*KEY_FONT_HEIGHT_RATIO),int(final_overlay_box_width*KEY_FONT_OVERLAY_WIDTH_RATIO)),MIN_KEY_FONT_SIZE),MAX_KEY_FONT_SIZE)
+    key_font = _get_font(PREFERRED_FONTS, key_font_size)
+    actual_title_w, actual_title_h = _get_text_measurement(draw, title, title_font)
+    disp_keys, actual_key_text_widths, total_keys_render_h, key_line_heights = [],[],0,[]
+    if keys:
+        temp_disp_keys=keys[:MAX_KEYS_TO_DISPLAY_OVERLAY-1]+[f"... ({len(keys)-(MAX_KEYS_TO_DISPLAY_OVERLAY-1)} more)"] if len(keys)>MAX_KEYS_TO_DISPLAY_OVERLAY else keys
+        for kt in temp_disp_keys:
+            disp_keys.append(kt); kw, kh = _get_text_measurement(draw, kt, key_font)
+            actual_key_text_widths.append(kw); key_line_heights.append(kh); total_keys_render_h += kh
+        if len(disp_keys)>1: total_keys_render_h += line_spacing*(len(disp_keys)-1)
+    title_bar_h = actual_title_h + 2*padding['title_y']; title_bar_x1=image.width-margin; title_bar_x0=title_bar_x1-final_overlay_box_width; title_bar_y0=margin; title_bar_y1=title_bar_y0+title_bar_h
+    draw.rectangle([(title_bar_x0,title_bar_y0),(title_bar_x1,title_bar_y1)],fill=title_bg_color)
+    available_width_for_title_text = final_overlay_box_width - 2*padding['title_x']
+    title_text_draw_x = title_bar_x0+padding['title_x'] + (available_width_for_title_text-actual_title_w)/2 if actual_title_w<=available_width_for_title_text else title_bar_x0+padding['title_x']
+    title_text_draw_y = title_bar_y0 + padding['title_y']
+    draw.text((title_text_draw_x, title_text_draw_y), title, font=title_font, fill=title_text_color)
+    if disp_keys:
+        key_list_box_h_ideal = total_keys_render_h + 2*padding['key_y']; key_list_x0, key_list_x1 = title_bar_x0, title_bar_x1
+        key_list_y0 = title_bar_y1; key_list_y1 = min(key_list_y0+key_list_box_h_ideal, image.height-margin)
+        current_key_list_box_h = key_list_y1 - key_list_y0
+        draw.rectangle([(key_list_x0,key_list_y0),(key_list_x1,key_list_y1)],fill=key_list_bg_color)
+        current_text_y = key_list_y0+padding['key_y']; available_text_width_for_keys = final_overlay_box_width-2*padding['key_x']
+        for i, key_text_item in enumerate(disp_keys):
+            if i>=len(key_line_heights): break
+            current_key_h=key_line_heights[i]
+            if current_text_y+current_key_h > key_list_y0+current_key_list_box_h-padding['key_y']:
+                _, ellipsis_h = _get_text_measurement(draw,"...",key_font)
+                if current_text_y+ellipsis_h <= key_list_y0+current_key_list_box_h-padding['key_y']:
+                    ellipsis_w, _ = _get_text_measurement(draw,"...",key_font)
+                    draw.text((key_list_x0+(final_overlay_box_width-ellipsis_w)/2, current_text_y), "...", font=key_font, fill=ellipsis_color)
+                break
+            original_key_text_w = actual_key_text_widths[i]; text_to_draw = key_text_item
+            if original_key_text_w > available_text_width_for_keys:
+                temp_text = key_text_item
+                while _get_text_measurement(draw, temp_text+"...", key_font)[0] > available_text_width_for_keys and len(temp_text) > 0: temp_text=temp_text[:-1]
+                text_to_draw = temp_text+"..." if len(temp_text)<len(key_text_item) else temp_text
+            final_key_text_w, _ = _get_text_measurement(draw, text_to_draw, key_font)
+            key_text_draw_x = key_list_x0+padding['key_x']+max(0,(available_text_width_for_keys-final_key_text_w)/2)
+            text_color_to_use = ellipsis_color if "..." in text_to_draw or f"... ({len(keys)-(MAX_KEYS_TO_DISPLAY_OVERLAY-1)} more)"==key_text_item else key_text_color
+            draw.text((key_text_draw_x, current_text_y), text_to_draw, font=key_font, fill=text_color_to_use)
+            current_text_y += current_key_h
+            if i < len(disp_keys)-1: current_text_y += line_spacing
+    return set_pil_image_format_to_png(img_overlayed)
+# --- END KV to Image Functions ---
+
+
 # --- Helper Functions ---
 def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str, list[dict]]:
     if not retrieved_insights_list:
@@ -691,8 +935,21 @@ def load_memories_from_file(filepath: str | None):
     return added_count, format_error_count, save_error_count
 
 
-# --- UI Functions for Rules and Memories (ui_refresh_..., ui_download_..., ui_upload_...) ---
+# --- UI Functions for Rules and Memories (Text and Image) ---
+def convert_kb_to_kv_string(rules: list[str], memories: list[dict]) -> str:
+    """Serializes rules and memories into a single key-value string for image embedding."""
+    lines = ["# iLearn Knowledge Base Export", f"# Exported on: {datetime.utcnow().isoformat()}Z"]
+    lines.append("\n# --- RULES ---")
+    for i, rule_text in enumerate(rules):
+        lines.append(f"rule_{i+1} = {json.dumps(rule_text)}")
+
+    lines.append("\n# --- MEMORIES ---")
+    for i, mem_dict in enumerate(memories):
+        lines.append(f"memory_{i+1} = {json.dumps(mem_dict)}")
+    return "\n".join(lines)
+
 def ui_refresh_rules_display_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."
+def ui_refresh_memories_display_fn(): return get_all_memories_cached() or []
 
 def ui_download_rules_action_fn():
     rules_content = "\n\n---\n\n".join(get_all_rules_cached())
@@ -715,10 +972,8 @@ def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
     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:
@@ -739,15 +994,10 @@ def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
                     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."
-
+    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)
@@ -755,23 +1005,18 @@ def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
         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")
@@ -792,42 +1037,28 @@ def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
     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]
+            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
+                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
+            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))
+                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."
-
+                    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"]):
@@ -835,308 +1066,343 @@ def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
             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
+            logger.warning(f"Memories Upload: Skipped invalid memory object structure: {str(mem_data)[:100]}"); format_error_count += 1
         progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} memories...")
-
     msg = f"Memories Upload: Processed {total_to_process} objects. Added: {added_count}, Format/Structure Errors: {format_error_count}, Save Errors: {save_error_count}."
     logger.info(msg); return msg
 
 def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
-    # --- DEMO MODE CHANGE ---
     if DEMO_MODE:
         gr.Warning("Saving edited rules is disabled in Demo Mode.")
         return "Saving edited rules is disabled in Demo Mode."
-
-    if not edited_rules_text.strip():
-        return "No rules text to save."
-
+    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."
-
+    if not potential_rules: return "No rules found to process from editor."
     added, skipped, errors = 0, 0, 0
     unique_rules_to_process = sorted(list(set(filter(None, [r.strip() for r in potential_rules]))))
-
     total_unique = len(unique_rules_to_process)
     if total_unique == 0: return "No unique, non-empty rules found in editor text."
-
     progress(0, desc=f"Saving {total_unique} unique rules from editor...")
-
     for idx, rule_text in enumerate(unique_rules_to_process):
         success, status_msg = add_rule_entry(rule_text)
         if success: added += 1
         elif status_msg == "duplicate": skipped += 1
         else: errors += 1
         progress((idx + 1) / total_unique, desc=f"Processed {idx+1}/{total_unique} rules...")
-
     return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors} from {total_unique} unique rules in text."
 
+def ui_download_kb_as_image_fn(password: str, progress=gr.Progress()):
+    """Generates and provides a downloadable image with embedded KB data."""
+    progress(0, desc="Fetching knowledge base...")
+    rules, memories = get_all_rules_cached(), get_all_memories_cached()
+    if not rules and not memories:
+        gr.Warning("Knowledge base is empty. Nothing to download.")
+        return None
+    
+    progress(0.2, desc="Serializing data...")
+    kv_string = convert_kb_to_kv_string(rules, memories)
+    data_bytes = kv_string.encode('utf-8')
+    
+    if password and password.strip():
+        progress(0.4, desc="Encrypting data...")
+        try: 
+            data_bytes = encrypt_data(data_bytes, password.strip())
+            gr.Info("Data encrypted successfully.")
+        except Exception as e: 
+            logger.error(f"KB ImgDL: Encrypt failed: {e}")
+            gr.Error(f"Encryption failed: {e}")
+            return None
+
+    progress(0.6, desc="Generating carrier image...")
+    carrier_image = generate_brain_carrier_image(w=800, h=800, msg="iLearn Knowledge Base")
+    
+    try: 
+        progress(0.7, desc="Embedding data...")
+        embedded_image = embed_data_in_image(carrier_image, data_bytes)
+    except ValueError as e: 
+        logger.error(f"KB ImgDL: Embed failed: {e}")
+        gr.Error(f"Data is too large for this image size: {e}")
+        return None
+
+    progress(0.8, desc="Adding visual overlay...")
+    keys_for_overlay = [f"Rule Count: {len(rules)}", f"Memory Count: {len(memories)}", "---"]
+    for r in rules[:5]:
+        match = re.search(r"\](.*)", r, re.DOTALL)
+        rule_content = match.group(1).strip() if match else r
+        keys_for_overlay.append(f"Rule: {rule_content[:40]}...")
+    if len(rules) > 5: keys_for_overlay.append("...")
+    
+    title_overlay = "Encrypted Data" if password and password.strip() else "Embedded Data"
+    final_image = draw_key_list_dropdown_overlay(embedded_image, keys=keys_for_overlay, title=title_overlay)
+
+    progress(0.9, desc="Saving final image...")
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
+            final_image.save(tmpfile, format="PNG")
+            tmp_path = tmpfile.name
+        
+        progress(1.0, desc="Download triggered!")
+        gr.Info("Download should start automatically.")
+        return tmp_path
+    except Exception as e: 
+        logger.error(f"KB ImgDL: Save failed: {e}")
+        gr.Error(f"Failed to save final image: {e}")
+        return None
+
+def ui_upload_kb_from_image_fn(uploaded_image_filepath: str, password: str, progress=gr.Progress()):
+    """Extracts KB data from an uploaded image and adds it to the system."""
+    if DEMO_MODE: 
+        gr.Warning("Uploading is disabled in Demo Mode.")
+        return "Upload disabled in Demo Mode."
+    if not uploaded_image_filepath: 
+        return "No image file provided or pasted."
+        
+    progress(0, desc="Loading image...")
+    try: 
+        img = Image.open(uploaded_image_filepath)
+    except Exception as e: 
+        logger.error(f"KB ImgUL: Open fail: {e}")
+        return f"Error: Could not open image file: {e}"
+
+    progress(0.2, desc="Extracting data from image...")
+    try:
+        extracted_bytes = extract_data_from_image(img)
+        if not extracted_bytes: return "No data found embedded in the image."
+    except ValueError as e: 
+        logger.error(f"KB ImgUL: Extract fail: {e}")
+        return f"Error extracting data: {e}"
+    except Exception as e: 
+        logger.error(f"KB ImgUL: Extract error: {e}", exc_info=True)
+        return f"Unexpected extraction error: {e}"
+    
+    kv_string = ""
+    try:
+        if extracted_bytes[:20].decode('utf-8', errors='ignore').strip().startswith("# iLearn"):
+             kv_string = extracted_bytes.decode('utf-8')
+             progress(0.4, desc="Parsing data...")
+        elif password and password.strip():
+            progress(0.3, desc="Attempting decryption...")
+            kv_string = decrypt_data(extracted_bytes, password.strip()).decode('utf-8')
+            progress(0.4, desc="Parsing decrypted data...")
+        else: return "Data appears encrypted, but no password was provided."
+    except (UnicodeDecodeError, InvalidTag, ValueError) as e:
+        if isinstance(e, (InvalidTag, ValueError)): return f"Decryption Failed. Check password/file. Details: {e}"
+        return "Data is binary and requires a password for decryption."
+    except Exception as e: 
+        logger.error(f"KB ImgUL: Decrypt error: {e}", exc_info=True)
+        return f"Unexpected decryption error: {e}"
+
+    if not kv_string: return "Could not get data from image (after potential decryption)."
+    try: 
+        kv_dict = parse_kv_string_to_dict(kv_string)
+    except Exception as e: 
+        logger.error(f"KB ImgUL: Parse fail: {e}")
+        return f"Error parsing data: {e}"
+    if not kv_dict: return "Parsed data is empty."
+
+    rules_to_add, memories_to_add = [], []
+    for key, value in kv_dict.items():
+        if key.startswith("rule_"):
+            try: rules_to_add.append(json.loads(value))
+            except: logger.warning(f"KB ImgUL: Bad rule format for key {key}")
+        elif key.startswith("memory_"):
+            try:
+                mem_dict = json.loads(value)
+                if isinstance(mem_dict, dict) and all(k in mem_dict for k in ['user_input', 'bot_response', 'metrics']):
+                    memories_to_add.append(mem_dict)
+            except: logger.warning(f"KB ImgUL: Bad memory format for key {key}")
+    
+    added_rules, skip_r, err_r, added_mems, err_m = 0, 0, 0, 0, 0
+    total = len(rules_to_add) + len(memories_to_add)
+    progress(0.5, desc=f"Adding {len(rules_to_add)} rules...")
+    for i, rule in enumerate(rules_to_add):
+        s, m = add_rule_entry(rule)
+        if s: added_rules += 1
+        elif m == "duplicate": skip_r += 1
+        else: err_r += 1
+        if total > 0: progress(0.5 + (0.4 * ((i+1)/total)) if total else 0)
+    
+    progress(0.9, desc=f"Adding {len(memories_to_add)} memories...")
+    for i, mem in enumerate(memories_to_add):
+        s, _ = add_memory_entry(mem['user_input'], mem['metrics'], mem['bot_response'])
+        if s: added_mems += 1
+        else: err_m += 1
+        if total > 0: progress(0.9 + (0.1 * ((i+1)/total)) if total else 0)
+        
+    progress(1.0, desc="Upload complete!")
+    msg = f"Upload Complete. Rules - Add: {added_rules}, Skip: {skip_r}, Err: {err_r}. Mems - Add: {added_mems}, Err: {err_m}."
+    logger.info(f"Image KB Upload: {msg}")
+    return msg
+
 def app_load_fn():
     logger.info("App loading. Initializing systems...")
     initialize_memory_system()
     logger.info("Memory system initialized.")
-
-    # --- Load Rules from File ---
     rules_added, rules_skipped, rules_errors = load_rules_from_file(LOAD_RULES_FILE)
     rules_load_msg = f"Rules: Added {rules_added}, Skipped {rules_skipped}, Errors {rules_errors} from {LOAD_RULES_FILE or 'None'}."
     logger.info(rules_load_msg)
-
-    # --- Load Memories from File ---
     mems_added, mems_format_errors, mems_save_errors = load_memories_from_file(LOAD_MEMORIES_FILE)
     mems_load_msg = f"Memories: Added {mems_added}, Format Errors {mems_format_errors}, Save Errors {mems_save_errors} from {LOAD_MEMORIES_FILE or 'None'}."
     logger.info(mems_load_msg)
-
     final_status = f"AI Systems Initialized. {rules_load_msg} {mems_load_msg} Ready."
-
-    # Initial population of all relevant UI components AFTER loading
-    rules_on_load = ui_refresh_rules_display_fn()
-    mems_on_load = ui_refresh_memories_display_fn()
-
-    # Return values for outputs defined in demo.load
-    return (
-        final_status, # agent_stat_tb
-        rules_on_load, # rules_disp_ta
-        mems_on_load,  # mems_disp_json
-        gr.Markdown(visible=False), # detect_out_md (initial state)
-        gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True), # fmt_report_tb (initial state)
-        gr.DownloadButton(interactive=False, value=None, visible=False), # dl_report_btn (initial state)
-    )
+    rules_on_load, mems_on_load = ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn()
+    return (final_status, rules_on_load, mems_on_load, gr.Markdown(visible=False),
+            gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True),
+            gr.DownloadButton(interactive=False, value=None, visible=False))
 
 
 # --- Gradio UI Definition ---
-with gr.Blocks(
-    theme=gr.themes.Soft(),
-    css="""
-        .gr-button { margin: 5px; }
-        .gr-textbox, .gr-text-area, .gr-dropdown, .gr-json { border-radius: 8px; }
-        .gr-group { border: 1px solid #e0e0e0; border-radius: 8px; padding: 10px; }
-        .gr-row { gap: 10px; }
-        .gr-tab { border-radius: 8px; }
-        .status-text { font-size: 0.9em; color: #555; }
-        .gr-json { max-height: 300px; overflow-y: auto; } /* Added scrolling for JSON */
+# --- Gradio UI Definition ---
+with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-textbox, .gr-text-area, .gr-dropdown, .gr-json { border-radius: 8px; } .gr-group { border: 1px solid #e0e0e0; border-radius: 8px; padding: 10px; } .gr-row { gap: 10px; } .gr-tab { border-radius: 8px; } .status-text { font-size: 0.9em; color: #555; } .gr-json { max-height: 400px; overflow-y: auto; }") as demo:
+    
+    # --- MODIFICATION: JavaScript for 1-Click Download ---
+    # This JS function will be triggered when the hidden file component changes.
+    # It receives the file data object from Gradio, which contains a web-accessible URL.
+    js_download_func = """
+    (file) => {
+        if (file && file.url) {
+            const link = document.createElement('a');
+            link.href = file.url;
+            link.download = `iLearn_KB_${new Date().toISOString().split('T')[0]}.png`;
+            document.body.appendChild(link);
+            link.click();
+            document.body.removeChild(link);
+        }
+    }
     """
-) as demo:
-    # --- DEMO MODE CHANGE ---
-    gr.Markdown(
-        f"""
-        # 🤖 AI Research Agent {'(DEMO MODE)' if DEMO_MODE else ''}
-        Your intelligent assistant for research and knowledge management
-        ### Special thanks to [Groq](https://groq.com) for their blazing fast inference
-        
-        """,
-        elem_classes=["header"]
-    )
-
-    is_sqlite = MEMORY_STORAGE_BACKEND == "SQLITE"
-    is_hf_dataset = MEMORY_STORAGE_BACKEND == "HF_DATASET"
-
+    
+    gr.Markdown(f"# 🤖 iLearn: An Autonomous Learning Agent {'(DEMO MODE)' if DEMO_MODE else ''}", elem_classes=["header"])
+    is_sqlite, is_hf_dataset = (MEMORY_STORAGE_BACKEND == "SQLITE"), (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
-        )
+        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"]
-            )
+            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"
-                )
-                available_providers = get_available_providers()
-                default_provider = available_providers[2] if available_providers else None
-                prov_sel_dd = gr.Dropdown(
-                    label="AI Provider", choices=available_providers,
-                    value=default_provider, interactive=True
-                )
+                api_key_tb = gr.Textbox(label="AI Provider API Key (Override)", type="password", placeholder="Uses .env if blank")
+                available_providers = get_available_providers(); default_provider = available_providers[0] if "groq" not in available_providers else "groq"
+                prov_sel_dd = gr.Dropdown(label="AI Provider", choices=available_providers, value=default_provider, interactive=True)
                 default_model_display = get_default_model_display_name_for_provider(default_provider) if default_provider else None
-                model_sel_dd = gr.Dropdown(
-                    label="AI Model",
-                    choices=get_model_display_names_for_provider(default_provider) if default_provider else [],
-                    value=default_model_display,
-                    interactive=True
-                )
+                model_sel_dd = gr.Dropdown(label="AI Model", choices=get_model_display_names_for_provider(default_provider) if default_provider else [], value=default_model_display, interactive=True)
             with gr.Group():
-                gr.Markdown("### System Prompt")
-                sys_prompt_tb = gr.Textbox(
-                    label="System Prompt Base", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True
-                )
-            if MEMORY_STORAGE_BACKEND == "RAM":
-                save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")
+                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://huggingface.co/spaces/Space-Share/bucket/resolve/main/images/pfp.webp"),
-                            show_copy_button=True, render_markdown=True, sanitize_html=True
-                        )
+                        main_chat_disp = gr.Chatbot(label=None, height=400, bubble_full_width=False,avatar_images=(None, "https://huggingface.co/spaces/Space-Share/bucket/resolve/main/images/pfp.webp"), 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
-                            )
+                            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
-                            )
+                            fmt_report_tb = gr.Textbox(label="Full AI Response", lines=8, interactive=True, show_copy_button=True)
+                            dl_report_btn = gr.DownloadButton("Download Report", value=None, interactive=False, visible=False)
                             detect_out_md = gr.Markdown(visible=False)
 
                 with gr.TabItem("🧠 Knowledge Base"):
                     with gr.Row(equal_height=True):
                         with gr.Column():
-                            gr.Markdown("### 📜 Rules Management")
-                            rules_disp_ta = gr.TextArea(
-                                label="Current Rules", lines=10,
-                                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.")
+                            gr.Markdown("### 📜 Rules Management (Text)"); rules_disp_ta = gr.TextArea(label="Current Rules", lines=10, placeholder="Rules will appear here.", interactive=True)
+                            gr.Markdown("To edit rules, modify text and click Save, or upload a file.")
                             save_edited_rules_btn = gr.Button("💾 Save Edited Text", variant="primary", interactive=not DEMO_MODE)
                             with gr.Row(variant="compact"):
-                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules", value=None)
-                                clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop", visible=not DEMO_MODE)
-                            # --- DEMO MODE CHANGE ---
-                            upload_rules_fobj = gr.File(
-                                label="Upload Rules File (.txt with '---' separators, or .jsonl of rule strings)",
-                                file_types=[".txt", ".jsonl"],
-                                interactive=not DEMO_MODE
-                            )
-                            rules_stat_tb = gr.Textbox(
-                                label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"]
-                            )
-
+                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules (.txt)", value=None); clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop", visible=not DEMO_MODE)
+                            upload_rules_fobj = gr.File(label="Upload Rules File (.txt/.jsonl)", file_types=[".txt", ".jsonl"], interactive=not DEMO_MODE)
+                            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", value=[]
-                            )
-                            gr.Markdown("To add memories, upload a .jsonl or .json file.")
+                            gr.Markdown("### 📚 Memories Management (Text)"); mems_disp_json = gr.JSON(label="Current Memories", value=[])
+                            gr.Markdown("To add memories, upload a 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", visible=not DEMO_MODE)
-                            # --- DEMO MODE CHANGE ---
-                            upload_mems_fobj = gr.File(
-                                label="Upload Memories File (.jsonl of memory objects, or .json array of objects)",
-                                file_types=[".jsonl", ".json"],
-                                interactive=not DEMO_MODE
-                            )
-                            mems_stat_tb = gr.Textbox(
-                                label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"]
-                            )
-
+                                dl_mems_btn = gr.DownloadButton("⬇️ Download Memories (.jsonl)", value=None); clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop", visible=not DEMO_MODE)
+                            upload_mems_fobj = gr.File(label="Upload Memories File (.json/.jsonl)", file_types=[".jsonl", ".json"], interactive=not DEMO_MODE)
+                            mems_stat_tb = gr.Textbox(label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"])
+                    with gr.Group():
+                        gr.Markdown("### 🖼️ Image-based Knowledge Base Import/Export")
+                        gr.Markdown("Save or load the entire knowledge base (rules and memories) as a single PNG image. A password can be used for AES-256 encryption.")
+                        with gr.Row():
+                            with gr.Column(scale=2):
+                                kb_img_password_tb = gr.Textbox(label="Password (optional for encryption)", type="password", placeholder="Leave blank for no encryption")
+                                kb_img_status_tb = gr.Textbox(label="Image Operation Status", interactive=False, lines=2)
+                            with gr.Column(scale=1):
+                                dl_kb_img_btn = gr.Button("⬇️ Download KB as Image", variant="secondary")
+                                upload_kb_img_fobj = gr.Image(label="⬆️ Upload or Paste KB Image", type="filepath", sources=["upload", "clipboard"], interactive=not DEMO_MODE)
+                                # --- MODIFICATION: This hidden component is the key to the solution ---
+                                hidden_downloader = gr.File(visible=False, label="File Downloader")
+
+    # --- Event Wiring ---
     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)
+        models_dyn = get_model_display_names_for_provider(sel_prov_dyn); def_model_dyn = get_default_model_display_name_for_provider(sel_prov_dyn)
         return gr.Dropdown(choices=models_dyn, value=def_model_dyn, interactive=True)
-
     prov_sel_dd.change(fn=dyn_upd_model_dd, inputs=prov_sel_dd, outputs=model_sel_dd)
 
-    # Inputs for the main chat submission function
     chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
-    # Outputs for the main chat submission function (includes knowledge base displays)
     chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn, rules_disp_ta, mems_disp_json]
-
     chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}
+    send_btn.click(**chat_event_args); user_msg_tb.submit(**chat_event_args)
 
-    send_btn.click(**chat_event_args)
-    user_msg_tb.submit(**chat_event_args)
-
-    # Rules Management events
     dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False)
+    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)
 
-    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, show_progress=False)
-
-    upload_mems_fobj.upload(
-        fn=ui_upload_memories_action_fn,
-        inputs=[upload_mems_fobj],
-        outputs=[mems_stat_tb],
+    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)
+
+    # --- MODIFICATION: Corrected Image KB event wiring ---
+    # 1. The button click now runs the Python function and outputs the file path to the hidden_downloader.
+    dl_kb_img_btn.click(
+        fn=ui_download_kb_as_image_fn, 
+        inputs=[kb_img_password_tb], 
+        outputs=[hidden_downloader], 
         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)
+    )
+    # 2. When the hidden_downloader's value changes, it triggers the JavaScript function for download.
+    hidden_downloader.change(
+        fn=None, # No python function needed here
+        inputs=[hidden_downloader], 
+        outputs=None, 
+        js=js_download_func
+    )
+    
+    upload_kb_img_fobj.upload(
+        fn=ui_upload_kb_from_image_fn, 
+        inputs=[upload_kb_img_fobj, kb_img_password_tb], 
+        outputs=[kb_img_status_tb], 
+        show_progress="full"
+    ).then(
+        fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta
+    ).then(
+        fn=ui_refresh_memories_display_fn, outputs=mems_disp_json
+    )
 
-    # FAISS save button visibility and action (RAM backend only)
     if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
         def save_faiss_action_with_feedback_sidebar_fn():
-            try:
-                save_faiss_indices_to_disk()
-                gr.Info("Attempted to save FAISS indices to disk.")
-            except Exception as e:
-                logger.error(f"Error saving FAISS indices: {e}", exc_info=True)
-                gr.Error(f"Error saving FAISS indices: {e}")
-
+            try: save_faiss_indices_to_disk(); gr.Info("Attempted to save FAISS indices to disk.")
+            except Exception as e: logger.error(f"Error saving FAISS indices: {e}", exc_info=True); gr.Error(f"Error saving FAISS indices: {e}")
         save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None, show_progress=False)
 
-    
-    # --- Initial Load Event ---
-    app_load_outputs = [
-        agent_stat_tb,
-        rules_disp_ta,
-        mems_disp_json,
-        detect_out_md,
-        fmt_report_tb,
-        dl_report_btn
-    ]
+    app_load_outputs = [agent_stat_tb, rules_disp_ta, mems_disp_json, detect_out_md, fmt_report_tb, dl_report_btn]
     demo.load(fn=app_load_fn, inputs=None, outputs=app_load_outputs, show_progress="full")
 
 
 if __name__ == "__main__":
-    logger.info(f"Starting Gradio AI Research Mega Agent (v6.5 - Direct UI Update & Core Learning Memories, Memory: {MEMORY_STORAGE_BACKEND})...")
+    logger.info(f"Starting Gradio AI Research Mega Agent (v9.1 - Correct 1-Click JS Download, 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, mcp_server=True)
+    demo.queue().launch(server_name=app_server, server_port=app_port, debug=app_debug, share=app_share)
     logger.info("Gradio application shut down.")