From d372bec70427fe2d568e052ce5e1529e2ad41da9 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 13 Sep 2026 21:37:12 -0400 Subject: [PATCH] refac Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com> --- backend/open_webui/utils/files.py | 4 ++- backend/open_webui/utils/middleware.py | 45 +++++++++++++++++++++++--- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py index 676af24afe..f982da1675 100644 --- a/backend/open_webui/utils/files.py +++ b/backend/open_webui/utils/files.py @@ -32,6 +32,7 @@ from open_webui.storage.provider import Storage BASE64_IMAGE_URL_PREFIX = re.compile(r'data:image/\w+;base64,', re.IGNORECASE) MARKDOWN_IMAGE_URL_PATTERN = re.compile(r'!\[(.*?)\]\((.+?)\)', re.IGNORECASE) +FILE_CONTENT_URL_PATTERN = re.compile(r'^/api/v1/files/([^/?#]+)/content(?:[?#]|$)') # Extension-based MIME fallback, only used when ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK is True. _IMAGE_MIME_FALLBACK = { @@ -90,7 +91,8 @@ async def get_image_base64_from_url(url: str, user=None) -> Optional[str]: else: # Non-URL string — treat as file_id. Delegate to the canonical # file-ID resolver which enforces ownership/access checks. - return await get_image_base64_from_file_id(url, user=user) + file_id_match = FILE_CONTENT_URL_PATTERN.match(url) + return await get_image_base64_from_file_id(file_id_match.group(1) if file_id_match else url, user=user) except Exception: return None diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 4523e842ad..a3b3632f4c 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -995,6 +995,29 @@ def extract_base64_images(value: Any, files: list) -> Any: return value +async def store_tool_result_image(request, image_url, metadata, user): + """Keep saved tool images out of chat JSON, falling back to inline data if storage fails.""" + metadata = metadata or {} + if ( + not isinstance(image_url, str) + or not image_url.startswith('data:image/') + or not is_saved_chat_id(metadata.get('chat_id')) + ): + return image_url + + try: + stored_url = await get_file_url_from_base64( + request, + image_url, + {key: metadata.get(key) for key in ('chat_id', 'message_id', 'session_id')}, + user, + ) + return stored_url or image_url + except Exception: + log.warning('Could not store tool image; retaining inline image') + return image_url + + async def process_tool_result( request, tool_function_name, @@ -1425,6 +1448,12 @@ async def chat_completion_tools_handler( ) if tool_result_files: + for file_item in tool_result_files: + if file_item.get('type') == 'image': + file_item['url'] = await store_tool_result_image( + request, file_item.get('url'), metadata, user + ) + await event_emitter( { 'type': 'files', @@ -2113,7 +2142,7 @@ async def convert_url_images_to_base64(form_data, user=None): new_content = [] for item in content: - if not isinstance(item, dict) or item.get('type') != 'image_url': + if not isinstance(item, dict) or item.get('type') not in ('image_url', 'input_image'): new_content.append(item) continue @@ -2130,13 +2159,15 @@ async def convert_url_images_to_base64(form_data, user=None): try: base64_data = await get_image_base64_from_url(image_url, user=user) - if base64_data: + if base64_data and isinstance(image_url_data, str): + new_content.append({**item, 'image_url': base64_data}) + elif base64_data: image_url_payload = {'url': base64_data} if isinstance(image_url_data, dict) and image_url_data.get('detail'): image_url_payload['detail'] = image_url_data['detail'] new_content.append( { - 'type': 'image_url', + 'type': item['type'], 'image_url': image_url_payload, } ) @@ -3321,7 +3352,8 @@ async def drain_approved_tool_calls(request, form_data, user, model, metadata) - display_files = [] for file_item in result.get('files', []): if file_item.get('type') == 'image' and file_item.get('url', '').startswith('data:'): - output_parts.append({'type': 'input_image', 'image_url': file_item['url']}) + image_url = await store_tool_result_image(request, file_item['url'], metadata, user) + output_parts.append({'type': 'input_image', 'image_url': image_url}) else: display_files.append(file_item) @@ -5907,7 +5939,8 @@ async def streaming_chat_response_handler(response, ctx): for file_item in result.get('files', []): if file_item.get('type') == 'image' and file_item.get('url', '').startswith('data:'): # LLM-only: add as input_image part, not frontend display output. - output_parts.append({'type': 'input_image', 'image_url': file_item['url']}) + image_url = await store_tool_result_image(request, file_item['url'], metadata, user) + output_parts.append({'type': 'input_image', 'image_url': image_url}) else: # Frontend display (MCP images, audio, etc.) display_files.append(file_item) @@ -6071,6 +6104,8 @@ async def streaming_chat_response_handler(response, ctx): } ) + new_form_data = await convert_url_images_to_base64(new_form_data, user=user) + if filter_functions: new_form_data, _ = await process_filter_functions( request=request,