refactored models

grai-io · Oct 23, 2023 · df6ab34 · df6ab34
1 parent fd09f7c
commit df6ab34
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 42 deletions.
diff --git a/grai-server/app/grAI/migrations/0001_initial.py b/grai-server/app/grAI/migrations/0001_initial.py
@@ -1,4 +1,4 @@
-# Generated by Django 4.2.6 on 2023-10-22 16:55
+# Generated by Django 4.2.6 on 2023-10-23 14:52
 
 from django.db import migrations, models
 import django.db.models.deletion
@@ -16,9 +16,24 @@ class Migration(migrations.Migration):
         migrations.CreateModel(
             name="UserChat",
             fields=[
-                ("id", models.AutoField(primary_key=True, serialize=False)),
-                ("chat_id", models.UUIDField(default=uuid.uuid4, editable=False)),
-                ("message_number", models.IntegerField(default=0)),
+                ("chat_id", models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                (
+                    "membership",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="membership",
+                        to="workspaces.membership",
+                    ),
+                ),
+            ],
+        ),
+        migrations.CreateModel(
+            name="Message",
+            fields=[
+                ("id", models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("message", models.TextField()),
+                ("visible", models.BooleanField()),
                 (
                     "role",
                     models.CharField(
@@ -27,28 +42,21 @@ class Migration(migrations.Migration):
                         max_length=255,
                     ),
                 ),
-                ("message", models.TextField()),
                 ("created_at", models.DateTimeField(auto_now_add=True)),
                 (
-                    "membership",
+                    "chat",
                     models.ForeignKey(
-                        on_delete=django.db.models.deletion.CASCADE,
-                        related_name="membership",
-                        to="workspaces.membership",
+                        on_delete=django.db.models.deletion.CASCADE, related_name="messages", to="grAI.userchat"
                     ),
                 ),
             ],
-            options={
-                "indexes": [
-                    models.Index(fields=["chat_id"], name="grAI_userch_chat_id_730147_idx"),
-                    models.Index(fields=["membership", "created_at"], name="grAI_userch_members_65b04e_idx"),
-                ],
-            },
         ),
-        migrations.AddConstraint(
+        migrations.AddIndex(
             model_name="userchat",
-            constraint=models.UniqueConstraint(
-                fields=("chat_id", "message_number"), name="Message numbers are unique for each chat"
-            ),
+            index=models.Index(fields=["chat_id"], name="grAI_userch_chat_id_730147_idx"),
+        ),
+        migrations.AddIndex(
+            model_name="message",
+            index=models.Index(fields=["created_at"], name="grAI_messag_created_7c8856_idx"),
         ),
     ]
diff --git a/grai-server/app/grAI/models.py b/grai-server/app/grAI/models.py
@@ -6,37 +6,29 @@
 class UserChat(models.Model):
     """ """
 
+    chat_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+    membership = models.ForeignKey("workspaces.Membership", related_name="membership", on_delete=models.CASCADE)
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    class Meta:
+        indexes = [
+            models.Index(fields=["chat_id"]),
+        ]
+
+
+class Message(models.Model):
     USER = "user"
     AGENT = "agent"
     SYSTEM = "system"
     ROLES = [(USER, "user"), (AGENT, "agent"), (SYSTEM, "system")]
 
-    id = models.AutoField(primary_key=True)
-    membership = models.ForeignKey("workspaces.Membership", related_name="membership", on_delete=models.CASCADE)
-    chat_id = models.UUIDField(default=uuid.uuid4, editable=False)
-    message_number = models.IntegerField(default=0)
-    role = models.CharField(max_length=255, choices=ROLES, default=USER)
+    chat = models.ForeignKey(UserChat, on_delete=models.CASCADE, related_name="messages")
     message = models.TextField()
+    visible = models.BooleanField()
+    role = models.CharField(max_length=255, choices=ROLES, default=USER)
     created_at = models.DateTimeField(auto_now_add=True)
 
-    def save(self, *args, **kwargs):
-        if self.chat_id is None:
-            self.message_number = None
-        elif self.message_number is None:
-            max_message_number = UserChat.objects.filter(membership=self.membership, chat_id=self.chat_id).aggregate(
-                models.Max("message_number")
-            )["message_number__max"]
-            self.message_number = max_message_number + 1
-        super().save(*args, **kwargs)
-
     class Meta:
-        constraints = [
-            models.UniqueConstraint(
-                fields=["chat_id", "message_number"],
-                name="Message numbers are unique for each chat",
-            ),
-        ]
         indexes = [
-            models.Index(fields=["chat_id"]),
-            models.Index(fields=["membership", "created_at"]),
+            models.Index(fields=["created_at"]),
         ]