Update app.py

app.py

@@ -1,15 +1,7 @@
 import os
-# --- CONFIGURATION TOGGLES ---
-# Set these values to configure the application's behavior.
-
-# Set to True to disable destructive actions (clearing all data, saving edited rules, and all uploads).
 DEMO_MODE = False
-# Select the storage backend: "HF_DATASET", "SQLITE", or "RAM".
-# This will override the .env file setting for STORAGE_BACKEND.
 MEMORY_STORAGE_TYPE = "RAM" 
 
-# If using HF_DATASET, specify the repository names here.
-# These will override the .env file settings.
 HF_DATASET_MEMORY_REPO = "broadfield-dev/ai-brain"
 HF_DATASET_RULES_REPO = "broadfield-dev/ai-rules"
 
@@ -41,6 +33,7 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from cryptography.hazmat.primitives import hashes
 from cryptography.exceptions import InvalidTag
 # --- End New Imports ---
+import random
 
 load_dotenv() # Load .env file, but our settings above will take precedence if set.
 
@@ -90,7 +83,11 @@ 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 convert_pil_to_png_bytes(image: Image.Image) -> bytes:
+    """Saves a PIL image to an in-memory buffer as PNG and returns the raw bytes."""
+    with io.BytesIO() as buffer:
+        image.save(buffer, format="PNG")
+        return buffer.getvalue()
 def _get_font(preferred_fonts, base_size):
     fp = None
     safe_base_size = int(base_size)
@@ -167,74 +164,81 @@ def parse_kv_string_to_dict(kv_str:str)->dict:
         dd[k]=v
     return dd
 
-def generate_brain_carrier_image(w=800, h=800, msg="iLearn Knowledge Base") -> Image.Image:
+def generate_brain_carrier_image(w=800, h=800) -> Image.Image:
     """
-    Generates a carrier image with a radial gradient, a central text-based icon, and a message.
+    Generates a high-quality carrier image with a multi-layered, procedural starfield
+    using NumPy for performance and visual appeal.
     """
-    # --- 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... )"
+    # --- 1. Create the Gradient Background with NumPy ---
+    center_x, center_y = w / 2, h / 2
+    y_coords, x_coords = np.mgrid[0:h, 0:w]
 
-    # --- Setup ---
-    img = Image.new("RGBA", (w, h), bg_outer_color)
+    # Calculate distance of each pixel from the center
+    distance = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
+    max_distance = np.sqrt(center_x**2 + center_y**2)
+    
+    # Normalize distance to a 0-1 range
+    distance_norm = distance / max_distance
+    
+    # Define colors and create the gradient array
+    bg_center_color = np.array([20, 25, 40])  # Deeper blue center
+    bg_outer_color = np.array([0, 0, 0])      # Black edges
+    
+    # Interpolate colors across all pixels at once (vectorized)
+    # The [..., np.newaxis] part is for broadcasting the color channels
+    gradient = bg_outer_color + (bg_center_color - bg_outer_color) * (1 - distance_norm[..., np.newaxis])
+    
+    # Convert the NumPy array to a PIL Image to start drawing on it
+    img = Image.fromarray(gradient.astype(np.uint8), 'RGB')
     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 
+    # --- 2. Draw a Multi-Layered Starfield ---
     
-    # 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)
+    # Layer 1: Distant, tiny stars (for depth)
+    num_distant_stars = int((w * h) / 200)
+    for _ in range(num_distant_stars):
+        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
+        brightness = random.randint(30, 90)
+        draw.point((x, y), fill=(brightness, brightness, int(brightness * 1.1))) # Slightly blue tint
+
+    # Layer 2: Main stars with glow, size, and color variation
+    num_main_stars = int((w * h) / 1000)
+    star_colors = [
+        (255, 255, 255),  # White
+        (220, 230, 255),  # Light Blue
+        (255, 240, 220),  # Faint Yellow
+    ]
+    
+    for _ in range(num_main_stars):
+        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
+        dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
+        dist_ratio = min(dist_from_center / max_distance, 1.0)
+        
+        # Base size and brightness increase with distance from center
+        size = 0.5 + (2.5 * (dist_ratio ** 2))
+        brightness = 120 + (135 * (dist_ratio ** 1.5))
+        
+        # Select a random base color
+        color = random.choice(star_colors)
+        
+        # Apply brightness to the selected color
+        final_color = tuple(int(c * (brightness / 255.0)) for c in color)
+        
+        # Simulate a soft glow by drawing a larger, dimmer circle first
+        glow_size = size * 3
+        glow_color = tuple(int(c * 0.3) for c in final_color) # Much dimmer
+        draw.ellipse([x - glow_size, y - glow_size, x + glow_size, y + glow_size], fill=glow_color)
+
+        # Simulate a "twinkle" effect for some stars by drawing a cross
+        if random.random() < 0.15: # 15% chance to twinkle
+            draw.line([x-size, y, x+size, y], fill=final_color, width=1)
+            draw.line([x, y-size, x, y+size], fill=final_color, width=1)
+        else:
+            # Draw the main star on top of the glow
+            draw.ellipse([x - size, y - size, x + size, y + size], fill=final_color)
+            
+    return img
 
-    # --- 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'):
@@ -256,63 +260,44 @@ def _get_text_measurement(draw_obj, text_str, font_obj):
         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,[]
+    """Draws overlays on the image using the 'KeyLock' style."""
+    img_overlayed = image.copy().convert("RGBA")
+    draw = ImageDraw.Draw(img_overlayed, "RGBA")
+    width, height = img_overlayed.size
+
+    overlay_color = (15, 23, 42, 190)
+    title_color = (226, 232, 240)
+    key_color = (148, 163, 184)
+
+    font_bold = _get_font(PREFERRED_FONTS, 30)
+    font_regular = _get_font(PREFERRED_FONTS, 15)
+
+    draw.rectangle([0, 20, width, 80], fill=overlay_color)
+    draw.text((width / 2, 50), title, fill=title_color, font=font_bold, anchor="ms")
+
     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 ---
+        box_padding = 15
+        line_spacing = 6
+        text_start_x = 35
+        lines = keys
+
+        line_heights = [_get_text_measurement(draw, line, font_regular)[1] for line in lines]
+        total_text_height = sum(line_heights) + (len(lines) - 1) * line_spacing
+        box_height = total_text_height + (box_padding * 2)
+        box_y0 = height - box_height - 20
+
+        draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
+        current_y = box_y0 + box_padding
+        
+        for i, key_text in enumerate(lines):
+            draw.text((text_start_x, current_y), key_text, fill=key_color, font=font_regular)
+            if i < len(line_heights):
+                current_y += line_heights[i] + line_spacing
+
+    final_image_rgb = Image.new("RGB", img_overlayed.size, (0, 0, 0))
+    final_image_rgb.paste(img_overlayed, (0, 0), img_overlayed)
+    
+    return final_image_rgb
 
 
 # --- Helper Functions ---
@@ -383,7 +368,7 @@ def process_user_interaction_gradio(user_input: str, provider_name: str, model_d
         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_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)
@@ -595,18 +580,16 @@ Combine all findings into a single, valid XML structure as specified in the syst
                 else:
                     logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}' from XML.")
 
-            # After processing all rules, if there were significant learnings, add a special memory
             if significant_learnings_summary:
                 learning_digest = "SYSTEM CORE LEARNING DIGEST:\n" + "\n".join(significant_learnings_summary)
-                # Create a synthetic metrics object for this system memory
                 system_metrics = {
                     "takeaway": "Core knowledge refined.",
-                    "response_success_score": 1.0, # Assuming successful internal update
+                    "response_success_score": 1.0,
                     "future_confidence_score": 1.0,
                     "type": "SYSTEM_REFLECTION"
                 }
                 add_memory_entry(
-                    user_input="SYSTEM_INTERNAL_REFLECTION_TRIGGER", # Fixed identifier
+                    user_input="SYSTEM_INTERNAL_REFLECTION_TRIGGER",
                     metrics=system_metrics,
                     bot_response=learning_digest
                 )
@@ -623,9 +606,8 @@ Combine all findings into a single, valid XML structure as specified in the syst
 def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_name: str, sel_model_disp_name: str, ui_api_key: str|None, cust_sys_prompt: str):
     global current_chat_session_history
     cleared_input, updated_gr_hist, status_txt = "", list(gr_hist_list), "Initializing..."
-    # Initialize all potential output components with default/current values
-    updated_rules_text = ui_refresh_rules_display_fn() # Get current rules state
-    updated_mems_json = ui_refresh_memories_display_fn() # Get current memories state
+    updated_rules_text = ui_refresh_rules_display_fn()
+    updated_mems_json = ui_refresh_memories_display_fn()
     def_detect_out_md = gr.Markdown(visible=False)
     def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True)
     def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
@@ -633,21 +615,18 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
     if not user_msg_txt.strip():
         status_txt = "Error: Empty message."
         updated_gr_hist.append((user_msg_txt or "(Empty)", status_txt))
-        # Ensure all outputs are provided on early exit
         yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
         return
 
     updated_gr_hist.append((user_msg_txt, "<i>Thinking...</i>"))
-    # Initial yield to update chat UI with thinking message and show current knowledge base state
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
 
     internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
-    # History truncation logic (keep MAX_HISTORY_TURNS pairs + optional system prompt)
     hist_len_check = MAX_HISTORY_TURNS * 2
     if internal_hist and internal_hist[0]["role"] == "system": hist_len_check +=1
     if len(internal_hist) > hist_len_check:
          current_chat_session_history = ([internal_hist[0]] if internal_hist[0]["role"] == "system" else []) + internal_hist[-(MAX_HISTORY_TURNS * 2):]
-         internal_hist = list(current_chat_session_history) # Use truncated history for current turn processing
+         internal_hist = list(current_chat_session_history)
 
     final_bot_resp_acc, insights_used_parsed = "", []
     temp_dl_file_path = None
@@ -659,21 +638,18 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
             if upd_type == "status":
                 status_txt = upd_data
                 if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
-                    # Update the status alongside the streaming message
                     updated_gr_hist[-1] = (user_msg_txt, f"{curr_bot_disp_msg} <i>{status_txt}</i>" if curr_bot_disp_msg else f"<i>{status_txt}</i>")
             elif upd_type == "response_chunk":
                 curr_bot_disp_msg += upd_data
                 if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
-                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg) # Update chat with streamed chunk
+                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg)
             elif upd_type == "final_response_and_insights":
                 final_bot_resp_acc, insights_used_parsed = upd_data["response"], upd_data["insights_used"]
                 status_txt = "Response generated. Processing learning..."
-                # Ensure the final chat message reflects the full response
                 if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
                 if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                     updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg or "(No text)")
 
-                # Update detailed response box and download button
                 def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg, interactive=True, show_copy_button=True)
 
                 if curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
@@ -688,15 +664,11 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
                 else:
                     def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
 
-                # Update insights display
                 insights_md_content = "### Insights Considered (Pre-Response):\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
                 def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True if insights_used_parsed else False)
 
-            # Yield intermediate updates for the UI
-            # Pass the *current* state of rules and memories display components
             yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
 
-            # Stop processing generator after final_response_and_insights
             if upd_type == "final_response_and_insights": break
 
     except Exception as e:
@@ -710,29 +682,23 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
         def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
         def_detect_out_md = gr.Markdown(value="*Error processing request.*", visible=True)
 
-        # Provide the current state of rules/memories on error path yield
         current_rules_text_on_error = ui_refresh_rules_display_fn()
         current_mems_json_on_error = ui_refresh_memories_display_fn()
 
         yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_on_error, current_mems_json_on_error)
-        # Clean up temp file if created before error
         if temp_dl_file_path and os.path.exists(temp_dl_file_path):
             try: os.unlink(temp_dl_file_path)
             except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path} after error: {e_unlink}")
-        return # Exit the function after error handling
+        return
 
-    # --- Post-Interaction Learning ---
     if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
-        # Add the successful turn to the internal history
         current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
-        # History truncation again after adding
         hist_len_check = MAX_HISTORY_TURNS * 2
         if current_chat_session_history and current_chat_session_history[0]["role"] == "system": hist_len_check +=1
         if len(current_chat_session_history) > hist_len_check:
             current_chat_session_history = ([current_chat_session_history[0]] if current_chat_session_history[0]["role"] == "system" else []) + current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]
 
         status_txt = "<i>[Performing post-interaction learning...]</i>"
-        # Yield status before synchronous learning
         current_rules_text_before_learn = ui_refresh_rules_display_fn()
         current_mems_json_before_learn = ui_refresh_memories_display_fn()
         yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_before_learn, current_mems_json_before_learn)
@@ -753,21 +719,14 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
 
     elif final_bot_resp_acc.startswith("Error:"):
         status_txt = final_bot_resp_acc
-        # If it was an error response from the generator, it's already in updated_gr_hist[-1]
-        # The other output components (fmt_report_tb, dl_btn, detect_out_md) are already set by the generator loop or default state
     else:
         status_txt = "Processing finished; no valid response or error occurred during main phase."
 
-
-    # Final yield after learning (or error handling)
-    # This final yield updates the UI one last time with the true final status
-    # AND crucially refreshes the Rules and Memories displays in case they changed during learning.
     updated_rules_text = ui_refresh_rules_display_fn()
     updated_mems_json = ui_refresh_memories_display_fn()
 
     yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
 
-    # Clean up the temporary download file after the final yield
     if temp_dl_file_path and os.path.exists(temp_dl_file_path):
         try: os.unlink(temp_dl_file_path)
         except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path}: {e_unlink}")
@@ -775,10 +734,9 @@ def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_na
 
 # --- Startup Loading Functions ---
 def load_rules_from_file(filepath: str | None):
-    """Loads rules from a local file (.txt or .jsonl) and adds them to the system."""
     if not filepath:
         logger.info("LOAD_RULES_FILE environment variable not set. Skipping rules loading from file.")
-        return 0, 0, 0 # added, skipped, errors
+        return 0, 0, 0
 
     if not os.path.exists(filepath):
         logger.warning(f"LOAD_RULES: Specified rules file not found: {filepath}. Skipping loading.")
@@ -792,7 +750,7 @@ def load_rules_from_file(filepath: str | None):
             content = f.read()
     except Exception as e:
         logger.error(f"LOAD_RULES: Error reading file {filepath}: {e}", exc_info=False)
-        return 0, 0, 1 # Indicate read error
+        return 0, 0, 1
 
     if not content.strip():
         logger.info(f"LOAD_RULES: File {filepath} is empty. Skipping loading.")
@@ -802,7 +760,6 @@ def load_rules_from_file(filepath: str | None):
 
     if file_name_lower.endswith(".txt"):
         potential_rules = content.split("\n\n---\n\n")
-        # Also handle simple line breaks if '---' separator is not used
         if len(potential_rules) == 1 and "\n" in content:
              potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
     elif file_name_lower.endswith(".jsonl"):
@@ -810,7 +767,6 @@ def load_rules_from_file(filepath: str | None):
             line = line.strip()
             if line:
                 try:
-                    # Expecting each line to be a JSON string containing the rule text
                     rule_text_in_json_string = json.loads(line)
                     if isinstance(rule_text_in_json_string, str):
                         potential_rules.append(rule_text_in_json_string)
@@ -822,7 +778,7 @@ def load_rules_from_file(filepath: str | None):
                     error_count +=1
     else:
         logger.error(f"LOAD_RULES: Unsupported file type for rules: {filepath}. Must be .txt or .jsonl")
-        return 0, 0, 1 # Indicate type error
+        return 0, 0, 1
 
     valid_potential_rules = [r.strip() for r in potential_rules if r.strip()]
     total_to_process = len(valid_potential_rules)
@@ -832,7 +788,7 @@ def load_rules_from_file(filepath: str | None):
         return 0, 0, 0
     elif total_to_process == 0 and error_count > 0:
          logger.warning(f"LOAD_RULES: No valid rule segments found to process. Encountered {error_count} parsing/format errors in {filepath}.")
-         return 0, 0, error_count # Indicate only errors
+         return 0, 0, error_count
 
     logger.info(f"LOAD_RULES: Attempting to add {total_to_process} potential rules from {filepath}...")
     for idx, rule_text in enumerate(valid_potential_rules):
@@ -849,10 +805,9 @@ def load_rules_from_file(filepath: str | None):
     return added_count, skipped_count, error_count
 
 def load_memories_from_file(filepath: str | None):
-    """Loads memories from a local file (.json or .jsonl) and adds them to the system."""
     if not filepath:
         logger.info("LOAD_MEMORIES_FILE environment variable not set. Skipping memories loading from file.")
-        return 0, 0, 0 # added, format_errors, save_errors
+        return 0, 0, 0
 
     if not os.path.exists(filepath):
         logger.warning(f"LOAD_MEMORIES: Specified memories file not found: {filepath}. Skipping loading.")
@@ -866,7 +821,7 @@ def load_memories_from_file(filepath: str | None):
             content = f.read()
     except Exception as e:
         logger.error(f"LOAD_MEMORIES: Error reading file {filepath}: {e}", exc_info=False)
-        return 0, 1, 0 # Indicate read error
+        return 0, 1, 0
 
     if not content.strip():
         logger.info(f"LOAD_MEMORIES: File {filepath} is empty. Skipping loading.")
@@ -880,7 +835,6 @@ def load_memories_from_file(filepath: str | None):
             if isinstance(parsed_json, list):
                 memory_objects_to_process = parsed_json
             elif isinstance(parsed_json, dict):
-                 # If it's a single object, process it as a list of one
                  memory_objects_to_process = [parsed_json]
             else:
                 logger.warning(f"LOAD_MEMORIES (.json): File content is not a JSON list or object in {filepath}. Type: {type(parsed_json)}")
@@ -899,7 +853,7 @@ def load_memories_from_file(filepath: str | None):
                     format_error_count += 1
     else:
         logger.error(f"LOAD_MEMORIES: Unsupported file type for memories: {filepath}. Must be .json or .jsonl")
-        return 0, 1, 0 # Indicate type error
+        return 0, 1, 0
 
     total_to_process = len(memory_objects_to_process)
 
@@ -913,18 +867,11 @@ def load_memories_from_file(filepath: str | None):
 
     logger.info(f"LOAD_MEMORIES: Attempting to add {total_to_process} memory objects from {filepath}...")
     for idx, mem_data in enumerate(memory_objects_to_process):
-        # Validate minimum structure
         if isinstance(mem_data, dict) and all(k in mem_data for k in ["user_input", "bot_response", "metrics"]):
-            # Add entry without generating new embeddings if possible (assuming file contains embeddings)
-            # NOTE: The current add_memory_entry function *always* generates embeddings.
-            # If performance is an issue with large files, memory_logic might need
-            # an optimized bulk import function that reuses existing embeddings or
-            # generates them in batches. For now, we use the existing function.
-            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"]) # add_memory_entry needs user_input, metrics, bot_response
+            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"])
             if success:
                 added_count += 1
             else:
-                # add_memory_entry currently doesn't return detailed error status
                 logger.warning(f"LOAD_MEMORIES: Failed to save memory object from {filepath} (segment {idx+1}). Data: {str(mem_data)[:100]}")
                 save_error_count += 1
         else:
@@ -936,18 +883,23 @@ def load_memories_from_file(filepath: str | None):
 
 
 # --- 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."""
+def convert_kb_to_kv_string(rules: list[str], memories: list[dict], include_rules: bool, include_memories: bool) -> str:
+    """Serializes rules and/or memories into a single key-value string for image embedding based on user selection."""
     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)}")
+    
+    if include_rules:
+        lines.append("\n# --- RULES ---")
+        for i, rule_text in enumerate(rules):
+            lines.append(f"rule_{i+1} = {json.dumps(rule_text)}")
+
+    if include_memories:
+        lines.append("\n# --- MEMORIES ---")
+        for i, mem_dict in enumerate(memories):
+            lines.append(f"memory_{i+1} = {json.dumps(mem_dict)}")
 
-    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 []
 
@@ -1093,66 +1045,8 @@ def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
         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."
@@ -1253,26 +1147,174 @@ def app_load_fn():
             gr.DownloadButton(interactive=False, value=None, visible=False))
 
 
-# --- Gradio UI Definition ---
-# --- 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:
+# --- Create a placeholder image for the examples ---
+# This makes the script self-contained and runnable without needing a separate file.
+placeholder_filename = "placeholder_image.png"
+try:
+    if not os.path.exists(placeholder_filename):
+        img = Image.new('RGB', (200, 100), color='darkblue')
+        draw = Image.Draw(img)
+        try:
+            font = _get_font(PREFERRED_FONTS, 14)
+            draw.text((10, 45), "Placeholder KB Image", font=font, fill='white')
+        except Exception:
+            draw.text((10, 45), "Placeholder", fill='white')
+        img.save(placeholder_filename)
+        logger.info(f"Created '{placeholder_filename}' for Gradio examples.")
+except Exception as e:
+    logger.error(f"Could not create placeholder image. The examples may not load correctly. Error: {e}")
+
+
+def ui_download_kb_as_image_fn(password: str, progress=gr.Progress()):
+    """
+    Generates a KB image and returns both the image object for display
+    and a file path for a download button.
+    """
+    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 create.")
+        # Return updates to hide the components if they were previously visible
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "Knowledge base is empty."
+
+    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 gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+    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 gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+    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="Preparing final image and download file...")
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
+            final_image.save(tmpfile, format="PNG")
+            tmp_path = tmpfile.name
+
+        progress(1.0, desc="Image created!")
+        gr.Info("Image created and is ready for download or copy.")
+        # Return updates to show the components with the new data
+        return gr.update(value=final_image, visible=True), gr.update(value=tmp_path, visible=True), "Success! Image created."
+    except Exception as e:
+        logger.error(f"KB ImgDL: Save failed: {e}")
+        gr.Error(f"Failed to save final image: {e}")
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+
+
+
+def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.Progress()):
+    """
+    Generates a KB image and returns a file path to both the display and download components.
+    """
+    include_rules = "Include Rules" in content_to_include
+    include_memories = "Include Memories" in content_to_include
     
-    # --- 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);
-        }
-    }
+    if not include_rules and not include_memories:
+        gr.Warning("Nothing selected to save.")
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "Nothing selected to save."
+
+    progress(0, desc="Fetching knowledge base...")
+    rules = get_all_rules_cached() if include_rules else []
+    memories = get_all_memories_cached() if include_memories else []
+
+    if not rules and not memories:
+        gr.Warning("Knowledge base is empty or selected content is empty.")
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "No content to save."
+
+    progress(0.2, desc="Serializing data...")
+    kv_string = convert_kb_to_kv_string(rules, memories, include_rules, include_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())
+        except Exception as e:
+            logger.error(f"KB ImgDL: Encrypt failed: {e}")
+            return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+    progress(0.6, desc="Generating carrier image...")
+    carrier_image = generate_brain_carrier_image(w=800, h=800)
+
+    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}")
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+    progress(0.8, desc="Adding visual overlay...")
+    keys_for_overlay = []
+    if include_rules: keys_for_overlay.append(f"Rule Count: {len(rules)}")
+    if include_memories: keys_for_overlay.append(f"Memory Count: {len(memories)}")
+    
+    title_overlay = "Encrypted KB" if password and password.strip() else "iLearn KB"
+    final_image = draw_key_list_dropdown_overlay(embedded_image, keys=keys_for_overlay, title=title_overlay)
+
+    progress(0.9, desc="Preparing final image and download file...")
+    try:
+        # Create a temporary file and save the image as a PNG.
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
+            final_image.save(tmpfile, format="PNG")
+            tmp_path = tmpfile.name
+        progress(1.0, desc="Image created!")
+        # Return the FILE PATH to both components.
+        return gr.update(value=tmp_path, visible=True), gr.update(value=tmp_path, visible=True), "Success! Image created."
+    except Exception as e:
+        logger.error(f"KB ImgDL: Save failed: {e}")
+        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"
+
+def ui_load_from_sources_fn(image_filepath: str, rules_file_obj: object, mems_file_obj: object, password: str, progress=gr.Progress()):
+    """
+    Loads data from one of the available sources with precedence: Image > Rules File > Memories File.
     """
+    if image_filepath:
+        progress(0.1, desc="Image source detected. Starting image processing...")
+        return ui_upload_kb_from_image_fn(image_filepath, password, progress)
+    
+    if rules_file_obj:
+        progress(0.1, desc="Rules file detected. Starting rules import...")
+        return ui_upload_rules_action_fn(rules_file_obj, progress)
+        
+    if mems_file_obj:
+        progress(0.1, desc="Memories file detected. Starting memories import...")
+        return ui_upload_memories_action_fn(mems_file_obj, progress)
+        
+    return "No file or image uploaded. Please provide a source file to load."
+
     
+# --- 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:
+
     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"):
@@ -1282,63 +1324,86 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-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[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)
-            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.Tabs():
+        with gr.TabItem("💬 Chat & Research"):
+            with gr.Row():
+                with gr.Sidebar():
+                    gr.Markdown("## ⚙️ Configuration")
                     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)
-                        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):
+                        gr.Markdown("### AI Model Settings")
+                        api_key_tb = gr.Textbox(label="AI Provider API Key (Override)", type="password", placeholder="Uses .env if blank")
+                        available_providers = get_available_providers(); default_provider = available_providers[0] if "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)
+                    with gr.Group():
+                        gr.Markdown("### System Prompt"); sys_prompt_tb = gr.Textbox(label="System Prompt Base", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
+
+                with gr.Column(scale=3):
+                    gr.Markdown("### AI Chat Interface")
+                    main_chat_disp = gr.Chatbot(label=None, height=450, 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)
+                        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():
+                rules_stat_tb = gr.Textbox(label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"])
+                mems_stat_tb = gr.Textbox(label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"])
+
+            with gr.Tabs():
+                with gr.TabItem("🎛️ System"):
+                    gr.Markdown("View and directly manage the current rules and memories in the system.")
+                    with gr.Row(equal_height=False, variant='compact'):
                         with gr.Column():
-                            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 (.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"])
+                            gr.Markdown("### 📜 Current Rules")
+                            rules_disp_ta = gr.TextArea(label=None, lines=15, placeholder="Rules will appear here.", interactive=True)
+                            save_edited_rules_btn = gr.Button("💾 Save Edited Rules", variant="primary", interactive=not DEMO_MODE)
+                            clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop", visible=not DEMO_MODE)
                         with gr.Column():
-                            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 (.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"])
+                            gr.Markdown("### 📚 Current Memories")
+                            mems_disp_json = gr.JSON(label=None, value=[], scale=1)
+                            clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop", visible=not DEMO_MODE)
+
+                with gr.TabItem("💾 Save KB"):
+                    gr.Markdown("Export the current knowledge base as text files or as a single, portable PNG image.")
+                    with gr.Row():
+                        with gr.Column():
+                            gr.Markdown("### Text File Export")
+                            dl_rules_btn = gr.DownloadButton("⬇️ Download Rules (.txt)", value=None)
+                            dl_mems_btn = gr.DownloadButton("⬇️ Download Memories (.jsonl)", value=None)
+                            gr.Row()
+                            if MEMORY_STORAGE_BACKEND == "RAM": save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")
+
+
+                        with gr.Column():
+                            gr.Markdown("### Image Export")
+                            with gr.Group():
+                                save_kb_password_tb = gr.Textbox(label="Password (optional for encryption)", type="password")
+                                save_kb_include_cbg = gr.CheckboxGroup(label="Content to Include", choices=["Include Rules", "Include Memories"], value=["Include Rules", "Include Memories"])
+                                create_kb_img_btn = gr.Button("✨ Create KB Image", variant="secondary")
+                            # Removed type="pil" to allow Gradio to handle the raw PNG bytes correctly
+                            kb_image_display_output = gr.Image(label="Generated Image (Right-click to copy)", visible=False)
+                            kb_image_download_output = gr.DownloadButton("⬇️ Download Image File", visible=False)
+
+                with gr.TabItem("📂 Load KB"):
+                    gr.Markdown("Import rules, memories, or a full KB from local files or a portable PNG image.")
+                    load_status_tb = gr.Textbox(label="Load Operation Status", interactive=False, lines=2)
+                    load_kb_password_tb = gr.Textbox(label="Password (for decrypting images)", type="password")
+                    
                     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.")
+                        gr.Markdown("#### Sources (Priority: Image > Rules File > Memories File)")
                         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")
+                             upload_kb_img_fobj = gr.Image(label="1. Image Source", type="filepath", sources=["upload", "clipboard"], interactive=not DEMO_MODE)
+                             upload_rules_fobj = gr.File(label="2. Rules File Source (.txt/.jsonl)", file_types=[".txt", ".jsonl"], interactive=not DEMO_MODE)
+                             upload_mems_fobj = gr.File(label="3. Memories File Source (.json/.jsonl)", file_types=[".jsonl", ".json"], interactive=not DEMO_MODE)
+                    
+                    load_master_btn = gr.Button("⬆️ Load from Sources", variant="primary", interactive=not DEMO_MODE)
+                    
 
     # --- Event Wiring ---
     def dyn_upd_model_dd(sel_prov_dyn: str):
@@ -1346,40 +1411,32 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
         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 Tab
     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, 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)
 
-    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False)
+    # KB Tab -> System
     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)
-
-    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], 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], 
+    # KB Tab -> Save KB
+    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False)
+    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn, show_progress=False)
+    create_kb_img_btn.click(
+        fn=ui_create_kb_image_fn,
+        inputs=[save_kb_password_tb, save_kb_include_cbg],
+        outputs=[kb_image_display_output, kb_image_download_output, load_status_tb],
         show_progress="full"
     )
-    # 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], 
+
+    # KB Tab -> Load KB
+    load_master_btn.click(
+        fn=ui_load_from_sources_fn,
+        inputs=[upload_kb_img_fobj, upload_rules_fobj, upload_mems_fobj, load_kb_password_tb],
+        outputs=[load_status_tb],
         show_progress="full"
     ).then(
         fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta
@@ -1387,16 +1444,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
         fn=ui_refresh_memories_display_fn, outputs=mems_disp_json
     )
 
+    # Sidebar FAISS button
     if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
         def save_faiss_action_with_feedback_sidebar_fn():
             try: save_faiss_indices_to_disk(); gr.Info("Attempted to save FAISS indices to disk.")
             except Exception as e: logger.error(f"Error saving FAISS indices: {e}", exc_info=True); gr.Error(f"Error saving FAISS indices: {e}")
         save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None, show_progress=False)
 
+    # App Load
     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 (v9.1 - Correct 1-Click JS Download, Memory: {MEMORY_STORAGE_BACKEND})...")
     app_port = int(os.getenv("GRADIO_PORT", 7860))