diff --git a/python/semantic_kernel/connectors/faiss.py b/python/semantic_kernel/connectors/faiss.py index 1f4c92e70226..a3845361e1db 100644 --- a/python/semantic_kernel/connectors/faiss.py +++ b/python/semantic_kernel/connectors/faiss.py @@ -64,6 +64,18 @@ def _create_index(field: VectorStoreField) -> faiss.Index: raise VectorStoreInitializationException(f"Index with {field.index_kind} is not supported.") +def _vector_value(record: dict[str, Any], vector_field: VectorStoreField) -> Any: + """Read a vector from a serialized record under either its storage or logical name. + + The dict path serializes with the field's storage_name, while a pydantic + model_dump() and a custom to_dict() may keep the logical name. + """ + storage_name = vector_field.storage_name or vector_field.name + if storage_name in record: + return record[storage_name] + return record.get(vector_field.name) + + class FaissCollection(InMemoryCollection[TKey, TModel], Generic[TKey, TModel]): """Create a Faiss collection. @@ -168,7 +180,7 @@ async def ensure_collection_exists( async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: """Upsert records.""" for vector_field in self.definition.vector_fields: - vectors_to_add = [record.get(vector_field.storage_name or vector_field.name) for record in records] + vectors_to_add = [_vector_value(record, vector_field) for record in records] vectors = np.array(vectors_to_add, dtype=np.float32) if not self.indexes[vector_field.name].is_trained: raise VectorStoreOperationException( @@ -180,7 +192,16 @@ async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence self.indexes[vector_field.name].add(vectors) # type: ignore start = len(self.indexes_key_map[vector_field.name]) for i, record in enumerate(records): - key = record[self.definition.key_field.name] + # Records are serialized in whichever shape the data model and + # definition produce: logical names from a pydantic model_dump(), + # storage names from the dict path, or a custom to_dict(). Use the + # key name the serialization actually recorded, like InMemory does. + key_field_name = getattr(record, "key_field_name", self._key_field_storage_name) + if key_field_name not in record: + key_field_name = ( + self._key_field_storage_name if self._key_field_storage_name in record else self._key_field_name + ) + key = record[key_field_name] self.indexes_key_map[vector_field.name][key] = start + i return await super()._inner_upsert(records, **kwargs) @@ -228,9 +249,18 @@ async def _inner_search( # first we create the vector to search with np_vector = np.array(vector, dtype=np.float32).reshape(1, -1) # then do the actual vector search - distances, indexes = self.indexes[field.name].search( - np_vector, min(options.top, self.indexes[field.name].ntotal) - ) # type: ignore[call-arg] + ntotal = self.indexes[field.name].ntotal + if ntotal == 0: + # An empty (or fully deleted) index has nothing to search: faiss + # asserts k > 0. Resolve the filter and build the empty result from + # the same helper as the populated path so an invalid filter still + # raises, as it does for a populated index and for InMemory. + filtered_records = self._get_filtered_records(options) + return KernelSearchResults( + results=self._get_vector_search_results_from_results(return_list, options), + total_count=0 if options and options.include_total_count else None, + ) + distances, indexes = self.indexes[field.name].search(np_vector, min(options.top, ntotal)) # type: ignore[call-arg] # since Faiss indexes do not contain the full records, # we get the filtered records, this is a dict of the records that match the search filters # and use that to get the actual records diff --git a/python/tests/unit/connectors/memory/test_faiss.py b/python/tests/unit/connectors/memory/test_faiss.py index f0f4e7375ffc..ad28b0edacdb 100644 --- a/python/tests/unit/connectors/memory/test_faiss.py +++ b/python/tests/unit/connectors/memory/test_faiss.py @@ -1,11 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. +from typing import Annotated + import faiss -from pytest import fixture, mark, raises +from pydantic import BaseModel +from pytest import approx, fixture, mark, raises from semantic_kernel.connectors.faiss import FaissCollection, FaissStore -from semantic_kernel.data.vector import DistanceFunction, VectorStoreCollectionDefinition, VectorStoreField -from semantic_kernel.exceptions import VectorStoreInitializationException +from semantic_kernel.data.vector import ( + DistanceFunction, + VectorStoreCollectionDefinition, + VectorStoreField, + vectorstoremodel, +) +from semantic_kernel.exceptions import VectorStoreInitializationException, VectorStoreOperationException @fixture(scope="function") @@ -144,6 +152,120 @@ async def test_custom_index_upsert_and_search(store, data_model_def, use_index_d assert collection.inner_storage == {} +@fixture(scope="function") +def aliased_definition(data_model_def) -> VectorStoreCollectionDefinition: + """Only the key field declares a storage_name. + + This mirrors the InMemory test for the same behaviour and keeps the data and + vector fields on their logical names, so the record reaches _inner_upsert + under whichever serialization the record type produces. + """ + data_model_def.key_field.storage_name = "stored_id" + return data_model_def + + +@vectorstoremodel +class AliasedModel(BaseModel): + id: Annotated[str, VectorStoreField("key", storage_name="stored_id")] + content: Annotated[str, VectorStoreField("data")] = "test content" + vector: Annotated[ + list[float] | None, + VectorStoreField( + "vector", + dimensions=5, + index_kind="flat", + distance_function="dot_prod", + type="float", + ), + ] = None + + +@mark.parametrize("model_kind", ["dict", "pydantic"], ids=["dict", "pydantic"]) +async def test_upsert_with_key_storage_name(aliased_definition, model_kind): + """A key field with a storage_name must upsert whichever shape it serializes to.""" + model_type = {"dict": dict, "pydantic": AliasedModel}[model_kind] + collection = FaissCollection(record_type=model_type, definition=aliased_definition, collection_name="test-alias") + record = model_type(id="testid", content="test content", vector=[1.0, 0.0, 0.0, 0.0, 0.0]) + await collection.ensure_collection_exists() + + assert await collection.upsert(record) == "testid" + assert collection.indexes_key_map["vector"] == {"testid": 0} + assert set(collection.inner_storage) == {"testid"} + + results = await collection.search(vector=[1.0, 0.0, 0.0, 0.0, 0.0], top=1, include_vectors=True) + assert [result.record async for result in results.results] == [record] + await collection.ensure_collection_deleted() + + +@vectorstoremodel +class VectorAliasedModel(BaseModel): + id: Annotated[str, VectorStoreField("key", type="str")] + content: Annotated[str, VectorStoreField("data", type="str")] = "test content" + embedding: Annotated[ + list[float] | None, + VectorStoreField( + "vector", + storage_name="vec", + dimensions=5, + index_kind="flat", + distance_function="dot_prod", + type="float", + ), + ] = None + + +def vector_aliased_definition() -> VectorStoreCollectionDefinition: + return VectorStoreCollectionDefinition( + fields=[ + VectorStoreField("key", name="id", type="str"), + VectorStoreField("data", name="content", type="str"), + VectorStoreField( + "vector", + name="embedding", + storage_name="vec", + dimensions=5, + index_kind="flat", + distance_function="dot_prod", + type="float", + ), + ] + ) + + +async def test_upsert_with_vector_storage_name_pydantic_model(): + """A pydantic model_dump keeps the logical vector name; the vector must still resolve.""" + collection = FaissCollection( + record_type=VectorAliasedModel, definition=vector_aliased_definition(), collection_name="test-vec-alias" + ) + record = VectorAliasedModel(id="testid", embedding=[1.0, 0.0, 0.0, 0.0, 0.0]) + await collection.ensure_collection_exists() + + assert await collection.upsert(record) == "testid" + assert collection.indexes["embedding"].ntotal == 1 + # Assert the indexed vector itself, so a wrong-but-valid vector cannot pass. + assert collection.indexes["embedding"].reconstruct(0) == approx([1.0, 0.0, 0.0, 0.0, 0.0]) + + results = await collection.search(vector=[1.0, 0.0, 0.0, 0.0, 0.0], top=1, include_vectors=True) + assert [result.record async for result in results.results] == [record] + await collection.ensure_collection_deleted() + + +async def test_upsert_with_custom_to_dict_logical_names(aliased_definition): + """A custom to_dict that emits the logical key name must still resolve.""" + aliased_definition.to_dict = lambda record, **kwargs: { + "id": record["id"], + "content": record["content"], + "vector": record["vector"], + } + collection = FaissCollection(record_type=dict, definition=aliased_definition, collection_name="test-alias-logical") + record = {"id": "testid", "content": "test content", "vector": [1.0, 0.0, 0.0, 0.0, 0.0]} + await collection.ensure_collection_exists() + + assert await collection.upsert(record) == "testid" + assert collection.indexes_key_map["vector"] == {"testid": 0} + await collection.ensure_collection_deleted() + + async def test_upsert(faiss_collection): await faiss_collection.ensure_collection_exists() record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]} @@ -153,6 +275,37 @@ async def test_upsert(faiss_collection): await faiss_collection.ensure_collection_deleted() +async def test_search_on_empty_collection_returns_no_results(faiss_collection): + """An empty index has nothing to search; faiss would otherwise assert k > 0.""" + await faiss_collection.ensure_collection_exists() + + results = await faiss_collection.search(vector=[0.1, 0.2, 0.3, 0.4, 0.5], top=3, include_total_count=True) + assert [result async for result in results.results] == [] + assert results.total_count == 0 + await faiss_collection.ensure_collection_deleted() + + +async def test_search_on_empty_collection_still_validates_the_filter(faiss_collection): + """An empty index returns early, but an invalid filter must still raise.""" + await faiss_collection.ensure_collection_exists() + + with raises(VectorStoreOperationException): + await faiss_collection.search(vector=[0.1, 0.2, 0.3, 0.4, 0.5], top=3, filter="not valid (") + await faiss_collection.ensure_collection_deleted() + + +async def test_search_after_deleting_every_record_returns_no_results(faiss_collection): + """Deleting the last record empties the index; a later search must not raise.""" + await faiss_collection.ensure_collection_exists() + record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]} + await faiss_collection.upsert(record) + await faiss_collection.delete("testid") + + results = await faiss_collection.search(vector=record["vector"], top=3) + assert [result async for result in results.results] == [] + await faiss_collection.ensure_collection_deleted() + + async def test_get(faiss_collection): await faiss_collection.ensure_collection_exists() record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}