Initial vector db implementation

apps/utils/chunked.py

@@ -0,0 +1,32 @@
+from itertools import islice
+
+
+def chunked(it, n, collection=tuple):
+    """
+    >>> for nums in chunked(range(10), 4):
+    ...    print(nums)
+    ...
+    (0, 1, 2, 3)
+    (4, 5, 6, 7)
+    (8, 9)
+    >>> for nums in chunked(range(10), 4, list):
+    ...    print(nums)
+    ...
+    [0, 1, 2, 3]
+    [4, 5, 6, 7]
+    [8, 9]
+    """
+    itr = iter(it)
+    while True:
+        try:
+            items = take(n, itr, collection)
+        except StopIteration:
+            break
+        if not items:
+            break
+        yield items
+
+
+def take(n, iterable, collection=list):
+    # https://docs.python.org/2/library/itertools.html#recipes
+    return collection(islice(iterable, n))

apps/vectordb/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class VectorDbConfig(AppConfig):
+    name = "apps.vectordb"
+    label = "vectordb"

apps/vectordb/const.py

@@ -0,0 +1,9 @@
+META_EMBEDDING_ID = "embedding_id"
+
+META_EXPERIMENT_ID = "experiment_id"
+
+META_FILE_ID = "file_id"
+
+META_SEARCH_SCORE = "search_score"
+
+META_ALL = (META_EMBEDDING_ID, META_EXPERIMENT_ID, META_FILE_ID, META_SEARCH_SCORE)

apps/vectordb/migrations/0001_initial.py

@@ -0,0 +1,66 @@
+# Generated by Django 4.2.7 on 2024-03-22 12:11
+
+from django.db import migrations, models
+import django.db.models.deletion
+import pgvector.django
+
+
+class Migration(migrations.Migration):
+    initial = True
+
+    dependencies = [
+        ("files", "0001_initial"),
+        ("experiments", "0070_alter_consentform_name_alter_experiment_llm_and_more"),
+        ("teams", "0005_invitation_groups"),
+    ]
+
+    operations = [
+        pgvector.django.VectorExtension(),
+        migrations.CreateModel(
+            name="Embedding",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("embedding", pgvector.django.VectorField(dimensions=1536)),
+                ("document", models.TextField(null=True)),
+                ("metadata", models.JSONField(null=True)),
+                (
+                    "experiment",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="embeddings",
+                        to="experiments.experiment",
+                    ),
+                ),
+                (
+                    "file",
+                    models.ForeignKey(
+                        blank=True,
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="files.file",
+                    ),
+                ),
+                (
+                    "team",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="teams.team",
+                        verbose_name="Team",
+                    ),
+                ),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]

apps/vectordb/models.py

@@ -0,0 +1,14 @@
+from django.db import models
+from pgvector.django import VectorField
+
+from apps.teams.models import BaseTeamModel
+
+ADA_TOKEN_COUNT = 1536
+
+
+class Embedding(BaseTeamModel):
+    experiment = models.ForeignKey("experiments.Experiment", on_delete=models.CASCADE, related_name="embeddings")
+    embedding = VectorField(dimensions=ADA_TOKEN_COUNT)
+    document = models.TextField(null=True)  # noqa: DJ001
+    metadata = models.JSONField(null=True)
+    file = models.ForeignKey("files.File", on_delete=models.CASCADE, null=True, blank=True)

apps/vectordb/tests/test_vectorstore.py

@@ -0,0 +1,143 @@
+"""Test PGVector functionality.
+
+Ported from https://github.com/hwchase17/langchain/blob/master/tests/integration_tests/vectorstores/test_pgvector.py
+"""
+
+import pytest
+from langchain.docstore.document import Document
+from langchain_community.embeddings import DeterministicFakeEmbedding
+
+from apps.utils.factories.experiment import ExperimentFactory
+from apps.vectordb.const import META_ALL
+from apps.vectordb.models import ADA_TOKEN_COUNT
+from apps.vectordb.vectorstore import DistanceStrategy, PGVector
+
+
+@pytest.fixture()
+def experiment(db):
+    return ExperimentFactory()
+
+
+@pytest.mark.django_db()
+def test_vectorstore(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        experiment=experiment,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    _remove_meta_fields(output)
+    assert output == [Document(page_content="foo")]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_with_metadatas(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    _remove_meta_fields(output)
+    assert output == [Document(page_content="foo", metadata={"page": "0"})]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_with_metadatas_with_scores(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+        distance_strategy=DistanceStrategy.COSINE,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=1)
+    _remove_meta_fields(output)
+    assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_with_filter_match(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "0"})
+    _remove_meta_fields(output)
+    assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_with_filter_distant_match(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "2"})
+    _remove_meta_fields(output)
+    # ordering here is deterministic but random due to fake embeddings
+    assert output == [(Document(page_content="baz", metadata={"page": "2"}), 0.9290842232061864)]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_distant_match_cosine_ordering(experiment):
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+        distance_strategy=DistanceStrategy.COSINE,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=3)
+    _remove_meta_fields(output)
+    # ordering here is deterministic but random due to fake embeddings
+    assert output == [
+        (Document(page_content="foo", metadata={"page": "0"}), 0.0),
+        (Document(page_content="baz", metadata={"page": "2"}), 0.9290842232061864),
+        (Document(page_content="bar", metadata={"page": "1"}), 1.0246011368123038),
+    ]
+
+
+@pytest.mark.django_db()
+def test_vectorstore_with_filter_no_match(experiment):
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        embedding=DeterministicFakeEmbedding(size=ADA_TOKEN_COUNT),
+        metadatas=metadatas,
+        experiment=experiment,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "5"})
+    assert output == []
+
+
+def _remove_meta_fields(docs, fields=META_ALL):
+    for doc in docs:
+        if isinstance(doc, tuple):
+            # (doc, score)
+            doc = doc[0]
+        for field in fields:
+            doc.metadata.pop(field, None)