rakam-systems-vectorstore 0.1.1rc7__py3-none-any.whl

rakam_systems_vectorstore/components/vectorstore/migrations/0001_initial.py

@@ -0,0 +1,55 @@
+# Generated by Django for PgVectorStore
+
+from django.db import migrations, models
+import django.db.models.deletion
+import pgvector.django
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+    ]
+
+    operations = [
+        migrations.RunSQL(
+            "CREATE EXTENSION IF NOT EXISTS vector;",
+            reverse_sql="DROP EXTENSION IF EXISTS vector;"
+        ),
+        migrations.CreateModel(
+            name='Collection',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('name', models.CharField(max_length=255, unique=True)),
+                ('embedding_dim', models.IntegerField(default=384)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('updated_at', models.DateTimeField(auto_now=True)),
+            ],
+            options={
+                'db_table': 'application_collection',
+            },
+        ),
+        migrations.CreateModel(
+            name='NodeEntry',
+            fields=[
+                ('node_id', models.AutoField(primary_key=True, serialize=False)),
+                ('content', models.TextField()),
+                ('embedding', pgvector.django.VectorField(dimensions=384)),
+                ('source_file_uuid', models.CharField(max_length=255)),
+                ('position', models.IntegerField(blank=True, null=True)),
+                ('custom_metadata', models.JSONField(blank=True, default=dict)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('updated_at', models.DateTimeField(auto_now=True)),
+                ('collection', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='nodes', to='application.collection')),
+            ],
+            options={
+                'db_table': 'application_nodeentry',
+                'indexes': [
+                    models.Index(fields=['source_file_uuid'], name='application_source__idx'),
+                    models.Index(fields=['collection', 'source_file_uuid'], name='application_collect_idx'),
+                ],
+            },
+        ),
+    ]
+

rakam_systems_vectorstore/components/vectorstore/models.py

@@ -0,0 +1,10 @@
+"""
+Django models for Vector Store.
+
+This module re-exports models from pg_models to make them discoverable by Django's
+automatic model discovery mechanism, which looks for a models.py file.
+"""
+from rakam_systems_vectorstore.components.vectorstore.pg_models import Collection  # noqa: F401
+from rakam_systems_vectorstore.components.vectorstore.pg_models import NodeEntry  # noqa: F401
+
+__all__ = ["Collection", "NodeEntry"]

rakam_systems_vectorstore/components/vectorstore/pg_models.py

@@ -0,0 +1,97 @@
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from django.db import models
+from pgvector.django import VectorField
+
+if TYPE_CHECKING:
+    from rakam_systems_vectorstore.core import Node
+
+
+class Collection(models.Model):
+    """
+    Represents a collection of vector embeddings in the database.
+    Model with explicit app_label to avoid Django app registration issues.
+    """
+
+    name = models.CharField(max_length=255, unique=True)
+    embedding_dim = models.IntegerField(
+        default=384
+    )  # Default dimension, but dynamically set based on embedding model
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
+
+    class Meta:
+        app_label = "application"
+        db_table = "application_collection"
+
+    def __str__(self):
+        return f"Collection: {self.name}"
+
+
+class NodeEntry(models.Model):
+    """
+    Represents a stored node entry with content, metadata, and vector embedding.
+
+    Note: The embedding field uses no fixed dimension, allowing it to store vectors
+    of any dimension. This provides flexibility to use different embedding models
+    without needing to alter the database schema.
+    """
+
+    collection = models.ForeignKey(
+        Collection, on_delete=models.CASCADE, related_name="nodes"
+    )
+    content = models.TextField()
+    # No dimension specified - allows vectors of any dimension
+    embedding = VectorField(dimensions=None)
+
+    # Node metadata
+    node_id = models.AutoField(primary_key=True)
+    source_file_uuid = models.CharField(max_length=255)
+    position = models.IntegerField(null=True, blank=True)
+    custom_metadata = models.JSONField(default=dict, blank=True)
+
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
+
+    class Meta:
+        app_label = "application"
+        db_table = "application_nodeentry"
+        indexes = [
+            models.Index(fields=["source_file_uuid"],
+                         name="application_source__idx"),
+            models.Index(
+                fields=["collection", "source_file_uuid"], name="application_collect_idx"),
+        ]
+
+    def __str__(self):
+        return f"Node {self.node_id}: {self.content[:30]}..."
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the node entry to a dictionary."""
+        return {
+            "node_id": self.node_id,
+            "content": self.content,
+            "source_file_uuid": self.source_file_uuid,
+            "position": self.position,
+            "custom": self.custom_metadata,
+        }
+
+    def to_node(self) -> "Node":
+        """Convert the database entry to a Node object."""
+        from rakam_systems_vectorstore.core import Node, NodeMetadata
+
+        metadata = NodeMetadata(
+            source_file_uuid=self.source_file_uuid,
+            position=self.position,
+            custom=self.custom_metadata,
+        )
+        metadata.node_id = self.node_id
+
+        node = Node(content=self.content, metadata=metadata)
+        # Convert the database vector field to a numpy array
+        node.embedding = self.embedding
+
+        return node