diff --git a/llmstack/datasources/admin.py b/llmstack/datasources/admin.py index 151dd86bdbe..bada131ed85 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 +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) diff --git a/llmstack/datasources/migrations/0003_userfiles.py b/llmstack/datasources/migrations/0003_userfiles.py new file mode 100644 index 00000000000..1bda2cc091d --- /dev/null +++ b/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)), + ], + ), + ] diff --git a/llmstack/datasources/models.py b/llmstack/datasources/models.py index 8e27f6e2476..506ec8aeff0 100644 --- a/llmstack/datasources/models.py +++ b/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}) diff --git a/llmstack/server/settings.py b/llmstack/server/settings.py index 46520c031b9..9be1bc058f9 100644 --- a/llmstack/server/settings.py +++ b/llmstack/server/settings.py @@ -208,6 +208,9 @@ "base_url": GENERATEDFILES_URL, }, }, + "useruploads": { + "BACKEND": "django.core.files.storage.FileSystemStorage", + }, } # Default primary key field type