diff --git a/deployments/ai-stack/openwebui-tools/smart_image_gen.py b/deployments/ai-stack/openwebui-tools/smart_image_gen.py index 7de17a2..70cc74f 100644 --- a/deployments/ai-stack/openwebui-tools/smart_image_gen.py +++ b/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…")