Change model and update model util methods (#193)

Author: vegito22

llmstack/common/utils/utils.py

@@ -1,4 +1,5 @@
 import copy
+import hashlib
 import logging
 import re
 import time
@@ -103,6 +104,22 @@ def validate_parse_data_uri(
     return (mime_type, file_name, data)
 
 
+def generate_checksum(file_data, algorithm="sha256"):
+    """
+    Generates a checksum of file data using the specified algorithm.
+
+    Parameters:
+        file_data (bytes): The data of the file as bytes.
+        algorithm (str): The hashing algorithm to use (default is SHA-256).
+
+    Returns:
+        str: The checksum of the file data.
+    """
+    checksum = hashlib.new(algorithm)
+    checksum.update(file_data)
+    return checksum.hexdigest()
+
+
 def sanitize_dict_values(value):
     if isinstance(value, dict):
         return {k: sanitize_dict_values(v) for k, v in value.items()}

llmstack/datasources/migrations/0004_alter_userfiles_user.py

@@ -0,0 +1,21 @@
+# Generated by Django 4.2.10 on 2024-03-17 16:02
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('datasources', '0003_userfiles'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='userfiles',
+            name='user',
+            field=models.ForeignKey(help_text='User this asset belongs to', on_delete=django.db.models.deletion.DO_NOTHING, to=settings.AUTH_USER_MODEL),
+        ),
+    ]

llmstack/datasources/models.py

@@ -4,6 +4,8 @@
 
 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
@@ -198,7 +200,7 @@ def upload_to(instance, 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")
+    user = models.ForeignKey(User, on_delete=models.DO_NOTHING, help_text="User this asset belongs to")
     path = ""
     file = models.FileField(
         storage=select_storage,
@@ -223,10 +225,16 @@ def profile_uuid(self):
         return Profile.objects.get(user=self.user).uuid
 
 
-def create_from_bytes(user, file_bytes, filename, metadata=None):
+@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)
+    asset = UserFiles(user=user, path=path)
     asset.file.save(
         filename,
         ContentFile(file_bytes),
@@ -237,7 +245,9 @@ def create_from_bytes(user, file_bytes, filename, metadata=None):
     return asset
 
 
-def create_from_data_uri(user, data_uri, metadata={}):
+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})
+    return create_from_bytes(
+        user, file_bytes, file_name, {**metadata, "mime_type": mime_type, "file_name": file_name}, path=path
+    )