diff --git a/llmstack/assets/__init__.py b/llmstack/assets/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/llmstack/assets/admin.py b/llmstack/assets/admin.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/llmstack/assets/apps.py b/llmstack/assets/apps.py new file mode 100644 index 00000000000..aacc78df7f4 --- /dev/null +++ b/llmstack/assets/apps.py @@ -0,0 +1,6 @@ +from django.apps import AppConfig + + +class AssetsConfig(AppConfig): + name = "assets" + label = "assets" diff --git a/llmstack/assets/models.py b/llmstack/assets/models.py new file mode 100644 index 00000000000..47ea1c70221 --- /dev/null +++ b/llmstack/assets/models.py @@ -0,0 +1,49 @@ +import base64 +import uuid + +from django.db import models +from django.db.models.signals import pre_delete +from django.dispatch import receiver + + +class Assets(models.Model): + uuid = models.UUIDField(default=uuid.uuid4, editable=False, help_text="UUID of the asset") + path = None + file = None + metadata = models.JSONField( + default=dict, + help_text="Metadata for the asset", + null=True, + blank=True, + ) + created_at = models.DateTimeField(auto_now_add=True) + + @classmethod + def create_from_bytes(cls, file_bytes, filename, metadata=None, path=""): + from django.core.files.base import ContentFile + + asset = cls(path=path) + asset.file.save(filename, ContentFile(file_bytes)) + bytes_size = len(file_bytes) + asset.metadata = {**metadata, "file_size": bytes_size} + asset.save() + return asset + + @classmethod + def create_from_data_uri(cls, data_uri, metadata={}, path=""): + from llmstack.common.utils.utils import validate_parse_data_uri + + mime_type, file_name, file_data = validate_parse_data_uri(data_uri) + file_bytes = base64.b64decode(file_data) + return cls.create_from_bytes( + file_bytes, file_name, {**metadata, "mime_type": mime_type, "file_name": file_name}, path=path + ) + + class Meta: + abstract = True + + +@receiver(pre_delete) +def delete_file_on_delete(sender, instance, **kwargs): + if issubclass(sender, Assets) and instance.file: + instance.file.delete(False) diff --git a/llmstack/datasources/admin.py b/llmstack/datasources/admin.py index bada131ed85..151dd86bdbe 100644 --- a/llmstack/datasources/admin.py +++ b/llmstack/datasources/admin.py @@ -1,6 +1,6 @@ from django.contrib import admin -from .models import DataSource, DataSourceEntry, DataSourceType, UserFiles +from .models import DataSource, DataSourceEntry, DataSourceType class DataSourceEntryAdmin(admin.ModelAdmin): @@ -15,4 +15,3 @@ class DataSourceAdmin(admin.ModelAdmin): admin.site.register(DataSourceEntry, DataSourceEntryAdmin) admin.site.register(DataSourceType, DataSourceAdmin) admin.site.register(DataSource) -admin.site.register(UserFiles) diff --git a/llmstack/datasources/migrations/0005_delete_userfiles.py b/llmstack/datasources/migrations/0005_delete_userfiles.py new file mode 100644 index 00000000000..93c4311de4f --- /dev/null +++ b/llmstack/datasources/migrations/0005_delete_userfiles.py @@ -0,0 +1,16 @@ +# Generated by Django 4.2.10 on 2024-03-21 00:01 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('datasources', '0004_alter_userfiles_user'), + ] + + operations = [ + migrations.DeleteModel( + name='UserFiles', + ), + ] diff --git a/llmstack/datasources/models.py b/llmstack/datasources/models.py index d12f154ce35..214a54600bd 100644 --- a/llmstack/datasources/models.py +++ b/llmstack/datasources/models.py @@ -1,15 +1,11 @@ -import base64 import logging import uuid from django.contrib.auth.models import User from django.db import models -from django.db.models.signals import pre_delete -from django.dispatch import receiver from django.utils.timezone import now from llmstack.base.models import Profile -from llmstack.common.utils.utils import validate_parse_data_uri logger = logging.getLogger(__name__) @@ -196,58 +192,3 @@ def upload_to(instance, filename): filename, ] ) - - -class UserFiles(models.Model): - uuid = models.UUIDField(default=uuid.uuid4, editable=False, help_text="UUID of the asset") - user = models.ForeignKey(User, on_delete=models.DO_NOTHING, help_text="User this asset belongs to") - path = "" - file = models.FileField( - storage=select_storage, - upload_to=upload_to, - null=True, - blank=True, - ) - metadata = models.JSONField( - default=dict, - help_text="Metadata for the asset", - null=True, - blank=True, - ) - created_at = models.DateTimeField(auto_now_add=True) - - def __init__(self, *args, path="", **kwargs) -> None: - super(UserFiles, self).__init__(*args, **kwargs) - self.path = path - - @property - def profile_uuid(self): - return Profile.objects.get(user=self.user).uuid - - -@receiver(pre_delete, sender=UserFiles) -def delete_file_on_delete(sender, instance, **kwargs): - if instance.file: - instance.file.delete(False) - - -def create_from_bytes(user, file_bytes, filename, metadata=None, path=""): - from django.core.files.base import ContentFile - - asset = UserFiles(user=user, path=path) - asset.file.save( - filename, - ContentFile(file_bytes), - ) - bytes_size = len(file_bytes) - asset.metadata = {**metadata, "file_size": bytes_size} - asset.save() - return asset - - -def create_from_data_uri(user, data_uri, metadata={}, path=""): - mime_type, file_name, file_data = validate_parse_data_uri(data_uri) - file_bytes = base64.b64decode(file_data) - return create_from_bytes( - user, file_bytes, file_name, {**metadata, "mime_type": mime_type, "file_name": file_name}, path=path - ) diff --git a/llmstack/server/settings.py b/llmstack/server/settings.py index 9be1bc058f9..6cca2ce77d9 100644 --- a/llmstack/server/settings.py +++ b/llmstack/server/settings.py @@ -62,6 +62,7 @@ "llmstack.connections.apps.ConnectionsConfig", "llmstack.jobs.apps.JobsConfig", "llmstack.organizations.apps.OrganizationsConfig", + "llmstack.assets.apps.AssetsConfig", "flags", "allauth", "allauth.account", @@ -208,7 +209,7 @@ "base_url": GENERATEDFILES_URL, }, }, - "useruploads": { + "assets": { "BACKEND": "django.core.files.storage.FileSystemStorage", }, }