From 446cc0ac6063402a743e949f50612376ed5a8437 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 22 Dec 2025 00:39:05 +0400 Subject: [PATCH] refac --- backend/open_webui/routers/images.py | 65 ++++++++++++++++++++------ backend/open_webui/utils/files.py | 5 +- backend/open_webui/utils/middleware.py | 22 ++++++++- 3 files changed, 75 insertions(+), 17 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index ed4348414..91bc3f2d4 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -17,6 +17,8 @@ from fastapi.responses import FileResponse from open_webui.config import CACHE_DIR from open_webui.constants import ERROR_MESSAGES from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS + +from open_webui.models.chats import Chats from open_webui.routers.files import upload_file_handler, get_file_content_by_id from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.headers import include_user_info_headers @@ -510,14 +512,29 @@ def upload_image(request, image_data, content_type, metadata, user): process=False, user=user, ) + + if file_item and file_item.id: + # If chat_id and message_id are provided in metadata, link the file to the chat message + chat_id = metadata.get("chat_id") + message_id = metadata.get("message_id") + + if chat_id and message_id: + Chats.insert_chat_files( + chat_id=chat_id, + message_id=message_id, + file_ids=[file_item.id], + user_id=user.id, + ) + url = request.app.url_path_for("get_file_content_by_id", id=file_item.id) - return url + return file_item, url @router.post("/generations") async def image_generations( request: Request, form_data: CreateImageForm, + metadata: Optional[dict] = None, user=Depends(get_verified_user), ): # if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default @@ -535,6 +552,9 @@ async def image_generations( size = form_data.size width, height = tuple(map(int, size.split("x"))) + + metadata = metadata or {} + model = get_image_model(request) r = None @@ -595,7 +615,9 @@ async def image_generations( else: image_data, content_type = get_image_data(image["b64_json"]) - url = upload_image(request, image_data, content_type, data, user) + _, url = upload_image( + request, image_data, content_type, {**data, **metadata}, user + ) images.append({"url": url}) return images @@ -645,7 +667,9 @@ async def image_generations( image_data, content_type = get_image_data( image["bytesBase64Encoded"] ) - url = upload_image(request, image_data, content_type, data, user) + _, url = upload_image( + request, image_data, content_type, {**data, **metadata}, user + ) images.append({"url": url}) elif model.endswith(":generateContent"): for image in res["candidates"]: @@ -654,8 +678,12 @@ async def image_generations( image_data, content_type = get_image_data( part["inlineData"]["data"] ) - url = upload_image( - request, image_data, content_type, data, user + _, url = upload_image( + request, + image_data, + content_type, + {**data, **metadata}, + user, ) images.append({"url": url}) @@ -705,11 +733,11 @@ async def image_generations( } image_data, content_type = get_image_data(image["url"], headers) - url = upload_image( + _, url = upload_image( request, image_data, content_type, - form_data.model_dump(exclude_none=True), + {**form_data.model_dump(exclude_none=True), **metadata}, user, ) images.append({"url": url}) @@ -752,11 +780,11 @@ async def image_generations( for image in res["images"]: image_data, content_type = get_image_data(image) - url = upload_image( + _, url = upload_image( request, image_data, content_type, - {**data, "info": res["info"]}, + {**data, "info": res["info"], **metadata}, user, ) images.append({"url": url}) @@ -783,10 +811,13 @@ class EditImageForm(BaseModel): async def image_edits( request: Request, form_data: EditImageForm, + metadata: Optional[dict] = None, user=Depends(get_verified_user), ): size = None width, height = None, None + metadata = metadata or {} + if ( request.app.state.config.IMAGE_EDIT_SIZE and "x" in request.app.state.config.IMAGE_EDIT_SIZE @@ -904,7 +935,9 @@ async def image_edits( else: image_data, content_type = get_image_data(image["b64_json"]) - url = upload_image(request, image_data, content_type, data, user) + _, url = upload_image( + request, image_data, content_type, {**data, **metadata}, user + ) images.append({"url": url}) return images @@ -957,8 +990,12 @@ async def image_edits( image_data, content_type = get_image_data( part["inlineData"]["data"] ) - url = upload_image( - request, image_data, content_type, data, user + _, url = upload_image( + request, + image_data, + content_type, + {**data, **metadata}, + user, ) images.append({"url": url}) @@ -1035,11 +1072,11 @@ async def image_edits( } image_data, content_type = get_image_data(image_url, headers) - url = upload_image( + _, url = upload_image( request, image_data, content_type, - form_data.model_dump(exclude_none=True), + {**form_data.model_dump(exclude_none=True), **metadata}, user, ) images.append({"url": url}) diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py index 2221d1707..a37ecf31c 100644 --- a/backend/open_webui/utils/files.py +++ b/backend/open_webui/utils/files.py @@ -14,6 +14,8 @@ from typing import Optional from pathlib import Path from open_webui.storage.provider import Storage + +from open_webui.models.chats import Chats from open_webui.models.files import Files from open_webui.routers.files import upload_file_handler @@ -65,13 +67,14 @@ def get_image_url_from_base64(request, base64_image_string, metadata, user): # Extract base64 image data from the line image_data, content_type = get_image_data(base64_image_string) if image_data is not None: - image_url = upload_image( + _, image_url = upload_image( request, image_data, content_type, metadata, user, ) + return image_url return None diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 865159208..05f76b172 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -799,6 +799,10 @@ async def chat_image_generation_handler( images = await image_edits( request=request, form_data=EditImageForm(**{"prompt": prompt, "image": input_images}), + metadata={ + "chat_id": metadata.get("chat_id", None), + "message_id": metadata.get("message_id", None), + }, user=user, ) @@ -824,7 +828,7 @@ async def chat_image_generation_handler( } ) - system_message_content = "The requested image has been created and is now being shown to the user. Let them know that it has been generated." + system_message_content = "The requested image has been edited and created and is now being shown to the user. Let them know that it has been generated." except Exception as e: log.debug(e) @@ -883,6 +887,10 @@ async def chat_image_generation_handler( images = await image_generations( request=request, form_data=CreateImageForm(**{"prompt": prompt}), + metadata={ + "chat_id": metadata.get("chat_id", None), + "message_id": metadata.get("message_id", None), + }, user=user, ) @@ -2737,7 +2745,17 @@ async def process_chat_response( if ENABLE_CHAT_RESPONSE_BASE64_IMAGE_URL_CONVERSION: value = convert_markdown_base64_images( - request, value, metadata, user + request, + value, + { + "chat_id": metadata.get( + "chat_id", None + ), + "message_id": metadata.get( + "message_id", None + ), + }, + user, ) content = f"{content}{value}"