Return scores from storage providers

Author: tomusher

src/django_ai_core/contrib/index/storage/base.py

@@ -37,9 +37,10 @@ class BaseStorageDocument(VirtualModel):
     document_key: str
     content: str
     metadata: dict[str, Any]
+    score: float = 0
 
     class Meta:
-        fields = ["document_key", "content", "metadata"]
+        fields = ["document_key", "content", "metadata", "score"]
         storage_provider: "StorageProvider"
 
     def __str__(self):

src/django_ai_core/contrib/index/storage/pgvector/models.py

@@ -1,5 +1,20 @@
+from typing import Self, Sequence
+
 from django.db import models
-from pgvector.django import VectorField
+from pgvector.django import CosineDistance, VectorField
+
+
+class PgvectorEmbeddingQuerySet(models.QuerySet["BasePgVectorEmbedding"]):
+    def annotate_with_distance(
+        self,
+        query_vector: Sequence[float],
+    ) -> Self:
+        kwargs = {"distance": CosineDistance("vector", query_vector)}
+        return self.annotate(**kwargs)
+
+
+class PgvectorEmbeddingManager(models.Manager.from_queryset(PgvectorEmbeddingQuerySet)):
+    pass
 
 
 class BasePgVectorEmbedding(models.Model):
@@ -12,6 +27,8 @@ class BasePgVectorEmbedding(models.Model):
     content = models.TextField()
     metadata = models.JSONField(default=dict)
 
+    objects = PgvectorEmbeddingManager()
+
     class Meta:
         abstract = True
 

src/django_ai_core/contrib/index/storage/pgvector/provider.py

@@ -1,12 +1,10 @@
-from typing import TYPE_CHECKING, Generator, Type
-
-from pgvector.django import CosineDistance
+from typing import TYPE_CHECKING, Generator, Type, cast
 
 from ...schema import EmbeddedDocument
 from ..base import BaseStorageDocument, BaseStorageQuerySet, StorageProvider
 
 if TYPE_CHECKING:
-    from .models import BasePgVectorEmbedding
+    from .models import BasePgVectorEmbedding, PgvectorEmbeddingQuerySet
 
 
 class PgVectorQuerySet(BaseStorageQuerySet["PgVectorProvider"]):
@@ -18,6 +16,7 @@ def get_instance(self, val: "BasePgVectorEmbedding") -> BaseStorageDocument:
             document_key=val.document_key,
             content=val.content,
             metadata=val.metadata,
+            score=1 - val.distance,
         )
 
     def run_query(self) -> Generator[BaseStorageDocument, None, None]:
@@ -41,8 +40,9 @@ def run_query(self) -> Generator[BaseStorageDocument, None, None]:
             raise ValueError("Model class is required")
 
         queryset = model.objects.filter(index_name=storage_provider.index_name)
+        queryset = cast("PgvectorEmbeddingQuerySet", queryset)
 
-        queryset = queryset.order_by(CosineDistance("vector", embedding))
+        queryset = queryset.annotate_with_distance(embedding).order_by("distance")
 
         # Apply metadata filters if any
         for key, value in filter_map.items():

src/django_ai_core/contrib/index/storage/qdrant.py

@@ -1,22 +1,29 @@
+import uuid
+
 from qdrant_client import QdrantClient
 from qdrant_client.http import models as qdrant_models
 from qdrant_client.models import Distance
 
 from ..schema import EmbeddedDocument
 from .base import BaseStorageDocument, BaseStorageQuerySet, StorageProvider
 
+# Key used for storing original content in metadata
+CONTENT_METADATA_KEY = "dj_ai_core_content"
+
 
 class QdrantQuerySet(BaseStorageQuerySet["QdrantProvider"]):
     def get_instance(self, val) -> BaseStorageDocument:
         if self.model:
-            metadata = val["metadata"]
+            metadata = val.payload
+            content = metadata.pop(CONTENT_METADATA_KEY)
             return self.model(
-                document_key=val["key"],
-                content="",
+                document_key=metadata["document_key"],
+                content=content,
                 metadata=metadata,
+                score=val.score,
             )
         else:
-            return val
+            return val.payload
 
     def run_query(self):
         if not self.storage_provider:
@@ -51,7 +58,7 @@ def run_query(self):
             query_filter=qdrant_models.Filter(must=filters),
         )
 
-        for vector in response["vectors"]:
+        for vector in response:
             yield self.get_instance(vector)
 
 
@@ -89,7 +96,13 @@ def add(self, documents: list["EmbeddedDocument"]):
         for doc in documents:
             points.append(
                 qdrant_models.PointStruct(
-                    id=doc.document_key, vector=doc.vector, payload=doc.metadata
+                    id=str(uuid.uuid4()),
+                    vector=doc.vector,
+                    payload={
+                        **doc.metadata,
+                        CONTENT_METADATA_KEY: doc.content,
+                        "document_key": doc.document_key,
+                    },
                 )
             )
 
@@ -99,7 +112,16 @@ def delete(self, document_keys: list[str]):
         """Delete documents by their keys."""
         self.client.delete(
             collection_name=self.index_name,
-            points_selector=qdrant_models.PointIdsList(points=document_keys),
+            points_selector=qdrant_models.FilterSelector(
+                filter=qdrant_models.Filter(
+                    must=[
+                        qdrant_models.FieldCondition(
+                            key="document_key",
+                            match=qdrant_models.MatchAny(any=document_keys),
+                        )
+                    ]
+                )
+            ),
         )
 
     def clear(self):

src/django_ai_core/contrib/index/storage/s3vectors.py

@@ -16,6 +16,7 @@ def get_instance(self, val) -> BaseStorageDocument:
                 document_key=val["key"],
                 content=content,
                 metadata=metadata,
+                score=1 - val["distance"],
             )
         else:
             return val

tests/testapp/indexes.py

@@ -29,6 +29,15 @@
     storage_provider = S3VectorProvider(
         bucket_name="prototyping-vector-bucket", dimensions=1536
     )
+elif storage_provider_setting == "qdrant":
+    from django_ai_core.contrib.index.storage.qdrant import QdrantProvider
+
+    storage_provider = QdrantProvider(
+        host=settings.AI_CORE_TESTAPP_QDRANT_HOST,
+        port=6333,
+        api_key=settings.AI_CORE_TESTAPP_QDRANT_API_KEY,
+        dimensions=1536,
+    )
 else:
     from django_ai_core.contrib.index.storage.inmemory import InMemoryProvider
 

tests/testapp/settings.py

@@ -174,3 +174,6 @@
             "postgres://postgres:[email protected]:5432/postgres"
         ),
     }
+
+AI_CORE_TESTAPP_QDRANT_HOST = os.environ.get("AI_CORE_TESTAPP_QDRANT_HOST", None)
+AI_CORE_TESTAPP_QDRANT_API_KEY = os.environ.get("AI_CORE_TESTAPP_QDRANT_API_KEY", None)