Update app.py

app.py

@@ -1115,12 +1115,13 @@ def ui_upload_kb_from_image_fn(uploaded_image_filepath: str, password: str, prog
     if not uploaded_image_filepath: 
         return "No image file provided or pasted."
         
-    progress(0, desc="Loading image...")
+    progress(0, desc="Loading and standardizing image...")
     try: 
-        img = Image.open(uploaded_image_filepath)
+        img_temp = Image.open(uploaded_image_filepath)
+        img = set_pil_image_format_to_png(img_temp)
     except Exception as e: 
-        logger.error(f"KB ImgUL: Open fail: {e}")
-        return f"Error: Could not open image file: {e}"
+        logger.error(f"KB ImgUL: Open/Standardize fail: {e}")
+        return f"Error: Could not open or process image file: {e}"
 
     progress(0.2, desc="Extracting data from image...")
     try:
@@ -1144,11 +1145,12 @@ def ui_upload_kb_from_image_fn(uploaded_image_filepath: str, password: str, prog
             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}"
+        if "decryption" in str(e).lower() or isinstance(e, InvalidTag):
+            return f"Decryption Failed. Check password or file integrity. 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}"
+        logger.error(f"KB ImgUL: Decrypt/Parse error: {e}", exc_info=True)
+        return f"Unexpected error during decryption or parsing: {e}"
 
     if not kv_string: return "Could not get data from image (after potential decryption)."
     try: 
@@ -1178,7 +1180,7 @@ def ui_upload_kb_from_image_fn(uploaded_image_filepath: str, password: str, prog
         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)
+        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):
@@ -1224,7 +1226,6 @@ try:
 except Exception as e:
     logger.error(f"Could not create placeholder image. The examples may not load correctly. Error: {e}")
 
-
 def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.Progress()):
     include_rules = "Include Rules" in content_to_include
     include_memories = "Include Memories" in content_to_include
@@ -1233,7 +1234,7 @@ def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.P
         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...")
+    progress(0.1, desc="Fetching knowledge base...")
     rules = get_all_rules_cached() if include_rules else []
     memories = get_all_memories_cached() if include_memories else []
 
@@ -1246,35 +1247,35 @@ def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.P
     data_bytes = kv_string.encode('utf-8')
 
     if password and password.strip():
-        progress(0.4, desc="Encrypting data...")
+        progress(0.3, 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...")
+    progress(0.5, 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...")
+    progress(0.6, 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 Knowledge Base" if password and password.strip() else "iLearn Knowledge Base"
-    final_image = draw_key_list_dropdown_overlay(embedded_image, keys=keys_for_overlay, title=title_overlay)
+    title_overlay = "Encrypted KB" if password and password.strip() else "iLearn KB"
+    image_with_overlay = draw_key_list_dropdown_overlay(carrier_image, keys=keys_for_overlay, title=title_overlay)
+
+    try:
+        progress(0.8, desc="Embedding data into final image...")
+        final_image_with_data = embed_data_in_image(image_with_overlay, 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.9, desc="Preparing final image and download file...")
     try:
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
-            final_image.save(tmpfile, format="PNG")
+            final_image_with_data.save(tmpfile, format="PNG")
             tmp_path = tmpfile.name
         progress(1.0, desc="Image created!")
         return gr.update(value=tmp_path, visible=True), gr.update(value=tmp_path, visible=True), "Success! Image created."
@@ -1282,6 +1283,8 @@ def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.P
         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()):
     if image_filepath:
         progress(0.1, desc="Image source detected. Starting image processing...")
@@ -1326,12 +1329,12 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
             with gr.Row():
                 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)
+                    main_chat_disp = gr.Chatbot(label=None, height=450, bubble_full_width=False,avatar_images=(None, "https://raw.githubusercontent.com/gradio-app/gradio/main/guides/assets/logo.png"), show_copy_button=True, render_markdown=True, sanitize_html=True)
                     with gr.Row(variant="compact"):
                         user_msg_tb = gr.Textbox(show_label=False, placeholder="Ask your research question...", scale=7, lines=1, max_lines=3)
                         send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
                     with gr.Accordion("📝 Detailed Response & Research Log", open=True):
-                        research_log_html = gr.HTML(label="Research Log", value="<p><i>Waiting for a new task to begin...</i></p>")
+                        research_log_html = gr.HTML(label="Research Log", value="<div class='log-container'><p><i>Waiting for a new task to begin...</i></p></div>")
                         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)
 
@@ -1368,7 +1371,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
                                 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")
-                            kb_image_display_output = gr.Image(label="Generated Image (Right-click to copy)", visible=False)
+                            kb_image_display_output = gr.Image(label="Generated Image (Right-click to copy)", type="filepath", visible=False)
                             kb_image_download_output = gr.DownloadButton("⬇️ Download Image File", visible=False)
 
                 with gr.TabItem("📂 Load KB"):
@@ -1384,7 +1387,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
                     load_master_btn = gr.Button("⬆️ Load from Sources", variant="primary", interactive=not DEMO_MODE)
                     gr.Examples(
                     examples=[
-                        ["https://huggingface.co/spaces/Agents-MCP-Hackathon/iLearn/resolve/main/evolutions/e0.01.01.png", ""],
+                        [placeholder_filename, ""],
                     ],
                     inputs=[upload_kb_img_fobj, load_kb_password_tb],
                     label="Click an Example to Load Data"
@@ -1433,7 +1436,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { margin: 5px; } .gr-text
     app_load_outputs = [agent_stat_tb, rules_disp_ta, mems_disp_json, research_log_html, 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))