feat: support Vector Search

google/cloud/firestore_v1/_helpers.py

@@ -26,6 +26,7 @@
 
 from google.cloud import exceptions  # type: ignore
 from google.cloud._helpers import _datetime_to_pb_timestamp  # type: ignore
+from google.cloud.firestore_v1.vector import Vector
 from google.cloud.firestore_v1.types.write import DocumentTransform
 from google.cloud.firestore_v1 import transforms
 from google.cloud.firestore_v1 import types
@@ -160,7 +161,8 @@ def encode_value(value) -> types.document.Value:
 
     Args:
         value (Union[NoneType, bool, int, float, datetime.datetime, \
-            str, bytes, dict, ~google.cloud.Firestore.GeoPoint]): A native
+            str, bytes, dict, ~google.cloud.Firestore.GeoPoint, \
+            ~google.cloud.firestore_v1.vector.Vector]): A native
             Python value to convert to a protobuf field.
 
     Returns:
@@ -209,6 +211,9 @@ def encode_value(value) -> types.document.Value:
         value_pb = document.ArrayValue(values=value_list)
         return document.Value(array_value=value_pb)
 
+    if isinstance(value, Vector):
+        return encode_value(value.to_map_value())
+
     if isinstance(value, dict):
         value_dict = encode_dict(value)
         value_pb = document.MapValue(fields=value_dict)
@@ -331,7 +336,9 @@ def reference_value_to_document(reference_value, client) -> Any:
 
 def decode_value(
     value, client
-) -> Union[None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint]:
+) -> Union[
+    None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint, Vector
+]:
     """Converts a Firestore protobuf ``Value`` to a native Python value.
 
     Args:
@@ -382,7 +389,7 @@ def decode_value(
         raise ValueError("Unknown ``value_type``", value_type)
 
 
-def decode_dict(value_fields, client) -> dict:
+def decode_dict(value_fields, client) -> Union[dict, Vector]:
     """Converts a protobuf map of Firestore ``Value``-s.
 
     Args:
@@ -397,8 +404,12 @@ def decode_dict(value_fields, client) -> dict:
         of native Python values converted from the ``value_fields``.
     """
     value_fields_pb = getattr(value_fields, "_pb", value_fields)
+    res = {key: decode_value(value, client) for key, value in value_fields_pb.items()}
+
+    if res.get("__type__", None) == "__vector__":
+        return Vector(res["value"])
 
-    return {key: decode_value(value, client) for key, value in value_fields_pb.items()}
+    return res
 
 
 def get_doc_id(document_pb, expected_prefix) -> str:

google/cloud/firestore_v1/base_collection.py

@@ -19,9 +19,12 @@
 from google.api_core import retry as retries
 
 from google.cloud.firestore_v1 import _helpers
+from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
 from google.cloud.firestore_v1.document import DocumentReference
 from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
+from google.cloud.firestore_v1.base_vector_query import BaseVectorQuery
 from google.cloud.firestore_v1.base_query import QueryType
+from google.cloud.firestore_v1.vector import Vector
 
 
 from typing import (
@@ -46,6 +49,7 @@
     from google.cloud.firestore_v1.base_document import DocumentSnapshot
     from google.cloud.firestore_v1.transaction import Transaction
     from google.cloud.firestore_v1.field_path import FieldPath
+    from firestore_v1.vector_query import VectorQuery
 
 _AUTO_ID_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
 
@@ -120,6 +124,9 @@ def _query(self) -> QueryType:
     def _aggregation_query(self) -> BaseAggregationQuery:
         raise NotImplementedError
 
+    def _vector_query(self) -> BaseVectorQuery:
+        raise NotImplementedError
+
     def document(self, document_id: Optional[str] = None) -> DocumentReference:
         """Create a sub-document underneath the current collection.
 
@@ -539,6 +546,31 @@ def avg(self, field_ref: str | FieldPath, alias=None):
         """
         return self._aggregation_query().avg(field_ref, alias=alias)
 
+    def find_nearest(
+        self,
+        vector_field: str,
+        query_vector: Vector,
+        limit: int,
+        distance_measure: DistanceMeasure,
+    ) -> VectorQuery:
+        """
+        Finds the closest vector embeddings to the given query vector.
+
+        Args:
+            vector_field(str): An indexed vector field to search upon. Only documents which contain
+                vectors whose dimensionality match the query_vector can be returned.
+            query_vector(Vector): The query vector that we are searching on. Must be a vector of no more
+                than 2048 dimensions.
+            limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000.
+            distance_measure(:class:`DistanceMeasure`): The Distance Measure to use.
+
+        Returns:
+            :class`~firestore_v1.vector_query.VectorQuery`: the vector query.
+        """
+        return self._vector_query().find_nearest(
+            vector_field, query_vector, limit, distance_measure
+        )
+
 
 def _auto_id() -> str:
     """Generate a "random" automatically generated ID.

google/cloud/firestore_v1/base_query.py

@@ -33,6 +33,7 @@
 from google.cloud.firestore_v1 import document
 from google.cloud.firestore_v1 import field_path as field_path_module
 from google.cloud.firestore_v1 import transforms
+from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
 from google.cloud.firestore_v1.types import StructuredQuery
 from google.cloud.firestore_v1.types import query
 from google.cloud.firestore_v1.types import Cursor
@@ -51,11 +52,13 @@
     Union,
     TYPE_CHECKING,
 )
+from google.cloud.firestore_v1.vector import Vector
 
 # Types needed only for Type Hints
 from google.cloud.firestore_v1.base_document import DocumentSnapshot
 
 if TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.firestore_v1.base_vector_query import BaseVectorQuery
     from google.cloud.firestore_v1.field_path import FieldPath
 
 _BAD_DIR_STRING: str
@@ -972,6 +975,15 @@ def _to_protobuf(self) -> StructuredQuery:
             query_kwargs["limit"] = wrappers_pb2.Int32Value(value=self._limit)
         return query.StructuredQuery(**query_kwargs)
 
+    def find_nearest(
+        self,
+        vector_field: str,
+        queryVector: Vector,
+        limit: int,
+        distance_measure: DistanceMeasure,
+    ) -> BaseVectorQuery:
+        raise NotImplementedError
+
     def count(
         self, alias: str | None = None
     ) -> Type["firestore_v1.base_aggregation.BaseAggregationQuery"]:

google/cloud/firestore_v1/base_vector_query.py

@@ -0,0 +1,119 @@
+# Copyright 2024 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Classes for representing vector queries for the Google Cloud Firestore API.
+"""
+
+import abc
+
+from abc import ABC
+from enum import Enum
+from typing import Iterable, Optional, Tuple, Union
+from google.api_core import gapic_v1
+from google.api_core import retry as retries
+from google.cloud.firestore_v1.base_document import DocumentSnapshot
+from google.cloud.firestore_v1.types import query
+from google.cloud.firestore_v1.vector import Vector
+from google.cloud.firestore_v1 import _helpers
+
+
+class DistanceMeasure(Enum):
+    EUCLIDEAN = 1
+    COSINE = 2
+    DOT_PRODUCT = 3
+
+
+class BaseVectorQuery(ABC):
+    """Represents a vector query to the Firestore API."""
+
+    def __init__(self, nested_query) -> None:
+        self._nested_query = nested_query
+        self._collection_ref = nested_query._parent
+        self._vector_field: Optional[str] = None
+        self._query_vector: Optional[Vector] = None
+        self._limit: Optional[int] = None
+        self._distance_measure: Optional[DistanceMeasure] = None
+
+    @property
+    def _client(self):
+        return self._collection_ref._client
+
+    def _to_protobuf(self) -> query.StructuredQuery:
+        pb = query.StructuredQuery()
+
+        distance_measure_proto = None
+        if self._distance_measure == DistanceMeasure.EUCLIDEAN:
+            distance_measure_proto = (
+                query.StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN
+            )
+        elif self._distance_measure == DistanceMeasure.COSINE:
+            distance_measure_proto = (
+                query.StructuredQuery.FindNearest.DistanceMeasure.COSINE
+            )
+        elif self._distance_measure == DistanceMeasure.DOT_PRODUCT:
+            distance_measure_proto = (
+                query.StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT
+            )
+        else:
+            raise ValueError("Invalid distance_measure")
+
+        pb = self._nested_query._to_protobuf()
+        pb.find_nearest = query.StructuredQuery.FindNearest(
+            vector_field=query.StructuredQuery.FieldReference(
+                field_path=self._vector_field
+            ),
+            query_vector=_helpers.encode_value(self._query_vector),
+            distance_measure=distance_measure_proto,
+            limit=self._limit,
+        )
+        return pb
+
+    def _prep_stream(
+        self,
+        transaction=None,
+        retry: Union[retries.Retry, None, gapic_v1.method._MethodDefault] = None,
+        timeout: Optional[float] = None,
+    ) -> Tuple[dict, str, dict]:
+        parent_path, expected_prefix = self._collection_ref._parent_info()
+        request = {
+            "parent": parent_path,
+            "structured_query": self._to_protobuf(),
+            "transaction": _helpers.get_transaction_id(transaction),
+        }
+        kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
+
+        return request, expected_prefix, kwargs
+
+    @abc.abstractmethod
+    def get(
+        self,
+        transaction=None,
+        retry: retries.Retry = gapic_v1.method.DEFAULT,
+        timeout: Optional[float] = None,
+    ) -> Iterable[DocumentSnapshot]:
+        """Runs the vector query."""
+
+    def find_nearest(
+        self,
+        vector_field: str,
+        query_vector: Vector,
+        limit: int,
+        distance_measure: DistanceMeasure,
+    ):
+        """Finds the closest vector embeddings to the given query vector."""
+        self._vector_field = vector_field
+        self._query_vector = query_vector
+        self._limit = limit
+        self._distance_measure = distance_measure
+        return self

google/cloud/firestore_v1/collection.py

@@ -23,6 +23,7 @@
 )
 from google.cloud.firestore_v1 import query as query_mod
 from google.cloud.firestore_v1 import aggregation
+from google.cloud.firestore_v1 import vector_query
 from google.cloud.firestore_v1.watch import Watch
 from google.cloud.firestore_v1 import document
 from typing import Any, Callable, Generator, Tuple, Union
@@ -76,6 +77,14 @@ def _aggregation_query(self) -> aggregation.AggregationQuery:
         """
         return aggregation.AggregationQuery(self._query())
 
+    def _vector_query(self) -> vector_query.VectorQuery:
+        """VectorQuery factory.
+
+        Returns:
+            :class:`~google.cloud.firestore_v1.vector_query.VectorQuery`
+        """
+        return vector_query.VectorQuery(self._query())
+
     def add(
         self,
         document_data: dict,