From 141810039f17f1c0a92a36adfc382b5091d17b5d Mon Sep 17 00:00:00 2001
From: Jarrod Lowe <jarrod@lowe.kiwi>
Date: Mon, 28 Apr 2025 16:56:56 +1200
Subject: [PATCH] Add tags to S3 uploaded files

---
 backend/open_webui/routers/files.py    |  8 +++++++-
 backend/open_webui/storage/provider.py | 24 +++++++++++++++---------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py
index 8a2888d86..3189dd070 100644
--- a/backend/open_webui/routers/files.py
+++ b/backend/open_webui/routers/files.py
@@ -95,7 +95,13 @@ def upload_file(
         id = str(uuid.uuid4())
         name = filename
         filename = f"{id}_{filename}"
-        contents, file_path = Storage.upload_file(file.file, filename)
+        tags = {
+            "OpenWebUI-User-Email": user.email,
+            "OpenWebUI-User-Id": user.id,
+            "OpenWebUI-User-Name": user.name,
+            "OpenWebUI-File-Id": id,
+        }
+        contents, file_path = Storage.upload_file(file.file, filename, tags)
 
         file_item = Files.insert_new_file(
             user.id,
diff --git a/backend/open_webui/storage/provider.py b/backend/open_webui/storage/provider.py
index c5c0056cc..5855b0185 100644
--- a/backend/open_webui/storage/provider.py
+++ b/backend/open_webui/storage/provider.py
@@ -3,7 +3,7 @@ import shutil
 import json
 import logging
 from abc import ABC, abstractmethod
-from typing import BinaryIO, Tuple
+from typing import BinaryIO, Tuple, Dict
 
 import boto3
 from botocore.config import Config
@@ -44,7 +44,7 @@ class StorageProvider(ABC):
         pass
 
     @abstractmethod
-    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+    def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
         pass
 
     @abstractmethod
@@ -58,7 +58,7 @@ class StorageProvider(ABC):
 
 class LocalStorageProvider(StorageProvider):
     @staticmethod
-    def upload_file(file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+    def upload_file(file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
         contents = file.read()
         if not contents:
             raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
@@ -131,12 +131,18 @@ class S3StorageProvider(StorageProvider):
         self.bucket_name = S3_BUCKET_NAME
         self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
 
-    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+    def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
         """Handles uploading of the file to S3 storage."""
-        _, file_path = LocalStorageProvider.upload_file(file, filename)
+        _, file_path = LocalStorageProvider.upload_file(file, filename, tags)
+        tagging = {'TagSet': [{'Key': k, 'Value': v} for k, v in tags.items()]}
         try:
             s3_key = os.path.join(self.key_prefix, filename)
             self.s3_client.upload_file(file_path, self.bucket_name, s3_key)
+            self.s3_client.put_object_tagging(
+                Bucket=self.bucket_name,
+                Key=s3_key,
+                Tagging=tagging,
+            )
             return (
                 open(file_path, "rb").read(),
                 "s3://" + self.bucket_name + "/" + s3_key,
@@ -207,9 +213,9 @@ class GCSStorageProvider(StorageProvider):
             self.gcs_client = storage.Client()
         self.bucket = self.gcs_client.bucket(GCS_BUCKET_NAME)
 
-    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+    def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
         """Handles uploading of the file to GCS storage."""
-        contents, file_path = LocalStorageProvider.upload_file(file, filename)
+        contents, file_path = LocalStorageProvider.upload_file(file, filename, tags)
         try:
             blob = self.bucket.blob(filename)
             blob.upload_from_filename(file_path)
@@ -277,9 +283,9 @@ class AzureStorageProvider(StorageProvider):
             self.container_name
         )
 
-    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+    def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
         """Handles uploading of the file to Azure Blob Storage."""
-        contents, file_path = LocalStorageProvider.upload_file(file, filename)
+        contents, file_path = LocalStorageProvider.upload_file(file, filename, tags)
         try:
             blob_client = self.container_client.get_blob_client(filename)
             blob_client.upload_blob(contents, overwrite=True)