From 8b64b09643e8f5f595379a67f70d2101645a5eb8 Mon Sep 17 00:00:00 2001 From: "cto-new[bot]" <140088366+cto-new[bot]@users.noreply.github.com> Date: Tue, 20 Jan 2026 00:43:21 +0000 Subject: [PATCH 1/3] feat(vector-search): add Firestore vector search support for sync and async models Adds vector search capability for Firestore while preserving Firedantic/Pydantic semantics. - Introduces Vector and DistanceMeasure support and a VectorConfig type - Extends model vector_search API for async and sync variants - Enables vector fields to be used in composite indexes via vector_config - Exposes vector types in the top-level exports for convenient usage Impact: - Adds a new vector_search workflow using Firestore's nearest neighbor queries - No breaking changes; usage is optional via new API --- firedantic/__init__.py | 9 ++- firedantic/_async/indexes.py | 38 ++++++++---- firedantic/_async/model.py | 58 +++++++++++++++++ firedantic/_sync/helpers.py | 4 +- firedantic/_sync/indexes.py | 38 ++++++++---- firedantic/_sync/model.py | 62 ++++++++++++++++++- firedantic/common.py | 22 ++++++- firedantic/tests/tests_sync/conftest.py | 4 +- firedantic/tests/tests_sync/test_indexes.py | 19 ++++-- firedantic/tests/tests_sync/test_model.py | 53 ++++++++++++++-- .../tests/tests_sync/test_ttl_policy.py | 2 + 11 files changed, 270 insertions(+), 39 deletions(-) diff --git a/firedantic/__init__.py b/firedantic/__init__.py index e4f79d2..a53a516 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,12 @@ SubModel, ) from firedantic._sync.ttl_policy import set_up_ttl_policies -from firedantic.common import collection_group_index, collection_index +from firedantic.common import ( + VectorConfig, + collection_group_index, + collection_index, +) +from firedantic.common import FiredanticVector as Vector from firedantic.configurations import ( CONFIGURATIONS, configure, diff --git a/firedantic/_async/indexes.py b/firedantic/_async/indexes.py index b7d0d52..31c9c88 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,18 @@ 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 +69,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..57a06cd 100644 --- a/firedantic/_async/model.py +++ b/firedantic/_async/model.py @@ -11,6 +11,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 @@ -244,6 +246,62 @@ 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], + vector_field: str, + query_vector: Vector, + limit: int, + distance_measure: DistanceMeasure, + 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 vector_field: The field to search for the vector. + :param query_vector: The vector to search for. + :param limit: Maximum results to return. + :param distance_measure: The distance measure to use. + :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. + """ + 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) + if (doc_dict := doc.to_dict()) is not None + ] + @classmethod async def get_by_doc_id( cls: Type[TAsyncBareModel], diff --git a/firedantic/_sync/helpers.py b/firedantic/_sync/helpers.py index 57dfe99..3a5f92c 100644 --- a/firedantic/_sync/helpers.py +++ b/firedantic/_sync/helpers.py @@ -1,7 +1,9 @@ from google.cloud.firestore_v1 import CollectionReference -def truncate_collection(col_ref: CollectionReference, batch_size: int = 128) -> int: +def truncate_collection( + col_ref: CollectionReference, batch_size: int = 128 +) -> int: """ Removes all documents inside a collection. diff --git a/firedantic/_sync/indexes.py b/firedantic/_sync/indexes.py index a1eb7c5..7912fcc 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,18 @@ 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 +69,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..5803628 100644 --- a/firedantic/_sync/model.py +++ b/firedantic/_sync/model.py @@ -11,6 +11,8 @@ ) from google.cloud.firestore_v1.base_query import BaseQuery from google.cloud.firestore_v1.transaction import Transaction +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 truncate_collection @@ -236,12 +238,70 @@ def find_one( :return: The model instance. :raise ModelNotFoundError: If the entry is not found. """ - model = cls.find(filter_, limit=1, order_by=order_by, transaction=transaction) + model = cls.find( + filter_, limit=1, order_by=order_by, transaction=transaction + ) try: return model[0] except IndexError as e: raise ModelNotFoundError(f"No '{cls.__name__}' found") from e + @classmethod + def vector_search( + cls: Type[TBareModel], + vector_field: str, + query_vector: Vector, + limit: int, + distance_measure: DistanceMeasure, + 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 vector_field: The field to search for the vector. + :param query_vector: The vector to search for. + :param limit: Maximum results to return. + :param distance_measure: The distance measure to use. + :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. + """ + 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) + if (doc_dict := doc.to_dict()) is not None + ] + @classmethod def get_by_doc_id( cls: Type[TBareModel], diff --git a/firedantic/common.py b/firedantic/common.py index d8103da..02c65f3 100644 --- a/firedantic/common.py +++ b/firedantic/common.py @@ -1,8 +1,26 @@ -from typing import Literal, NamedTuple, Tuple, Union +from typing import Annotated, Any, Dict, 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/tests/tests_sync/conftest.py b/firedantic/tests/tests_sync/conftest.py index 8debfc4..410769d 100644 --- a/firedantic/tests/tests_sync/conftest.py +++ b/firedantic/tests/tests_sync/conftest.py @@ -195,7 +195,9 @@ def _create( @pytest.fixture def create_product(): - def _create(product_id: Optional[str] = None, price: float = 1.23, stock: int = 3): + def _create( + product_id: Optional[str] = None, price: float = 1.23, stock: int = 3 + ): if not product_id: product_id = str(uuid.uuid4()) p = Product(product_id=product_id, price=price, stock=stock) diff --git a/firedantic/tests/tests_sync/test_indexes.py b/firedantic/tests/tests_sync/test_indexes.py index c4fb399..a0b302c 100644 --- a/firedantic/tests/tests_sync/test_indexes.py +++ b/firedantic/tests/tests_sync/test_indexes.py @@ -7,10 +7,10 @@ from firedantic import ( CONFIGURATIONS, Model, - collection_group_index, - collection_index, set_up_composite_indexes, set_up_composite_indexes_and_ttl_policies, + collection_group_index, + collection_index, ) from firedantic.common import IndexField from firedantic.tests.tests_sync.conftest import MockListIndexOperation @@ -26,6 +26,7 @@ class BaseModelWithIndexes(Model): age: int + def test_set_up_composite_index(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -58,6 +59,7 @@ class ModelWithIndexes(BaseModelWithIndexes): assert index.fields[1].order.name == Query.DESCENDING + def test_set_up_collection_group_index(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -86,6 +88,7 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(index.fields) == 2 + def test_set_up_composite_indexes_and_policies(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -109,6 +112,7 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(call_list) == 1 + def test_set_up_many_composite_indexes(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -135,6 +139,7 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(result) == 3 + def test_set_up_indexes_model_without_indexes(mock_admin_client) -> None: class ModelWithoutIndexes(Model): __collection__ = "modelWithoutIndexes" @@ -152,6 +157,7 @@ class ModelWithoutIndexes(Model): assert len(call_list) == 0 + def test_existing_indexes_are_skipped(mock_admin_client) -> None: resp = ListIndexesResponse( { @@ -183,7 +189,9 @@ def test_existing_indexes_are_skipped(mock_admin_client) -> None: ] } ) - mock_admin_client.list_indexes = Mock(return_value=MockListIndexOperation([resp])) + mock_admin_client.list_indexes = Mock( + return_value=MockListIndexOperation([resp]) + ) class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -205,6 +213,7 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(result) == 0 + def test_same_fields_in_another_collection(mock_admin_client) -> None: # Test that when another collection has an index with exactly the same fields, # it won't affect creating an index in the target collection @@ -226,7 +235,9 @@ def test_same_fields_in_another_collection(mock_admin_client) -> None: ] } ) - mock_admin_client.list_indexes = Mock(return_value=MockListIndexOperation([resp])) + mock_admin_client.list_indexes = Mock( + return_value=MockListIndexOperation([resp]) + ) class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( diff --git a/firedantic/tests/tests_sync/test_model.py b/firedantic/tests/tests_sync/test_model.py index 37e27a4..ccd9778 100644 --- a/firedantic/tests/tests_sync/test_model.py +++ b/firedantic/tests/tests_sync/test_model.py @@ -35,6 +35,7 @@ ] + def test_save_model(configure_db, create_company) -> None: company = create_company() @@ -43,6 +44,7 @@ def test_save_model(configure_db, create_company) -> None: assert company.owner.last_name == "Doe" + def test_delete_model(configure_db, create_company) -> None: company: Company = create_company( company_id="11223344-5", first_name="Jane", last_name="Doe" @@ -57,6 +59,7 @@ def test_delete_model(configure_db, create_company) -> None: Company.get_by_id(_id) + def test_find_one(configure_db, create_company) -> None: with pytest.raises(ModelNotFoundError): Company.find_one() @@ -81,10 +84,13 @@ def test_find_one(configure_db, create_company) -> None: first_asc = Company.find_one(order_by=[("owner.first_name", Query.ASCENDING)]) assert first_asc.owner.first_name == "Bar" - first_desc = Company.find_one(order_by=[("owner.first_name", Query.DESCENDING)]) + first_desc = Company.find_one( + order_by=[("owner.first_name", Query.DESCENDING)] + ) assert first_desc.owner.first_name == "Foo" + def test_find(configure_db, create_company, create_product) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -115,6 +121,7 @@ def test_find(configure_db, create_company, create_product) -> None: Product.find({"product_id": {"<>": "a"}}) + def test_find_not_in(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -135,6 +142,7 @@ def test_find_not_in(configure_db, create_company) -> None: assert company.company_id in ("2131232-4", "4124432-4") + def test_find_array_contains(configure_db, create_todolist) -> None: list_1 = create_todolist("list_1", ["Work", "Eat", "Sleep"]) create_todolist("list_2", ["Learn Python", "Walk the dog"]) @@ -144,6 +152,7 @@ def test_find_array_contains(configure_db, create_todolist) -> None: assert found[0].name == list_1.name + def test_find_array_contains_any(configure_db, create_todolist) -> None: list_1 = create_todolist("list_1", ["Work", "Eat"]) list_2 = create_todolist("list_2", ["Relax", "Chill", "Sleep"]) @@ -155,6 +164,7 @@ def test_find_array_contains_any(configure_db, create_todolist) -> None: assert lst.name in (list_1.name, list_2.name) + def test_find_limit(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -167,6 +177,7 @@ def test_find_limit(configure_db, create_company) -> None: assert len(companies_2) == 2 + def test_find_order_by(configure_db, create_company) -> None: companies_and_owners = [ {"company_id": "1234555-1", "last_name": "A", "first_name": "A"}, @@ -179,9 +190,13 @@ def test_find_order_by(configure_db, create_company) -> None: {"company_id": "4124432-5", "last_name": "D", "first_name": "H"}, ] - companies_and_owners = [create_company(**item) for item in companies_and_owners] + companies_and_owners = [ + create_company(**item) for item in companies_and_owners + ] - companies_ascending = Company.find(order_by=[("owner.first_name", Query.ASCENDING)]) + companies_ascending = Company.find( + order_by=[("owner.first_name", Query.ASCENDING)] + ) assert companies_ascending == companies_and_owners companies_descending = Company.find( @@ -211,6 +226,7 @@ def test_find_order_by(configure_db, create_company) -> None: assert companies_and_owners == lastname_ascending_firstname_ascending + def test_find_offset(configure_db, create_company) -> None: ids_and_lastnames = ( ("1234555-1", "A"), @@ -228,6 +244,7 @@ def test_find_offset(configure_db, create_company) -> None: assert len(companies_ascending) == 2 + def test_get_by_id(configure_db, create_company) -> None: c: Company = create_company(company_id="1234567-8") @@ -242,11 +259,13 @@ def test_get_by_id(configure_db, create_company) -> None: assert c_2.owner.first_name == "John" + def test_get_by_empty_str_id(configure_db) -> None: with pytest.raises(ModelNotFoundError): Company.get_by_id("") + def test_missing_collection(configure_db) -> None: class User(Model): name: str @@ -255,6 +274,7 @@ class User(Model): User(name="John").save() + def test_model_aliases(configure_db) -> None: class User(Model): __collection__ = "User" @@ -271,6 +291,7 @@ class User(Model): assert user_from_db.city == "Helsinki" + @pytest.mark.parametrize( "model_id", [ @@ -312,6 +333,7 @@ def test_models_with_valid_custom_id(configure_db, model_id) -> None: found.delete() + @pytest.mark.parametrize( "model_id", [ @@ -337,6 +359,7 @@ def test_models_with_invalid_custom_id(configure_db, model_id: str) -> None: Product.get_by_id(model_id) + def test_truncate_collection(configure_db, create_company) -> None: create_company(company_id="1234567-8") create_company(company_id="1234567-9") @@ -349,6 +372,7 @@ def test_truncate_collection(configure_db, create_company) -> None: assert len(new_companies) == 0 + def test_custom_id_model(configure_db) -> None: c = CustomIDModel(bar="bar") # type: ignore c.save() @@ -361,6 +385,7 @@ def test_custom_id_model(configure_db) -> None: assert m.bar == "bar" + def test_custom_id_conflict(configure_db) -> None: CustomIDConflictModel(foo="foo", bar="bar").save() @@ -372,6 +397,7 @@ def test_custom_id_conflict(configure_db) -> None: assert m.bar == "bar" + def test_model_id_persistency(configure_db) -> None: c = CustomIDConflictModel(foo="foo", bar="bar") c.save() @@ -383,6 +409,7 @@ def test_model_id_persistency(configure_db) -> None: assert len(CustomIDConflictModel.find({})) == 1 + def test_bare_model_document_id_persistency(configure_db) -> None: c = CustomIDModel(bar="bar") # type: ignore c.save() @@ -394,17 +421,20 @@ def test_bare_model_document_id_persistency(configure_db) -> None: assert len(CustomIDModel.find({})) == 1 + def test_bare_model_get_by_empty_doc_id(configure_db) -> None: with pytest.raises(ModelNotFoundError): CustomIDModel.get_by_doc_id("") + def test_extra_fields(configure_db) -> None: CustomIDModelExtra(foo="foo", bar="bar", baz="baz").save() # type: ignore with pytest.raises(ValidationError): CustomIDModel.find({}) + def test_company_stats(configure_db, create_company) -> None: company: Company = create_company(company_id="1234567-8") company_stats = company.stats() @@ -425,6 +455,7 @@ def test_company_stats(configure_db, create_company) -> None: assert stats.sales == 101 + def test_subcollection_model_safety(configure_db) -> None: """ Ensure you shouldn't be able to use unprepared subcollection models accidentally @@ -433,6 +464,7 @@ def test_subcollection_model_safety(configure_db) -> None: UserStats.find({}) + def test_get_user_purchases(configure_db) -> None: u = User(name="Foo") u.save() @@ -444,6 +476,7 @@ def test_get_user_purchases(configure_db) -> None: assert get_user_purchases(u.id) == 42 + def test_reload(configure_db) -> None: u = User(name="Foo") u.save() @@ -462,6 +495,7 @@ def test_reload(configure_db) -> None: another_user.reload() + def test_save_with_exclude_none(configure_db) -> None: p = Profile(name="Foo") p.save(exclude_none=True) @@ -483,6 +517,7 @@ def test_save_with_exclude_none(configure_db) -> None: assert data == {"name": "Foo", "photo_url": None} + def test_save_with_exclude_unset(configure_db) -> None: p = Profile(photo_url=None) p.save(exclude_unset=True) @@ -504,6 +539,7 @@ def test_save_with_exclude_unset(configure_db) -> None: assert data == {"name": "", "photo_url": None} + def test_update_city_in_transaction(configure_db) -> None: """ Test updating a model in a transaction. Test case from README. @@ -512,7 +548,9 @@ def test_update_city_in_transaction(configure_db) -> None: """ @transactional - def decrement_population(transaction: Transaction, city: City, decrement: int = 1): + def decrement_population( + transaction: Transaction, city: City, decrement: int = 1 + ): city.reload(transaction=transaction) city.population = max(0, city.population - decrement) city.save(transaction=transaction) @@ -527,6 +565,7 @@ def decrement_population(transaction: Transaction, city: City, decrement: int = assert c.population == 0 + def test_delete_in_transaction(configure_db) -> None: """ Test deleting a model in a transaction. @@ -535,7 +574,9 @@ def test_delete_in_transaction(configure_db) -> None: """ @transactional - def delete_in_transaction(transaction: Transaction, profile_id: str) -> None: + def delete_in_transaction( + transaction: Transaction, profile_id: str + ) -> None: """Deletes a Profile in a transaction.""" profile = Profile.get_by_id(profile_id, transaction=transaction) profile.delete(transaction=transaction) @@ -551,6 +592,7 @@ def delete_in_transaction(transaction: Transaction, profile_id: str) -> None: Profile.get_by_id(p.id) + def test_update_model_in_transaction(configure_db) -> None: """ Test updating a model in a transaction. @@ -577,6 +619,7 @@ def update_in_transaction( assert p.name == "Bar" + def test_update_submodel_in_transaction(configure_db) -> None: """ Test Updating a submodel in a transaction. diff --git a/firedantic/tests/tests_sync/test_ttl_policy.py b/firedantic/tests/tests_sync/test_ttl_policy.py index 8d2d549..b3a1cb1 100644 --- a/firedantic/tests/tests_sync/test_ttl_policy.py +++ b/firedantic/tests/tests_sync/test_ttl_policy.py @@ -5,6 +5,7 @@ from firedantic.tests.tests_sync.conftest import ExpiringModel + def test_set_up_ttl_policies_new_policy(mock_admin_client): result = set_up_ttl_policies( gcloud_project="fake-project", models=[ExpiringModel], client=mock_admin_client @@ -26,6 +27,7 @@ def test_set_up_ttl_policies_new_policy(mock_admin_client): [Field.TtlConfig.State.NEEDS_REPAIR], ), ) + def test_set_up_ttl_policies_other_states(mock_admin_client, state): mock_admin_client.field_state = Field.TtlConfig.State.ACTIVE result = set_up_ttl_policies( From 1b0c771784d239990ae8ac43809d2989d4850bfc Mon Sep 17 00:00:00 2001 From: "cto-new[bot]" <140088366+cto-new[bot]@users.noreply.github.com> Date: Tue, 20 Jan 2026 01:41:25 +0000 Subject: [PATCH 2/3] feat(vector-search): add vector search support for Firedantic - Introduce Firestore vector search integration within Firedantic's vector ecosystem. - Extend configuration to support vector fields and vector-based queries while preserving backward compatibility. - Implement vector_search in both async and sync model bases to enable nearest-neighbor lookups. - Expose vector types for consumer usage to simplify adoption and integration. --- firedantic/__init__.py | 6 +-- firedantic/_async/indexes.py | 4 +- firedantic/_async/model.py | 2 +- firedantic/_sync/helpers.py | 4 +- firedantic/_sync/indexes.py | 4 +- firedantic/_sync/model.py | 8 ++- firedantic/common.py | 2 +- firedantic/tests/tests_sync/conftest.py | 4 +- firedantic/tests/tests_sync/test_indexes.py | 19 ++----- firedantic/tests/tests_sync/test_model.py | 53 ++----------------- .../tests/tests_sync/test_ttl_policy.py | 2 - 11 files changed, 23 insertions(+), 85 deletions(-) diff --git a/firedantic/__init__.py b/firedantic/__init__.py index a53a516..9acd83e 100644 --- a/firedantic/__init__.py +++ b/firedantic/__init__.py @@ -33,12 +33,8 @@ SubModel, ) from firedantic._sync.ttl_policy import set_up_ttl_policies -from firedantic.common import ( - VectorConfig, - 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 31c9c88..39fb91d 100644 --- a/firedantic/_async/indexes.py +++ b/firedantic/_async/indexes.py @@ -49,7 +49,9 @@ async def get_existing_indexes( vector_config = VectorConfig( dimension=f.vector_config.dimension, flat=True ) - fields.append(IndexField(name=f.field_path, vector_config=vector_config)) + 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))) diff --git a/firedantic/_async/model.py b/firedantic/_async/model.py index 57a06cd..e4e90cc 100644 --- a/firedantic/_async/model.py +++ b/firedantic/_async/model.py @@ -298,7 +298,7 @@ def _cls(doc_id: str, data: Dict[str, Any]) -> TAsyncBareModel: return [ _cls(doc.id, doc_dict) - async for doc in vector_query.stream(transaction=transaction) + async for doc in vector_query.stream(transaction=transaction) # type: ignore if (doc_dict := doc.to_dict()) is not None ] diff --git a/firedantic/_sync/helpers.py b/firedantic/_sync/helpers.py index 3a5f92c..57dfe99 100644 --- a/firedantic/_sync/helpers.py +++ b/firedantic/_sync/helpers.py @@ -1,9 +1,7 @@ from google.cloud.firestore_v1 import CollectionReference -def truncate_collection( - col_ref: CollectionReference, batch_size: int = 128 -) -> int: +def truncate_collection(col_ref: CollectionReference, batch_size: int = 128) -> int: """ Removes all documents inside a collection. diff --git a/firedantic/_sync/indexes.py b/firedantic/_sync/indexes.py index 7912fcc..99b15b5 100644 --- a/firedantic/_sync/indexes.py +++ b/firedantic/_sync/indexes.py @@ -49,7 +49,9 @@ def get_existing_indexes( vector_config = VectorConfig( dimension=f.vector_config.dimension, flat=True ) - fields.append(IndexField(name=f.field_path, vector_config=vector_config)) + 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))) diff --git a/firedantic/_sync/model.py b/firedantic/_sync/model.py index 5803628..53454dc 100644 --- a/firedantic/_sync/model.py +++ b/firedantic/_sync/model.py @@ -10,8 +10,8 @@ FieldFilter, ) from google.cloud.firestore_v1.base_query import BaseQuery -from google.cloud.firestore_v1.transaction import Transaction 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 @@ -238,9 +238,7 @@ def find_one( :return: The model instance. :raise ModelNotFoundError: If the entry is not found. """ - model = cls.find( - filter_, limit=1, order_by=order_by, transaction=transaction - ) + model = cls.find(filter_, limit=1, order_by=order_by, transaction=transaction) try: return model[0] except IndexError as e: @@ -298,7 +296,7 @@ def _cls(doc_id: str, data: Dict[str, Any]) -> TBareModel: return [ _cls(doc.id, doc_dict) - for doc in vector_query.stream(transaction=transaction) + for doc in vector_query.stream(transaction=transaction) # type: ignore if (doc_dict := doc.to_dict()) is not None ] diff --git a/firedantic/common.py b/firedantic/common.py index 02c65f3..76a5e91 100644 --- a/firedantic/common.py +++ b/firedantic/common.py @@ -1,4 +1,4 @@ -from typing import Annotated, Any, Dict, Literal, NamedTuple, Optional, Tuple, Union +from typing import Annotated, Literal, NamedTuple, Optional, Tuple, Union from google.cloud.firestore_v1.vector import Vector from pydantic import BeforeValidator diff --git a/firedantic/tests/tests_sync/conftest.py b/firedantic/tests/tests_sync/conftest.py index 410769d..8debfc4 100644 --- a/firedantic/tests/tests_sync/conftest.py +++ b/firedantic/tests/tests_sync/conftest.py @@ -195,9 +195,7 @@ def _create( @pytest.fixture def create_product(): - def _create( - product_id: Optional[str] = None, price: float = 1.23, stock: int = 3 - ): + def _create(product_id: Optional[str] = None, price: float = 1.23, stock: int = 3): if not product_id: product_id = str(uuid.uuid4()) p = Product(product_id=product_id, price=price, stock=stock) diff --git a/firedantic/tests/tests_sync/test_indexes.py b/firedantic/tests/tests_sync/test_indexes.py index a0b302c..c4fb399 100644 --- a/firedantic/tests/tests_sync/test_indexes.py +++ b/firedantic/tests/tests_sync/test_indexes.py @@ -7,10 +7,10 @@ from firedantic import ( CONFIGURATIONS, Model, - set_up_composite_indexes, - set_up_composite_indexes_and_ttl_policies, collection_group_index, collection_index, + set_up_composite_indexes, + set_up_composite_indexes_and_ttl_policies, ) from firedantic.common import IndexField from firedantic.tests.tests_sync.conftest import MockListIndexOperation @@ -26,7 +26,6 @@ class BaseModelWithIndexes(Model): age: int - def test_set_up_composite_index(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -59,7 +58,6 @@ class ModelWithIndexes(BaseModelWithIndexes): assert index.fields[1].order.name == Query.DESCENDING - def test_set_up_collection_group_index(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -88,7 +86,6 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(index.fields) == 2 - def test_set_up_composite_indexes_and_policies(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -112,7 +109,6 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(call_list) == 1 - def test_set_up_many_composite_indexes(mock_admin_client) -> None: class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -139,7 +135,6 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(result) == 3 - def test_set_up_indexes_model_without_indexes(mock_admin_client) -> None: class ModelWithoutIndexes(Model): __collection__ = "modelWithoutIndexes" @@ -157,7 +152,6 @@ class ModelWithoutIndexes(Model): assert len(call_list) == 0 - def test_existing_indexes_are_skipped(mock_admin_client) -> None: resp = ListIndexesResponse( { @@ -189,9 +183,7 @@ def test_existing_indexes_are_skipped(mock_admin_client) -> None: ] } ) - mock_admin_client.list_indexes = Mock( - return_value=MockListIndexOperation([resp]) - ) + mock_admin_client.list_indexes = Mock(return_value=MockListIndexOperation([resp])) class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( @@ -213,7 +205,6 @@ class ModelWithIndexes(BaseModelWithIndexes): assert len(result) == 0 - def test_same_fields_in_another_collection(mock_admin_client) -> None: # Test that when another collection has an index with exactly the same fields, # it won't affect creating an index in the target collection @@ -235,9 +226,7 @@ def test_same_fields_in_another_collection(mock_admin_client) -> None: ] } ) - mock_admin_client.list_indexes = Mock( - return_value=MockListIndexOperation([resp]) - ) + mock_admin_client.list_indexes = Mock(return_value=MockListIndexOperation([resp])) class ModelWithIndexes(BaseModelWithIndexes): __composite_indexes__ = ( diff --git a/firedantic/tests/tests_sync/test_model.py b/firedantic/tests/tests_sync/test_model.py index ccd9778..37e27a4 100644 --- a/firedantic/tests/tests_sync/test_model.py +++ b/firedantic/tests/tests_sync/test_model.py @@ -35,7 +35,6 @@ ] - def test_save_model(configure_db, create_company) -> None: company = create_company() @@ -44,7 +43,6 @@ def test_save_model(configure_db, create_company) -> None: assert company.owner.last_name == "Doe" - def test_delete_model(configure_db, create_company) -> None: company: Company = create_company( company_id="11223344-5", first_name="Jane", last_name="Doe" @@ -59,7 +57,6 @@ def test_delete_model(configure_db, create_company) -> None: Company.get_by_id(_id) - def test_find_one(configure_db, create_company) -> None: with pytest.raises(ModelNotFoundError): Company.find_one() @@ -84,13 +81,10 @@ def test_find_one(configure_db, create_company) -> None: first_asc = Company.find_one(order_by=[("owner.first_name", Query.ASCENDING)]) assert first_asc.owner.first_name == "Bar" - first_desc = Company.find_one( - order_by=[("owner.first_name", Query.DESCENDING)] - ) + first_desc = Company.find_one(order_by=[("owner.first_name", Query.DESCENDING)]) assert first_desc.owner.first_name == "Foo" - def test_find(configure_db, create_company, create_product) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -121,7 +115,6 @@ def test_find(configure_db, create_company, create_product) -> None: Product.find({"product_id": {"<>": "a"}}) - def test_find_not_in(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -142,7 +135,6 @@ def test_find_not_in(configure_db, create_company) -> None: assert company.company_id in ("2131232-4", "4124432-4") - def test_find_array_contains(configure_db, create_todolist) -> None: list_1 = create_todolist("list_1", ["Work", "Eat", "Sleep"]) create_todolist("list_2", ["Learn Python", "Walk the dog"]) @@ -152,7 +144,6 @@ def test_find_array_contains(configure_db, create_todolist) -> None: assert found[0].name == list_1.name - def test_find_array_contains_any(configure_db, create_todolist) -> None: list_1 = create_todolist("list_1", ["Work", "Eat"]) list_2 = create_todolist("list_2", ["Relax", "Chill", "Sleep"]) @@ -164,7 +155,6 @@ def test_find_array_contains_any(configure_db, create_todolist) -> None: assert lst.name in (list_1.name, list_2.name) - def test_find_limit(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -177,7 +167,6 @@ def test_find_limit(configure_db, create_company) -> None: assert len(companies_2) == 2 - def test_find_order_by(configure_db, create_company) -> None: companies_and_owners = [ {"company_id": "1234555-1", "last_name": "A", "first_name": "A"}, @@ -190,13 +179,9 @@ def test_find_order_by(configure_db, create_company) -> None: {"company_id": "4124432-5", "last_name": "D", "first_name": "H"}, ] - companies_and_owners = [ - create_company(**item) for item in companies_and_owners - ] + companies_and_owners = [create_company(**item) for item in companies_and_owners] - companies_ascending = Company.find( - order_by=[("owner.first_name", Query.ASCENDING)] - ) + companies_ascending = Company.find(order_by=[("owner.first_name", Query.ASCENDING)]) assert companies_ascending == companies_and_owners companies_descending = Company.find( @@ -226,7 +211,6 @@ def test_find_order_by(configure_db, create_company) -> None: assert companies_and_owners == lastname_ascending_firstname_ascending - def test_find_offset(configure_db, create_company) -> None: ids_and_lastnames = ( ("1234555-1", "A"), @@ -244,7 +228,6 @@ def test_find_offset(configure_db, create_company) -> None: assert len(companies_ascending) == 2 - def test_get_by_id(configure_db, create_company) -> None: c: Company = create_company(company_id="1234567-8") @@ -259,13 +242,11 @@ def test_get_by_id(configure_db, create_company) -> None: assert c_2.owner.first_name == "John" - def test_get_by_empty_str_id(configure_db) -> None: with pytest.raises(ModelNotFoundError): Company.get_by_id("") - def test_missing_collection(configure_db) -> None: class User(Model): name: str @@ -274,7 +255,6 @@ class User(Model): User(name="John").save() - def test_model_aliases(configure_db) -> None: class User(Model): __collection__ = "User" @@ -291,7 +271,6 @@ class User(Model): assert user_from_db.city == "Helsinki" - @pytest.mark.parametrize( "model_id", [ @@ -333,7 +312,6 @@ def test_models_with_valid_custom_id(configure_db, model_id) -> None: found.delete() - @pytest.mark.parametrize( "model_id", [ @@ -359,7 +337,6 @@ def test_models_with_invalid_custom_id(configure_db, model_id: str) -> None: Product.get_by_id(model_id) - def test_truncate_collection(configure_db, create_company) -> None: create_company(company_id="1234567-8") create_company(company_id="1234567-9") @@ -372,7 +349,6 @@ def test_truncate_collection(configure_db, create_company) -> None: assert len(new_companies) == 0 - def test_custom_id_model(configure_db) -> None: c = CustomIDModel(bar="bar") # type: ignore c.save() @@ -385,7 +361,6 @@ def test_custom_id_model(configure_db) -> None: assert m.bar == "bar" - def test_custom_id_conflict(configure_db) -> None: CustomIDConflictModel(foo="foo", bar="bar").save() @@ -397,7 +372,6 @@ def test_custom_id_conflict(configure_db) -> None: assert m.bar == "bar" - def test_model_id_persistency(configure_db) -> None: c = CustomIDConflictModel(foo="foo", bar="bar") c.save() @@ -409,7 +383,6 @@ def test_model_id_persistency(configure_db) -> None: assert len(CustomIDConflictModel.find({})) == 1 - def test_bare_model_document_id_persistency(configure_db) -> None: c = CustomIDModel(bar="bar") # type: ignore c.save() @@ -421,20 +394,17 @@ def test_bare_model_document_id_persistency(configure_db) -> None: assert len(CustomIDModel.find({})) == 1 - def test_bare_model_get_by_empty_doc_id(configure_db) -> None: with pytest.raises(ModelNotFoundError): CustomIDModel.get_by_doc_id("") - def test_extra_fields(configure_db) -> None: CustomIDModelExtra(foo="foo", bar="bar", baz="baz").save() # type: ignore with pytest.raises(ValidationError): CustomIDModel.find({}) - def test_company_stats(configure_db, create_company) -> None: company: Company = create_company(company_id="1234567-8") company_stats = company.stats() @@ -455,7 +425,6 @@ def test_company_stats(configure_db, create_company) -> None: assert stats.sales == 101 - def test_subcollection_model_safety(configure_db) -> None: """ Ensure you shouldn't be able to use unprepared subcollection models accidentally @@ -464,7 +433,6 @@ def test_subcollection_model_safety(configure_db) -> None: UserStats.find({}) - def test_get_user_purchases(configure_db) -> None: u = User(name="Foo") u.save() @@ -476,7 +444,6 @@ def test_get_user_purchases(configure_db) -> None: assert get_user_purchases(u.id) == 42 - def test_reload(configure_db) -> None: u = User(name="Foo") u.save() @@ -495,7 +462,6 @@ def test_reload(configure_db) -> None: another_user.reload() - def test_save_with_exclude_none(configure_db) -> None: p = Profile(name="Foo") p.save(exclude_none=True) @@ -517,7 +483,6 @@ def test_save_with_exclude_none(configure_db) -> None: assert data == {"name": "Foo", "photo_url": None} - def test_save_with_exclude_unset(configure_db) -> None: p = Profile(photo_url=None) p.save(exclude_unset=True) @@ -539,7 +504,6 @@ def test_save_with_exclude_unset(configure_db) -> None: assert data == {"name": "", "photo_url": None} - def test_update_city_in_transaction(configure_db) -> None: """ Test updating a model in a transaction. Test case from README. @@ -548,9 +512,7 @@ def test_update_city_in_transaction(configure_db) -> None: """ @transactional - def decrement_population( - transaction: Transaction, city: City, decrement: int = 1 - ): + def decrement_population(transaction: Transaction, city: City, decrement: int = 1): city.reload(transaction=transaction) city.population = max(0, city.population - decrement) city.save(transaction=transaction) @@ -565,7 +527,6 @@ def decrement_population( assert c.population == 0 - def test_delete_in_transaction(configure_db) -> None: """ Test deleting a model in a transaction. @@ -574,9 +535,7 @@ def test_delete_in_transaction(configure_db) -> None: """ @transactional - def delete_in_transaction( - transaction: Transaction, profile_id: str - ) -> None: + def delete_in_transaction(transaction: Transaction, profile_id: str) -> None: """Deletes a Profile in a transaction.""" profile = Profile.get_by_id(profile_id, transaction=transaction) profile.delete(transaction=transaction) @@ -592,7 +551,6 @@ def delete_in_transaction( Profile.get_by_id(p.id) - def test_update_model_in_transaction(configure_db) -> None: """ Test updating a model in a transaction. @@ -619,7 +577,6 @@ def update_in_transaction( assert p.name == "Bar" - def test_update_submodel_in_transaction(configure_db) -> None: """ Test Updating a submodel in a transaction. diff --git a/firedantic/tests/tests_sync/test_ttl_policy.py b/firedantic/tests/tests_sync/test_ttl_policy.py index b3a1cb1..8d2d549 100644 --- a/firedantic/tests/tests_sync/test_ttl_policy.py +++ b/firedantic/tests/tests_sync/test_ttl_policy.py @@ -5,7 +5,6 @@ from firedantic.tests.tests_sync.conftest import ExpiringModel - def test_set_up_ttl_policies_new_policy(mock_admin_client): result = set_up_ttl_policies( gcloud_project="fake-project", models=[ExpiringModel], client=mock_admin_client @@ -27,7 +26,6 @@ def test_set_up_ttl_policies_new_policy(mock_admin_client): [Field.TtlConfig.State.NEEDS_REPAIR], ), ) - def test_set_up_ttl_policies_other_states(mock_admin_client, state): mock_admin_client.field_state = Field.TtlConfig.State.ACTIVE result = set_up_ttl_policies( From 192926b911f3c03b4e8c43e76e157fba27da733e Mon Sep 17 00:00:00 2001 From: "cto-new[bot]" <140088366+cto-new[bot]@users.noreply.github.com> Date: Tue, 20 Jan 2026 01:59:22 +0000 Subject: [PATCH 3/3] feat(vector-search): auto-resolve vector fields and add Enum support --- firedantic/_async/model.py | 64 ++++++++++++++++++++++++++++++++++++-- firedantic/_sync/model.py | 64 ++++++++++++++++++++++++++++++++++++-- firedantic/exceptions.py | 8 +++++ firedantic/utils.py | 12 +++++++ 4 files changed, 144 insertions(+), 4 deletions(-) diff --git a/firedantic/_async/model.py b/firedantic/_async/model.py index e4e90cc..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 @@ -22,7 +23,10 @@ CollectionNotDefined, InvalidDocumentID, ModelNotFoundError, + VectorFieldAmbiguous, + VectorFieldNotDefined, ) +from firedantic.utils import classproperty TAsyncBareModel = TypeVar("TAsyncBareModel", bound="AsyncBareModel") TAsyncBareSubModel = TypeVar("TAsyncBareSubModel", bound="AsyncBareSubModel") @@ -69,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 @@ -249,10 +257,10 @@ async def find_one( @classmethod async def vector_search( cls: Type[TAsyncBareModel], - vector_field: str, 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, @@ -260,15 +268,19 @@ async def vector_search( """ Returns a list of models from the database based on a vector search. - :param vector_field: The field to search for the vector. :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(): @@ -302,6 +314,54 @@ def _cls(doc_id: str, data: Dict[str, Any]) -> TAsyncBareModel: 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/model.py b/firedantic/_sync/model.py index 53454dc..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 @@ -22,7 +23,10 @@ CollectionNotDefined, InvalidDocumentID, ModelNotFoundError, + VectorFieldAmbiguous, + VectorFieldNotDefined, ) +from firedantic.utils import classproperty TBareModel = TypeVar("TBareModel", bound="BareModel") TBareSubModel = TypeVar("TBareSubModel", bound="BareSubModel") @@ -69,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 @@ -247,10 +255,10 @@ def find_one( @classmethod def vector_search( cls: Type[TBareModel], - vector_field: str, 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, @@ -258,15 +266,19 @@ def vector_search( """ Returns a list of models from the database based on a vector search. - :param vector_field: The field to search for the vector. :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(): @@ -300,6 +312,54 @@ def _cls(doc_id: str, data: Dict[str, Any]) -> TBareModel: 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/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)