mirror of
https://github.com/open-webui/open-webui
synced 2025-06-26 18:26:48 +00:00
Add tags to S3 uploaded files
This commit is contained in:
parent
852d9dcbe9
commit
141810039f
@ -95,7 +95,13 @@ def upload_file(
|
|||||||
id = str(uuid.uuid4())
|
id = str(uuid.uuid4())
|
||||||
name = filename
|
name = filename
|
||||||
filename = f"{id}_{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(
|
file_item = Files.insert_new_file(
|
||||||
user.id,
|
user.id,
|
||||||
|
@ -3,7 +3,7 @@ import shutil
|
|||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from typing import BinaryIO, Tuple
|
from typing import BinaryIO, Tuple, Dict
|
||||||
|
|
||||||
import boto3
|
import boto3
|
||||||
from botocore.config import Config
|
from botocore.config import Config
|
||||||
@ -44,7 +44,7 @@ class StorageProvider(ABC):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@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
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
@ -58,7 +58,7 @@ class StorageProvider(ABC):
|
|||||||
|
|
||||||
class LocalStorageProvider(StorageProvider):
|
class LocalStorageProvider(StorageProvider):
|
||||||
@staticmethod
|
@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()
|
contents = file.read()
|
||||||
if not contents:
|
if not contents:
|
||||||
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
|
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
|
||||||
@ -131,12 +131,18 @@ class S3StorageProvider(StorageProvider):
|
|||||||
self.bucket_name = S3_BUCKET_NAME
|
self.bucket_name = S3_BUCKET_NAME
|
||||||
self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
|
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."""
|
"""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:
|
try:
|
||||||
s3_key = os.path.join(self.key_prefix, filename)
|
s3_key = os.path.join(self.key_prefix, filename)
|
||||||
self.s3_client.upload_file(file_path, self.bucket_name, s3_key)
|
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 (
|
return (
|
||||||
open(file_path, "rb").read(),
|
open(file_path, "rb").read(),
|
||||||
"s3://" + self.bucket_name + "/" + s3_key,
|
"s3://" + self.bucket_name + "/" + s3_key,
|
||||||
@ -207,9 +213,9 @@ class GCSStorageProvider(StorageProvider):
|
|||||||
self.gcs_client = storage.Client()
|
self.gcs_client = storage.Client()
|
||||||
self.bucket = self.gcs_client.bucket(GCS_BUCKET_NAME)
|
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."""
|
"""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:
|
try:
|
||||||
blob = self.bucket.blob(filename)
|
blob = self.bucket.blob(filename)
|
||||||
blob.upload_from_filename(file_path)
|
blob.upload_from_filename(file_path)
|
||||||
@ -277,9 +283,9 @@ class AzureStorageProvider(StorageProvider):
|
|||||||
self.container_name
|
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."""
|
"""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:
|
try:
|
||||||
blob_client = self.container_client.get_blob_client(filename)
|
blob_client = self.container_client.get_blob_client(filename)
|
||||||
blob_client.upload_blob(contents, overwrite=True)
|
blob_client.upload_blob(contents, overwrite=True)
|
||||||
|
Loading…
Reference in New Issue
Block a user