diff --git a/src/flint/db/migrations/0005_alter_conversationvector_vector.py b/src/flint/db/migrations/0005_alter_conversationvector_vector.py new file mode 100644 index 0000000..611f8ac --- /dev/null +++ b/src/flint/db/migrations/0005_alter_conversationvector_vector.py @@ -0,0 +1,18 @@ +# Generated by Django 4.2.4 on 2023-09-11 17:01 + +from django.db import migrations +import pgvector.django + + +class Migration(migrations.Migration): + dependencies = [ + ("db", "0004_conversationvector"), + ] + + operations = [ + migrations.AlterField( + model_name="conversationvector", + name="vector", + field=pgvector.django.VectorField(dimensions=1024), + ), + ] diff --git a/src/flint/db/migrations/0006_conversationvector_vector_index.py b/src/flint/db/migrations/0006_conversationvector_vector_index.py new file mode 100644 index 0000000..9d28ecb --- /dev/null +++ b/src/flint/db/migrations/0006_conversationvector_vector_index.py @@ -0,0 +1,19 @@ +# Generated by Django 4.2.4 on 2023-09-11 17:01 + +from django.db import migrations +import pgvector.django + + +class Migration(migrations.Migration): + dependencies = [ + ("db", "0005_alter_conversationvector_vector"), + ] + + operations = [ + migrations.AddIndex( + model_name="conversationvector", + index=pgvector.django.IvfflatIndex( + fields=["vector"], lists=100, name="vector_index", opclasses=["vector_cosine_ops"] + ), + ), + ] diff --git a/src/flint/db/models.py b/src/flint/db/models.py index df91661..5282f4e 100644 --- a/src/flint/db/models.py +++ b/src/flint/db/models.py @@ -1,4 +1,5 @@ from flint.manage import init_django + init_django() import uuid @@ -9,7 +10,8 @@ from django.db.models.signals import post_save from django.dispatch import receiver -from pgvector.django import VectorField +from pgvector.django import VectorField, IvfflatIndex + class BaseModel(models.Model): created_at = models.DateTimeField(auto_now_add=True) @@ -23,7 +25,8 @@ def __str__(self): def __repr__(self): return self.__str__() - + + class KhojUser(BaseModel): id = models.AutoField(primary_key=True) user = models.OneToOneField(User, on_delete=models.CASCADE) @@ -36,12 +39,14 @@ class Meta: def __str__(self): return f"{self.__class__.__name__}({self.id})" - + + @receiver(post_save, sender=User) def create_khoj_user(sender, instance, created, **kwargs): if created: KhojUser.objects.create(user=instance) + @receiver(post_save, sender=User) def save_khoj_user(sender, instance, **kwargs): instance.khojuser.save() @@ -59,16 +64,20 @@ class Meta: def __str__(self): return f"{self.__class__.__name__}({self.id})" - + + class ConversationVector(BaseModel): id = models.BigAutoField(primary_key=True) conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, null=False, related_name="vectors") compiled = models.TextField() - vector = VectorField(1024) + vector = VectorField(dimensions=1024) class Meta: db_table = "conversation_vector" ordering = ["-created_at"] + indexes = [ + IvfflatIndex(name="vector_index", fields=["vector"], lists=100, opclasses=["vector_cosine_ops"]), + ] def __str__(self): - return f"{self.__class__.__name__}({self.id})" \ No newline at end of file + return f"{self.__class__.__name__}({self.id})"