From ffb9e739753f1d096531d5084fc524e642744266 Mon Sep 17 00:00:00 2001
From: Rodrigo Agundez <rragundez@gmail.com>
Date: Fri, 7 Feb 2025 08:32:06 +0900
Subject: [PATCH] Save image metadata to DB

---
 backend/open_webui/routers/files.py  | 36 +++++++++++++---------------
 backend/open_webui/routers/images.py |  9 ++-----
 2 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py
index 7160c2e86..7cf7a942f 100644
--- a/backend/open_webui/routers/files.py
+++ b/backend/open_webui/routers/files.py
@@ -3,30 +3,22 @@ import os
 import uuid
 from pathlib import Path
 from typing import Optional
-from pydantic import BaseModel
-import mimetypes
 from urllib.parse import quote
 
-from open_webui.storage.provider import Storage
-
+from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile, status
+from fastapi.responses import FileResponse, StreamingResponse
+from open_webui.constants import ERROR_MESSAGES
+from open_webui.env import SRC_LOG_LEVELS
 from open_webui.models.files import (
     FileForm,
     FileModel,
     FileModelResponse,
     Files,
 )
-from open_webui.routers.retrieval import process_file, ProcessFileForm
-
-from open_webui.config import UPLOAD_DIR
-from open_webui.env import SRC_LOG_LEVELS
-from open_webui.constants import ERROR_MESSAGES
-
-
-from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status, Request
-from fastapi.responses import FileResponse, StreamingResponse
-
-
+from open_webui.routers.retrieval import ProcessFileForm, process_file
+from open_webui.storage.provider import Storage
 from open_webui.utils.auth import get_admin_user, get_verified_user
+from pydantic import BaseModel
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["MODELS"])
@@ -41,7 +33,10 @@ router = APIRouter()
 
 @router.post("/", response_model=FileModelResponse)
 def upload_file(
-    request: Request, file: UploadFile = File(...), user=Depends(get_verified_user)
+    request: Request,
+    file: UploadFile = File(...),
+    user=Depends(get_verified_user),
+    file_metadata: dict = {},
 ):
     log.info(f"file.content_type: {file.content_type}")
     try:
@@ -61,6 +56,7 @@ def upload_file(
                     "id": id,
                     "filename": name,
                     "path": file_path,
+                    "data": file_metadata,
                     "meta": {
                         "name": name,
                         "content_type": file.content_type,
@@ -126,7 +122,7 @@ async def delete_all_files(user=Depends(get_admin_user)):
             Storage.delete_all_files()
         except Exception as e:
             log.exception(e)
-            log.error(f"Error deleting files")
+            log.error("Error deleting files")
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail=ERROR_MESSAGES.DEFAULT("Error deleting files"),
@@ -248,7 +244,7 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
                 )
         except Exception as e:
             log.exception(e)
-            log.error(f"Error getting file content")
+            log.error("Error getting file content")
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail=ERROR_MESSAGES.DEFAULT("Error getting file content"),
@@ -279,7 +275,7 @@ async def get_html_file_content_by_id(id: str, user=Depends(get_verified_user)):
                 )
         except Exception as e:
             log.exception(e)
-            log.error(f"Error getting file content")
+            log.error("Error getting file content")
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail=ERROR_MESSAGES.DEFAULT("Error getting file content"),
@@ -355,7 +351,7 @@ async def delete_file_by_id(id: str, user=Depends(get_verified_user)):
                 Storage.delete_file(file.path)
             except Exception as e:
                 log.exception(e)
-                log.error(f"Error deleting files")
+                log.error("Error deleting files")
                 raise HTTPException(
                     status_code=status.HTTP_400_BAD_REQUEST,
                     detail=ERROR_MESSAGES.DEFAULT("Error deleting files"),
diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py
index da4bf8a17..1dcaca866 100644
--- a/backend/open_webui/routers/images.py
+++ b/backend/open_webui/routers/images.py
@@ -411,7 +411,7 @@ def load_url_image_data(url, headers=None):
         return None
 
 
-def upload_image(request, data, image_data, content_type, user):
+def upload_image(request, image_metadata, image_data, content_type, user):
     image_format = mimetypes.guess_extension(content_type)
     file = UploadFile(
         file=io.BytesIO(image_data),
@@ -420,12 +420,7 @@ def upload_image(request, data, image_data, content_type, user):
             "content-type": content_type,
         },
     )
-    file_item = upload_file(request, file, user)
-    file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.filename}.json")
-
-    with open(file_body_path, "w") as f:
-        json.dump(data, f)
-
+    file_item = upload_file(request, file, user, file_metadata=image_metadata)
     url = request.app.url_path_for("get_file_content_by_id", id=file_item.id)
     return url