From b0401e1f36a465dd133dc57701c04e7e8eda98dc Mon Sep 17 00:00:00 2001 From: Jarrod Lowe Date: Mon, 12 May 2025 16:52:33 +1200 Subject: [PATCH 1/2] Store pyodide output image files in the storage system --- backend/open_webui/models/chats.py | 49 +++++++++++++++ backend/open_webui/models/files.py | 10 +++ backend/open_webui/routers/files.py | 10 +-- backend/open_webui/utils/middleware.py | 86 ++++++++------------------ 4 files changed, 89 insertions(+), 66 deletions(-) diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 4b4f37197..48c3c38c7 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -6,6 +6,7 @@ from typing import Optional from open_webui.internal.db import Base, get_db from open_webui.models.tags import TagModel, Tag, Tags +from open_webui.models.files import Files from open_webui.env import SRC_LOG_LEVELS from pydantic import BaseModel, ConfigDict @@ -296,6 +297,14 @@ class ChatTable: .update({"share_id": shared_chat.id}) ) db.commit() + + # Make sure all the output files are shared. There names are GUIDs + # and they don't show up in search for everyone, so these can still + # only be accessed if you know their GUID and are a valid user. + for fileId in chat.meta.get("outputFileIds", []): + log.debug(f"Setting shared on file {fileId}") + Files.update_file_access_control_by_id(fileId, {"shared": True}) + return shared_chat if (shared_result and result) else None def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]: @@ -908,5 +917,45 @@ class ChatTable: except Exception: return False + def add_output_file_id_to_chat(self, id: str, file_id: str) -> Optional[ChatModel]: + """Adds a new file ID to the outputFileIds list in the chat's metadata.""" + try: + with get_db() as db: + chat = db.query(Chat).filter(Chat.id == id).with_for_update().first() + if chat is None: + return None + + output_file_ids = chat.meta.get("outputFileIds", []) + if file_id in output_file_ids: + return ChatModel.model_validate(chat) + + output_file_ids.append(file_id) + chat.meta = { + **chat.meta, + "outputFileIds": output_file_ids + } + chat.updated_at = int(time.time()) + + db.commit() + return ChatModel.model_validate(chat) + except Exception as e: + log.error(f"Error adding output file ID: {e}") + return None + + def get_output_file_ids_by_chat_id(self, id: str) -> list[str]: + """ + Gets all file IDs from the outputFileIds list in the chat's metadata. + """ + try: + with get_db() as db: + chat = db.query(Chat).filter(Chat.id == id).first() + if chat is None: + return [] + + return chat.meta.get("outputFileIds", []) + except Exception as e: + log.error(f"Error getting output file IDs: {e}") + return [] + Chats = ChatTable() diff --git a/backend/open_webui/models/files.py b/backend/open_webui/models/files.py index 6f1511cd1..e6d7bbc5f 100644 --- a/backend/open_webui/models/files.py +++ b/backend/open_webui/models/files.py @@ -211,6 +211,16 @@ class FilesTable: except Exception: return None + def update_file_access_control_by_id(self, id: str, access_control: dict) -> Optional[FileModel]: + with get_db() as db: + try: + file = db.query(File).filter_by(id=id).first() + file.access_control = {**(file.access_control if file.access_control else {}), **access_control} + db.commit() + return FileModel.model_validate(file) + except Exception as e: + return None + def delete_file_by_id(self, id: str) -> bool: with get_db() as db: try: diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index 475905da1..cca3af3e5 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -60,19 +60,19 @@ def has_access_to_file( detail=ERROR_MESSAGES.NOT_FOUND, ) - has_access = False - knowledge_base_id = file.meta.get("collection_name") if file.meta else None + if file.access_control.get("shared", False): + return True + knowledge_base_id = file.meta.get("collection_name") if file.meta else None if knowledge_base_id: knowledge_bases = Knowledges.get_knowledge_bases_by_user_id( user.id, access_type ) for knowledge_base in knowledge_bases: if knowledge_base.id == knowledge_base_id: - has_access = True - break + return True - return has_access + return False ############################ diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 442dfba76..c3420fc97 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -3,6 +3,7 @@ import logging import sys import os import base64 +import io import asyncio from aiocache import cached @@ -18,7 +19,7 @@ from uuid import uuid4 from concurrent.futures import ThreadPoolExecutor -from fastapi import Request, HTTPException +from fastapi import Request, HTTPException, UploadFile from starlette.responses import Response, StreamingResponse @@ -41,6 +42,7 @@ from open_webui.routers.pipelines import ( process_pipeline_inlet_filter, process_pipeline_outlet_filter, ) +from open_webui.routers.files import upload_file from open_webui.utils.webhook import post_webhook @@ -2039,7 +2041,7 @@ async def process_chat_response( ) retries += 1 - log.debug(f"Attempt count: {retries}") + log.debug(f"Attempt count: {retries}, intepreter {request.app.state.config.CODE_INTERPRETER_ENGINE}") output = "" try: @@ -2088,73 +2090,35 @@ async def process_chat_response( "stdout": "Code interpreter engine not configured." } - log.debug(f"Code interpreter output: {output}") if isinstance(output, dict): - stdout = output.get("stdout", "") + for sourceField in ("stdout", "result"): + source = output.get(sourceField, "") - if isinstance(stdout, str): - stdoutLines = stdout.split("\n") - for idx, line in enumerate(stdoutLines): - if "data:image/png;base64" in line: - id = str(uuid4()) + if isinstance(source, str): + sourceLines = source.split("\n") + for idx, line in enumerate(sourceLines): + if "data:image/png;base64" in line: + # line looks like data:image/png;base64, + content_type = line.split(',')[0].split(';')[0].split(':')[1] + file_data = io.BytesIO(base64.b64decode(line.split(',')[1])) + file_name = f"image-{metadata['chat_id']}-{metadata['message_id']}-{sourceField}-{idx}.png" + file = UploadFile( + filename=file_name, + file=file_data, + headers={"content-type": content_type}, + ) + file_response = upload_file(request, file, user=user) + Chats.add_output_file_id_to_chat(metadata["chat_id"], file_response.id) - # ensure the path exists - os.makedirs( - os.path.join(CACHE_DIR, "images"), - exist_ok=True, - ) - - image_path = os.path.join( - CACHE_DIR, - f"images/{id}.png", - ) - - with open(image_path, "wb") as f: - f.write( - base64.b64decode( - line.split(",")[1] - ) + sourceLines[idx] = ( + f"![Output Image {idx}](/api/v1/files/{file_response.id}/content)" ) - stdoutLines[idx] = ( - f"![Output Image {idx}](/cache/images/{id}.png)" - ) + output[sourceField] = "\n".join(sourceLines) - output["stdout"] = "\n".join(stdoutLines) - - result = output.get("result", "") - - if isinstance(result, str): - resultLines = result.split("\n") - for idx, line in enumerate(resultLines): - if "data:image/png;base64" in line: - id = str(uuid4()) - - # ensure the path exists - os.makedirs( - os.path.join(CACHE_DIR, "images"), - exist_ok=True, - ) - - image_path = os.path.join( - CACHE_DIR, - f"images/{id}.png", - ) - - with open(image_path, "wb") as f: - f.write( - base64.b64decode( - line.split(",")[1] - ) - ) - - resultLines[idx] = ( - f"![Output Image {idx}](/cache/images/{id}.png)" - ) - - output["result"] = "\n".join(resultLines) except Exception as e: + log.exception(e) output = str(e) content_blocks[-1]["output"] = output From 607d01b107c4bee9dc2c2406c7826aa753566d1e Mon Sep 17 00:00:00 2001 From: Jarrod Lowe Date: Mon, 12 May 2025 17:08:28 +1200 Subject: [PATCH 2/2] Reformatting --- backend/open_webui/models/chats.py | 5 +---- backend/open_webui/models/files.py | 9 ++++++-- backend/open_webui/utils/middleware.py | 30 ++++++++++++++++++++------ 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 48c3c38c7..f95a1aa30 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -930,10 +930,7 @@ class ChatTable: return ChatModel.model_validate(chat) output_file_ids.append(file_id) - chat.meta = { - **chat.meta, - "outputFileIds": output_file_ids - } + chat.meta = {**chat.meta, "outputFileIds": output_file_ids} chat.updated_at = int(time.time()) db.commit() diff --git a/backend/open_webui/models/files.py b/backend/open_webui/models/files.py index e6d7bbc5f..4d8175211 100644 --- a/backend/open_webui/models/files.py +++ b/backend/open_webui/models/files.py @@ -211,11 +211,16 @@ class FilesTable: except Exception: return None - def update_file_access_control_by_id(self, id: str, access_control: dict) -> Optional[FileModel]: + def update_file_access_control_by_id( + self, id: str, access_control: dict + ) -> Optional[FileModel]: with get_db() as db: try: file = db.query(File).filter_by(id=id).first() - file.access_control = {**(file.access_control if file.access_control else {}), **access_control} + file.access_control = { + **(file.access_control if file.access_control else {}), + **access_control, + } db.commit() return FileModel.model_validate(file) except Exception as e: diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index c3420fc97..4200552c3 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -2041,7 +2041,9 @@ async def process_chat_response( ) retries += 1 - log.debug(f"Attempt count: {retries}, intepreter {request.app.state.config.CODE_INTERPRETER_ENGINE}") + log.debug( + f"Attempt count: {retries}, intepreter {request.app.state.config.CODE_INTERPRETER_ENGINE}" + ) output = "" try: @@ -2090,7 +2092,6 @@ async def process_chat_response( "stdout": "Code interpreter engine not configured." } - if isinstance(output, dict): for sourceField in ("stdout", "result"): source = output.get(sourceField, "") @@ -2100,16 +2101,31 @@ async def process_chat_response( for idx, line in enumerate(sourceLines): if "data:image/png;base64" in line: # line looks like data:image/png;base64, - content_type = line.split(',')[0].split(';')[0].split(':')[1] - file_data = io.BytesIO(base64.b64decode(line.split(',')[1])) + content_type = ( + line.split(",")[0] + .split(";")[0] + .split(":")[1] + ) + file_data = io.BytesIO( + base64.b64decode( + line.split(",")[1] + ) + ) file_name = f"image-{metadata['chat_id']}-{metadata['message_id']}-{sourceField}-{idx}.png" file = UploadFile( filename=file_name, file=file_data, - headers={"content-type": content_type}, + headers={ + "content-type": content_type + }, + ) + file_response = upload_file( + request, file, user=user + ) + Chats.add_output_file_id_to_chat( + metadata["chat_id"], + file_response.id, ) - file_response = upload_file(request, file, user=user) - Chats.add_output_file_id_to_chat(metadata["chat_id"], file_response.id) sourceLines[idx] = ( f"![Output Image {idx}](/api/v1/files/{file_response.id}/content)"