From aeb126b95da7e48e9d62dd58c33150394be3f6af Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 30 Aug 2026 23:46:02 -0400 Subject: [PATCH] refac --- backend/open_webui/routers/images.py | 39 +++++++++++++++----------- backend/open_webui/socket/main.py | 35 +++++++++++++++++++++-- backend/open_webui/tools/builtin.py | 14 +++++++-- backend/open_webui/utils/files.py | 3 +- backend/open_webui/utils/middleware.py | 20 ++++++------- 5 files changed, 78 insertions(+), 33 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index a15dd91498..51570000c5 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -552,7 +552,12 @@ async def upload_image(request, image_data, content_type, metadata, user, db=Non ) url = request.app.url_path_for('get_file_content_by_id', id=file_item.id) - return file_item, url + return file_item, { + 'id': file_item.id, + 'url': url, + 'name': (file_item.meta or {}).get('name') or file_item.filename, + 'content_type': (file_item.meta or {}).get('content_type'), + } @router.post('/generations') @@ -668,8 +673,8 @@ async def image_generations( else: image_data, content_type = await get_image_data(image['b64_json']) - _, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user) - images.append({'url': url}) + _, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user) + images.append(image_file) return images elif image_config.IMAGE_GENERATION_ENGINE == 'gemini': @@ -712,21 +717,21 @@ async def image_generations( if model.endswith(':predict'): for image in res['predictions']: image_data, content_type = await get_image_data(image['bytesBase64Encoded']) - _, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user) - images.append({'url': url}) + _, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user) + images.append(image_file) elif model.endswith(':generateContent'): for image in res['candidates']: for part in image['content']['parts']: if part.get('inlineData', {}).get('data'): image_data, content_type = await get_image_data(part['inlineData']['data']) - _, url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, {**data, **metadata}, user, ) - images.append({'url': url}) + images.append(image_file) return images @@ -776,14 +781,14 @@ async def image_generations( headers, trusted_base_url=image_config.COMFYUI_BASE_URL, ) - _, url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, {**form_data.model_dump(exclude_none=True), **metadata}, user, ) - images.append({'url': url}) + images.append(image_file) return images elif image_config.IMAGE_GENERATION_ENGINE == 'automatic1111' or image_config.IMAGE_GENERATION_ENGINE == '': # Automatic1111 holds one checkpoint instance-wide, so set_image_model @@ -823,14 +828,14 @@ async def image_generations( for image in res['images']: image_data, content_type = await get_image_data(image) - _, url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, {**data, 'info': res['info'], **metadata}, user, ) - images.append({'url': url}) + images.append(image_file) return images except Exception as e: error = e @@ -1039,8 +1044,8 @@ async def image_edits( else: image_data, content_type = await get_image_data(image['b64_json']) - _, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user) - images.append({'url': url}) + _, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user) + images.append(image_file) return images elif image_config.IMAGE_EDIT_ENGINE == 'gemini': @@ -1089,14 +1094,14 @@ async def image_edits( for part in image['content']['parts']: if part.get('inlineData', {}).get('data'): image_data, content_type = await get_image_data(part['inlineData']['data']) - _, url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, {**data, **metadata}, user, ) - images.append({'url': url}) + images.append(image_file) return images @@ -1173,14 +1178,14 @@ async def image_edits( headers, trusted_base_url=image_config.IMAGES_EDIT_COMFYUI_BASE_URL, ) - _, url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, {**form_data.model_dump(exclude_none=True), **metadata}, user, ) - images.append({'url': url}) + images.append(image_file) return images except Exception as e: diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 0998682dd8..d95158e0bd 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -938,20 +938,26 @@ async def _make_channel_emitter(request_info): state = {'last_emit_at': 0.0, 'output': []} THROTTLE_INTERVAL = 0.15 # ~6 updates/sec - async def _emit_channel_update(content: str, done: bool = False, output: list | None = None): + async def _emit_channel_update( + content: str, + done: bool = False, + output: list | None = None, + data: dict | None = None, + ): from open_webui.models.messages import MessageForm, Messages msg = await Messages.get_message_by_id(message_id) if not msg or msg.channel_id != channel_id: return - update_form = MessageForm(content=content, data={'output': output} if output else None) + update_data = data or ({'output': output} if output else None) + update_form = MessageForm(content=content, data=update_data) if done: # Merge done flag into existing meta (preserve model_id etc.) existing_meta = msg.meta or {} update_form = MessageForm( content=content, - data={'output': output} if output else None, + data=update_data, meta={**existing_meta, 'done': True}, ) @@ -1000,6 +1006,29 @@ async def _make_channel_emitter(request_info): state['last_emit_at'] = now await _emit_channel_update(content, False, state['output']) + elif event_type in ('files', 'chat:message:files'): + from open_webui.models.messages import Messages + + files = event_data.get('data', {}).get('files', []) + if not files: + return + + msg = await Messages.get_message_by_id(message_id) + if not msg or msg.channel_id != channel_id: + return + + existing_files = (msg.data or {}).get('files') + for file in files: + if isinstance(file, dict) and file.get('id'): + file['url'] = file['id'] + await Channels.add_file_to_channel_by_id(channel_id, file['id'], msg.user_id) + await Channels.set_file_message_id_in_channel_by_id(channel_id, file['id'], message_id) + + if isinstance(existing_files, list): + files.extend(existing_files) + + await _emit_channel_update(msg.content, data={'files': files}) + elif event_type == 'chat:message:error': error = event_data.get('data', {}).get('error', {}) error_content = error.get('content', 'An error occurred') if isinstance(error, dict) else str(error) diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py index e886dad315..853aaaa9d5 100644 --- a/backend/open_webui/tools/builtin.py +++ b/backend/open_webui/tools/builtin.py @@ -389,11 +389,16 @@ async def generate_image( images = await image_generations( request=__request__, form_data=CreateImageForm(prompt=prompt), + metadata=( + {'channel_id': __chat_id__.removeprefix('channel:'), 'message_id': __message_id__} + if isinstance(__chat_id__, str) and __chat_id__.startswith('channel:') + else None + ), user=user, ) # Prepare file entries for the images - image_files = [{'type': 'image', 'url': img['url']} for img in images] + image_files = [{'type': 'image', **img} for img in images] # Persist files to DB if chat context is available if is_saved_chat_id(__chat_id__) and __message_id__ and images: @@ -457,11 +462,16 @@ async def edit_image( images = await image_edits( request=__request__, form_data=EditImageForm(prompt=prompt, image=image_urls), + metadata=( + {'channel_id': __chat_id__.removeprefix('channel:'), 'message_id': __message_id__} + if isinstance(__chat_id__, str) and __chat_id__.startswith('channel:') + else None + ), user=user, ) # Prepare file entries for the images - image_files = [{'type': 'image', 'url': img['url']} for img in images] + image_files = [{'type': 'image', **img} for img in images] # Persist files to DB if chat context is available if is_saved_chat_id(__chat_id__) and __message_id__ and images: diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py index 74527a40b6..a247535032 100644 --- a/backend/open_webui/utils/files.py +++ b/backend/open_webui/utils/files.py @@ -86,13 +86,14 @@ async def get_image_url_from_base64(request, base64_image_string, metadata, user # Extract base64 image data from the line image_data, content_type = await get_image_data(base64_image_string) if image_data is not None: - _, image_url = await upload_image( + _, image_file = await upload_image( request, image_data, content_type, metadata, user, ) + image_url = image_file['url'] return image_url return None diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 6c18eb26ba..2f60d0aa9b 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -1771,6 +1771,12 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra if not chat_id or not isinstance(chat_id, str) or not __event_emitter__: return form_data + is_channel_chat = chat_id.startswith('channel:') + image_metadata = { + 'message_id': metadata.get('message_id', None), + **({'channel_id': chat_id.removeprefix('channel:')} if is_channel_chat else {'chat_id': chat_id}), + } + if not is_saved_chat_id(chat_id): message_list = form_data.get('messages', []) else: @@ -1815,10 +1821,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra 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), - }, + metadata=image_metadata, user=user, ) @@ -1836,7 +1839,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra 'files': [ { 'type': 'image', - 'url': image['url'], + **image, } for image in images ] @@ -1926,10 +1929,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra 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), - }, + metadata=image_metadata, user=user, ) @@ -1947,7 +1947,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra 'files': [ { 'type': 'image', - 'url': image['url'], + **image, } for image in images ]