Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 35 additions & 5 deletions python/semantic_kernel/connectors/faiss.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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(
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down
159 changes: 156 additions & 3 deletions python/tests/unit/connectors/memory/test_faiss.py
Original file line number Diff line number Diff line change
@@ -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")
Expand Down Expand Up @@ -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]}
Expand All @@ -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]}
Expand Down
Loading