feat(vector-search): enable Firestore vector search across sync and async models

firedantic/__init__.py

@@ -1,4 +1,6 @@
 # flake8: noqa
+from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
+
 from firedantic._async.helpers import truncate_collection as async_truncate_collection
 from firedantic._async.indexes import (
     set_up_composite_indexes as async_set_up_composite_indexes,
@@ -31,7 +33,8 @@
     SubModel,
 )
 from firedantic._sync.ttl_policy import set_up_ttl_policies
-from firedantic.common import collection_group_index, collection_index
+from firedantic.common import FiredanticVector as Vector
+from firedantic.common import VectorConfig, collection_group_index, collection_index
 from firedantic.configurations import (
     CONFIGURATIONS,
     configure,

firedantic/_async/indexes.py

@@ -1,5 +1,5 @@
 from logging import getLogger
-from typing import Iterable, List, Optional, Set, Type
+from typing import Any, Dict, Iterable, List, Optional, Set, Type
 
 from google.api_core.operation_async import AsyncOperation
 from google.cloud.firestore_admin_v1 import (
@@ -13,7 +13,7 @@
 
 from firedantic._async.model import AsyncBareModel
 from firedantic._async.ttl_policy import set_up_ttl_policies
-from firedantic.common import IndexDefinition, IndexField
+from firedantic.common import IndexDefinition, IndexField, VectorConfig
 
 logger = getLogger("firedantic")
 
@@ -41,12 +41,20 @@ async def get_existing_indexes(
         if not raw_index.name.startswith(path):
             continue
         query_scope = raw_index.query_scope.name
-        fields = tuple(
-            IndexField(name=f.field_path, order=f.order.name)  # noqa
-            for f in raw_index.fields
-            if f.field_path != "__name__"
-        )
-        indexes.add(IndexDefinition(query_scope=query_scope, fields=fields))
+        fields = []
+        for f in raw_index.fields:
+            if f.field_path == "__name__":
+                continue
+            if f.vector_config:
+                vector_config = VectorConfig(
+                    dimension=f.vector_config.dimension, flat=True
+                )
+                fields.append(
+                    IndexField(name=f.field_path, vector_config=vector_config)
+                )
+            else:
+                fields.append(IndexField(name=f.field_path, order=f.order.name))
+        indexes.add(IndexDefinition(query_scope=query_scope, fields=tuple(fields)))
     return indexes
 
 
@@ -63,16 +71,24 @@ async def create_composite_index(
     :param path: Index path in Firestore.
     :return: Operation that was launched to create the index.
     """
+    fields = []
+    for field in index.fields:
+        f_dict: Dict[str, Any] = {"field_path": field.name}
+        if field.order is not None:
+            f_dict["order"] = field.order
+        if field.vector_config is not None:
+            f_dict["vector_config"] = {"dimension": field.vector_config.dimension}
+            if field.vector_config.flat:
+                f_dict["vector_config"]["flat"] = {}
+        fields.append(f_dict)
+
     request = CreateIndexRequest(
         {
             "parent": path,
             "index": Index(
                 {
                     "query_scope": index.query_scope,
-                    "fields": [
-                        {"field_path": field[0], "order": field[1]}
-                        for field in list(index.fields)
-                    ],
+                    "fields": fields,
                 }
             ),
         }

firedantic/_async/model.py

@@ -1,4 +1,5 @@
 from abc import ABC
+from enum import Enum
 from logging import getLogger
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
 
@@ -11,6 +12,8 @@
 )
 from google.cloud.firestore_v1.async_query import AsyncQuery
 from google.cloud.firestore_v1.async_transaction import AsyncTransaction
+from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
+from google.cloud.firestore_v1.vector import Vector
 
 import firedantic.operators as op
 from firedantic import async_truncate_collection
@@ -20,7 +23,10 @@
     CollectionNotDefined,
     InvalidDocumentID,
     ModelNotFoundError,
+    VectorFieldAmbiguous,
+    VectorFieldNotDefined,
 )
+from firedantic.utils import classproperty
 
 TAsyncBareModel = TypeVar("TAsyncBareModel", bound="AsyncBareModel")
 TAsyncBareSubModel = TypeVar("TAsyncBareSubModel", bound="AsyncBareSubModel")
@@ -67,6 +73,10 @@ class AsyncBareModel(pydantic.BaseModel, ABC):
     Implements basic functionality for Pydantic models, such as save, delete, find etc.
     """
 
+    model_config = pydantic.ConfigDict(
+        ignored_types=(classproperty,), arbitrary_types_allowed=True
+    )
+
     __collection__: Optional[str] = None
     __document_id__: str
     __ttl_field__: Optional[str] = None
@@ -244,6 +254,114 @@ async def find_one(
         except IndexError as e:
             raise ModelNotFoundError(f"No '{cls.__name__}' found") from e
 
+    @classmethod
+    async def vector_search(
+        cls: Type[TAsyncBareModel],
+        query_vector: Vector,
+        limit: int,
+        distance_measure: DistanceMeasure,
+        vector_field: Optional[Union[str, Enum]] = None,
+        filter_: Optional[Dict[str, Union[str, dict]]] = None,
+        distance_result_field: Optional[str] = None,
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> List[TAsyncBareModel]:
+        """
+        Returns a list of models from the database based on a vector search.
+
+        :param query_vector: The vector to search for.
+        :param limit: Maximum results to return.
+        :param distance_measure: The distance measure to use.
+        :param vector_field: The field to search for the vector. If not specified, it will
+            be automatically resolved if the model has exactly one vector field.
+        :param filter_: The filter criteria.
+        :param distance_result_field: Optional field to store the distance result.
+        :param transaction: Optional transaction to use.
+        :return: List of found models.
+        :raise VectorFieldNotDefined: If no vector fields are defined in the model.
+        :raise VectorFieldAmbiguous: If multiple vector fields are found and none is specified.
+        """
+        vector_field = cls._resolve_vector_field(vector_field)
+        query: Union[AsyncQuery, AsyncCollectionReference] = cls._get_col_ref()
+        if filter_:
+            for key, value in filter_.items():
+                query = cls._add_filter(query, key, value)
+
+        vector_query = query.find_nearest(
+            vector_field=vector_field,
+            query_vector=query_vector,
+            limit=limit,
+            distance_measure=distance_measure,
+            distance_result_field=distance_result_field,
+        )
+
+        def _cls(doc_id: str, data: Dict[str, Any]) -> TAsyncBareModel:
+            if cls.__document_id__ in data:
+                logger.warning(
+                    "%s document ID %s contains conflicting %s in data with value %s",
+                    cls.__name__,
+                    doc_id,
+                    cls.__document_id__,
+                    data[cls.__document_id__],
+                )
+            data[cls.__document_id__] = doc_id
+            model = cls(**data)
+            setattr(model, cls.__document_id__, doc_id)
+            return model
+
+        return [
+            _cls(doc.id, doc_dict)
+            async for doc in vector_query.stream(transaction=transaction)  # type: ignore
+            if (doc_dict := doc.to_dict()) is not None
+        ]
+
+    @classmethod
+    def _get_vector_fields(cls) -> List[str]:
+        """
+        Returns a list of fields that are annotated as Vector.
+        """
+        return [
+            name
+            for name, field in cls.model_fields.items()
+            if field.annotation is Vector
+        ]
+
+    @classmethod
+    def get_vector_fields_enum(cls) -> Type[Enum]:
+        """
+        Returns an Enum of vector fields for this model.
+        """
+        fields = cls._get_vector_fields()
+        return Enum(f"{cls.__name__}VectorFields", {f: f for f in fields})  # type: ignore
+
+    @classproperty
+    def VectorFields(cls) -> Type[Enum]:  # type: ignore[no-redef]
+        """
+        Returns an Enum of vector fields for this model.
+        """
+        return cls.get_vector_fields_enum()
+
+    @classmethod
+    def _resolve_vector_field(
+        cls, vector_field: Optional[Union[str, Enum]] = None
+    ) -> str:
+        """
+        Resolves the vector field to use for search.
+        """
+        if vector_field:
+            if isinstance(vector_field, Enum):
+                return str(vector_field.value)
+            return vector_field
+
+        fields = cls._get_vector_fields()
+        if len(fields) == 0:
+            raise VectorFieldNotDefined(f"No vector fields defined in {cls.__name__}")
+        if len(fields) > 1:
+            raise VectorFieldAmbiguous(
+                f"Multiple vector fields found in {cls.__name__}: {', '.join(fields)}. "
+                "Please specify which one to use."
+            )
+        return fields[0]
+
     @classmethod
     async def get_by_doc_id(
         cls: Type[TAsyncBareModel],

firedantic/_sync/indexes.py

@@ -1,5 +1,5 @@
 from logging import getLogger
-from typing import Iterable, List, Optional, Set, Type
+from typing import Any, Dict, Iterable, List, Optional, Set, Type
 
 from google.api_core.operation import Operation
 from google.cloud.firestore_admin_v1 import (
@@ -13,7 +13,7 @@
 
 from firedantic._sync.model import BareModel
 from firedantic._sync.ttl_policy import set_up_ttl_policies
-from firedantic.common import IndexDefinition, IndexField
+from firedantic.common import IndexDefinition, IndexField, VectorConfig
 
 logger = getLogger("firedantic")
 
@@ -41,12 +41,20 @@ def get_existing_indexes(
         if not raw_index.name.startswith(path):
             continue
         query_scope = raw_index.query_scope.name
-        fields = tuple(
-            IndexField(name=f.field_path, order=f.order.name)  # noqa
-            for f in raw_index.fields
-            if f.field_path != "__name__"
-        )
-        indexes.add(IndexDefinition(query_scope=query_scope, fields=fields))
+        fields = []
+        for f in raw_index.fields:
+            if f.field_path == "__name__":
+                continue
+            if f.vector_config:
+                vector_config = VectorConfig(
+                    dimension=f.vector_config.dimension, flat=True
+                )
+                fields.append(
+                    IndexField(name=f.field_path, vector_config=vector_config)
+                )
+            else:
+                fields.append(IndexField(name=f.field_path, order=f.order.name))
+        indexes.add(IndexDefinition(query_scope=query_scope, fields=tuple(fields)))
     return indexes
 
 
@@ -63,16 +71,24 @@ def create_composite_index(
     :param path: Index path in Firestore.
     :return: Operation that was launched to create the index.
     """
+    fields = []
+    for field in index.fields:
+        f_dict: Dict[str, Any] = {"field_path": field.name}
+        if field.order is not None:
+            f_dict["order"] = field.order
+        if field.vector_config is not None:
+            f_dict["vector_config"] = {"dimension": field.vector_config.dimension}
+            if field.vector_config.flat:
+                f_dict["vector_config"]["flat"] = {}
+        fields.append(f_dict)
+
     request = CreateIndexRequest(
         {
             "parent": path,
             "index": Index(
                 {
                     "query_scope": index.query_scope,
-                    "fields": [
-                        {"field_path": field[0], "order": field[1]}
-                        for field in list(index.fields)
-                    ],
+                    "fields": fields,
                 }
             ),
         }