Add a new model for user assets (#190)

Author: vegito22

llmstack/datasources/admin.py

@@ -1,6 +1,6 @@
 from django.contrib import admin
 
-from .models import DataSource, DataSourceEntry, DataSourceType
+from .models import DataSource, DataSourceEntry, DataSourceType, UserFiles
 
 
 class DataSourceEntryAdmin(admin.ModelAdmin):
@@ -15,3 +15,4 @@ class DataSourceAdmin(admin.ModelAdmin):
 admin.site.register(DataSourceEntry, DataSourceEntryAdmin)
 admin.site.register(DataSourceType, DataSourceAdmin)
 admin.site.register(DataSource)
+admin.site.register(UserFiles)

llmstack/datasources/migrations/0003_userfiles.py

@@ -0,0 +1,29 @@
+# Generated by Django 4.2.10 on 2024-03-16 10:53
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import llmstack.datasources.models
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('datasources', '0002_datasource_config_and_more'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='UserFiles',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, help_text='UUID of the asset')),
+                ('file', models.FileField(blank=True, null=True, storage=llmstack.datasources.models.select_storage, upload_to=llmstack.datasources.models.upload_to)),
+                ('metadata', models.JSONField(blank=True, default=dict, help_text='Metadata for the asset', null=True)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('user', models.OneToOneField(help_text='User this asset belongs to', on_delete=django.db.models.deletion.DO_NOTHING, to=settings.AUTH_USER_MODEL)),
+            ],
+        ),
+    ]

llmstack/datasources/models.py

@@ -1,3 +1,4 @@
+import base64
 import logging
 import uuid
 
@@ -6,6 +7,7 @@
 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__)
 
@@ -176,3 +178,66 @@ def user_can_read(self, user) -> bool:
                 ).organization
             )
         return False
+
+
+def select_storage():
+    from django.core.files.storage import storages
+
+    return storages["useruploads"]
+
+
+def upload_to(instance, filename):
+    return "/".join(
+        [
+            str(instance.profile_uuid),
+            instance.path,
+            filename,
+        ]
+    )
+
+
+class UserFiles(models.Model):
+    uuid = models.UUIDField(default=uuid.uuid4, editable=False, help_text="UUID of the asset")
+    user = models.OneToOneField(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
+
+
+def create_from_bytes(user, file_bytes, filename, metadata=None):
+    from django.core.files.base import ContentFile
+
+    asset = UserFiles(user=user)
+    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={}):
+    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})

llmstack/server/settings.py

@@ -208,6 +208,9 @@
             "base_url": GENERATEDFILES_URL,
         },
     },
+    "useruploads": {
+        "BACKEND": "django.core.files.storage.FileSystemStorage",
+    },
 }
 
 # Default primary key field type