smart_image_gen v0.7.2: chat-DB fallback + diagnostic 'no image' msg

If __messages__ doesn't include the assistant's prior file attachments
(which is what the screenshot is showing), the new fallback queries
the chat by id via Chats.get_chat_by_id and walks every persisted
message for files. Open WebUI's socket handler always upserts files
onto the assistant message via {'files': files} so this path is
authoritative.

The 'No image found' return now includes diagnostic counts —
__files__, __messages__, messages_with_files, chat_id_present,
openwebui_runtime — so subsequent failures actually show what the
tool saw instead of being opaque.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

deployments/ai-stack/openwebui-tools/smart_image_gen.py

@@ -1,7 +1,7 @@
 """
 title: Smart Image Generator & Editor (ComfyUI)
 author: ai-stack
-version: 0.7.1
+version: 0.7.2
 description: Generate or edit images via ComfyUI with automatic SDXL
     checkpoint routing. Two methods — generate_image (txt2img) and
     edit_image (img2img on the user's most recently attached image). The
@@ -34,6 +34,7 @@ from pydantic import BaseModel, Field
 # falls back to emitting a markdown data-URI message.
 try:
     from fastapi import UploadFile
+    from open_webui.models.chats import Chats
     from open_webui.models.files import Files
     from open_webui.models.users import Users
     from open_webui.routers.files import upload_file_handler
@@ -405,6 +406,7 @@ def _read_file_dict(f: dict) -> Optional[bytes]:
 async def _extract_attached_image(
     files: Optional[list],
     messages: Optional[list],
+    metadata: Optional[dict],
     session: aiohttp.ClientSession,
 ) -> Optional[bytes]:
     """
@@ -460,7 +462,31 @@ async def _extract_attached_image(
         if data is not None:
             return data
 
-    # 4. Last-resort URL fetch (no auth — only works for public endpoints).
+    # 4. Pull the chat from the database directly. Open WebUI persists
+    # `files` on every message via the upsert in socket/main.py — so even
+    # if __messages__ doesn't hydrate the assistant-emitted attachments,
+    # the chat record does. This is the strongest fallback.
+    if _OPENWEBUI_RUNTIME and metadata:
+        chat_id = metadata.get("chat_id")
+        if chat_id:
+            try:
+                chat = Chats.get_chat_by_id(chat_id)
+                chat_data = getattr(chat, "chat", None) if chat else None
+                chat_messages = (chat_data or {}).get("messages", []) if isinstance(chat_data, dict) else []
+                for msg in reversed(chat_messages):
+                    if not isinstance(msg, dict):
+                        continue
+                    msg_files = msg.get("files") or []
+                    for f in msg_files:
+                        if not isinstance(f, dict) or not _file_dict_is_image(f):
+                            continue
+                        data = _read_file_dict(f)
+                        if data is not None:
+                            return data
+            except Exception:
+                pass
+
+    # 5. Last-resort URL fetch (no auth — only works for public endpoints).
     for source in [files or []] + [
         (msg.get("files") or []) for msg in reversed(messages or []) if isinstance(msg, dict)
     ]:
@@ -806,12 +832,24 @@ class Tools:
 
         async with aiohttp.ClientSession() as session:
             await emit("Looking for attached image…")
-            raw_in = await _extract_attached_image(__files__, __messages__, session)
+            raw_in = await _extract_attached_image(
+                __files__, __messages__, __metadata__, session,
+            )
             if raw_in is None:
+                msgs_with_files = sum(
+                    1 for m in (__messages__ or [])
+                    if isinstance(m, dict) and m.get("files")
+                )
+                chat_id_present = bool((__metadata__ or {}).get("chat_id"))
                 return (
-                    "No image found in the chat. Ask the user to attach the "
-                    "image they want edited (paperclip / drag-drop), or call "
-                    "generate_image instead if they want a new image."
+                    "No image found in the chat. Diagnostics: "
+                    f"__files__={len(__files__ or [])}, "
+                    f"__messages__={len(__messages__ or [])} "
+                    f"(of which {msgs_with_files} had a files field), "
+                    f"chat_id_present={chat_id_present}, "
+                    f"openwebui_runtime={_OPENWEBUI_RUNTIME}. "
+                    "Ask the user to attach the image they want edited "
+                    "(paperclip / drag-drop), or call generate_image instead."
                 )
 
             await emit("Uploading source to ComfyUI…")