diff --git a/firedantic/__init__.py b/firedantic/__init__.py index e4f79d2..9acd83e 100644 --- a/firedantic/__init__.py +++ b/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, diff --git a/firedantic/_async/indexes.py b/firedantic/_async/indexes.py index b7d0d52..39fb91d 100644 --- a/firedantic/_async/indexes.py +++ b/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, } ), } diff --git a/firedantic/_async/model.py b/firedantic/_async/model.py index 84251df..8556ae7 100644 --- a/firedantic/_async/model.py +++ b/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], diff --git a/firedantic/_sync/indexes.py b/firedantic/_sync/indexes.py index a1eb7c5..99b15b5 100644 --- a/firedantic/_sync/indexes.py +++ b/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, } ), } diff --git a/firedantic/_sync/model.py b/firedantic/_sync/model.py index 78ec372..f62490c 100644 --- a/firedantic/_sync/model.py +++ b/firedantic/_sync/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 @@ -10,7 +11,9 @@ FieldFilter, ) from google.cloud.firestore_v1.base_query import BaseQuery +from google.cloud.firestore_v1.base_vector_query import DistanceMeasure from google.cloud.firestore_v1.transaction import Transaction +from google.cloud.firestore_v1.vector import Vector import firedantic.operators as op from firedantic import truncate_collection @@ -20,7 +23,10 @@ CollectionNotDefined, InvalidDocumentID, ModelNotFoundError, + VectorFieldAmbiguous, + VectorFieldNotDefined, ) +from firedantic.utils import classproperty TBareModel = TypeVar("TBareModel", bound="BareModel") TBareSubModel = TypeVar("TBareSubModel", bound="BareSubModel") @@ -67,6 +73,10 @@ class BareModel(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 @@ -242,6 +252,114 @@ def find_one( except IndexError as e: raise ModelNotFoundError(f"No '{cls.__name__}' found") from e + @classmethod + def vector_search( + cls: Type[TBareModel], + 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[Transaction] = None, + ) -> List[TBareModel]: + """ + 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[BaseQuery, CollectionReference] = 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]) -> TBareModel: + 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) + 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 def get_by_doc_id( cls: Type[TBareModel], diff --git a/firedantic/common.py b/firedantic/common.py index d8103da..76a5e91 100644 --- a/firedantic/common.py +++ b/firedantic/common.py @@ -1,8 +1,26 @@ -from typing import Literal, NamedTuple, Tuple, Union +from typing import Annotated, Literal, NamedTuple, Optional, Tuple, Union + +from google.cloud.firestore_v1.vector import Vector +from pydantic import BeforeValidator OrderDirection = Union[Literal["ASCENDING"], Literal["DESCENDING"]] -IndexField = NamedTuple("IndexField", [("name", str), ("order", OrderDirection)]) +FiredanticVector = Annotated[ + Vector, + BeforeValidator(lambda v: v if isinstance(v, Vector) else Vector(v)), +] + + +class VectorConfig(NamedTuple): + dimension: int + flat: bool = True + + +class IndexField(NamedTuple): + name: str + order: Optional[OrderDirection] = None + vector_config: Optional[VectorConfig] = None + IndexDefinition = NamedTuple( "IndexDefinition", [("query_scope", str), ("fields", Tuple[IndexField, ...])] diff --git a/firedantic/exceptions.py b/firedantic/exceptions.py index 5f8dbd4..a742143 100644 --- a/firedantic/exceptions.py +++ b/firedantic/exceptions.py @@ -12,3 +12,11 @@ class ModelNotFoundError(ModelError): class CollectionNotDefined(ModelError): """Raised when the model collection is not defined.""" + + +class VectorFieldNotDefined(ModelError): + """Raised when a vector field is not defined in the model.""" + + +class VectorFieldAmbiguous(ModelError): + """Raised when there are multiple vector fields and none is specified.""" diff --git a/firedantic/utils.py b/firedantic/utils.py index e6e14e1..242c76a 100644 --- a/firedantic/utils.py +++ b/firedantic/utils.py @@ -8,3 +8,15 @@ def get_all_subclasses(cls) -> Iterator: for subclass in cls.__subclasses__(): yield subclass yield from get_all_subclasses(subclass) + + +class classproperty: + """ + A decorator that behaves like @property but for class methods. + """ + + def __init__(self, fget): + self.fget = fget + + def __get__(self, owner, cls): + return self.fget(cls)