diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 0ac53a023..58c3ef919 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]: @@ -951,5 +960,42 @@ 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..4d8175211 100644 --- a/backend/open_webui/models/files.py +++ b/backend/open_webui/models/files.py @@ -211,6 +211,21 @@ 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 ba6758671..65b58192b 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -62,19 +62,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 7b5659d51..98ddd41e7 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.routers.memories import query_memory, QueryMemoryForm from open_webui.utils.webhook import post_webhook @@ -2157,7 +2159,9 @@ 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: @@ -2206,73 +2210,49 @@ 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()) - - # 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( + 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] ) ) - - stdoutLines[idx] = ( - f"![Output Image {idx}](/cache/images/{id}.png)" - ) - - 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] - ) + 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, ) - resultLines[idx] = ( - f"![Output Image {idx}](/cache/images/{id}.png)" - ) + sourceLines[idx] = ( + f"![Output Image {idx}](/api/v1/files/{file_response.id}/content)" + ) + + output[sourceField] = "\n".join(sourceLines) - output["result"] = "\n".join(resultLines) except Exception as e: + log.exception(e) output = str(e) content_blocks[-1]["output"] = output