From 607d01b107c4bee9dc2c2406c7826aa753566d1e Mon Sep 17 00:00:00 2001 From: Jarrod Lowe Date: Mon, 12 May 2025 17:08:28 +1200 Subject: [PATCH] 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)"