# Copyright (c) Microsoft. All rights reserved. import faiss from pytest import 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 @fixture(scope="function") def data_model_def() -> VectorStoreCollectionDefinition: return VectorStoreCollectionDefinition( fields=[ VectorStoreField("key", name="id"), VectorStoreField("data", name="content"), VectorStoreField( "vector", name="vector", dimensions=5, index_kind="flat", distance_function="dot_prod", type="float", ), ] ) @fixture(scope="function") def store() -> FaissStore: return FaissStore() @fixture(scope="function") def faiss_collection(data_model_def): return FaissCollection(record_type=dict, definition=data_model_def, collection_name="test") async def test_store_get_collection(store, data_model_def): collection = store.get_collection(dict, definition=data_model_def, collection_name="test") assert collection.collection_name == "test" assert collection.record_type is dict assert collection.definition == data_model_def assert collection.inner_storage == {} @mark.parametrize( "dist", [ DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE, DistanceFunction.DOT_PROD, ], ) async def test_ensure_collection_exists(store, data_model_def, dist): for field in data_model_def.fields: if field.name == "vector": field.distance_function = dist collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) await collection.ensure_collection_exists() assert collection.inner_storage == {} assert collection.indexes assert collection.indexes["vector"] is not None async def test_ensure_collection_exists_incompatible_dist(store, data_model_def): for field in data_model_def.fields: if field.name == "vector": field.distance_function = "cosine_distance" collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) with raises(VectorStoreInitializationException): await collection.ensure_collection_exists() async def test_ensure_collection_exists_custom(store, data_model_def): index = faiss.IndexFlat(5) collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) await collection.ensure_collection_exists(index=index) assert collection.inner_storage == {} assert collection.indexes assert collection.indexes["vector"] is not None assert collection.indexes["vector"] == index assert collection.indexes["vector"].is_trained is True await collection.ensure_collection_deleted() async def test_ensure_collection_exists_custom_untrained(store, data_model_def): index = faiss.IndexIVFFlat(faiss.IndexFlat(5), 5, 10) collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) with raises(VectorStoreInitializationException): await collection.ensure_collection_exists(index=index) del index async def test_ensure_collection_exists_custom_dict(store, data_model_def): index = faiss.IndexFlat(5) collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) await collection.ensure_collection_exists(indexes={"vector": index}) assert collection.inner_storage == {} assert collection.indexes assert collection.indexes["vector"] is not None assert collection.indexes["vector"] == index 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]} key = await faiss_collection.upsert(record) assert key == "testid" assert faiss_collection.inner_storage == {"testid": record} 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]} await faiss_collection.upsert(record) result = await faiss_collection.get("testid") assert result["id"] == record["id"] assert result["content"] == record["content"] await faiss_collection.ensure_collection_deleted() async def test_get_missing(faiss_collection): await faiss_collection.ensure_collection_exists() result = await faiss_collection.get("testid") assert result is None await faiss_collection.ensure_collection_deleted() async def test_delete(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]} await faiss_collection.upsert(record) await faiss_collection.delete("testid") assert faiss_collection.inner_storage == {} await faiss_collection.ensure_collection_deleted() async def test_collection_exists(faiss_collection): assert await faiss_collection.collection_exists() is False await faiss_collection.ensure_collection_exists() assert await faiss_collection.collection_exists() is True await faiss_collection.ensure_collection_deleted() async def test_ensure_collection_deleted(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]} await faiss_collection.upsert(record) assert faiss_collection.inner_storage == {"testid": record} await faiss_collection.ensure_collection_deleted() assert faiss_collection.inner_storage == {} @mark.parametrize("dist", [DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE, DistanceFunction.DOT_PROD]) async def test_ensure_collection_exists_and_search(faiss_collection, dist): for field in faiss_collection.definition.fields: if field.name == "vector": field.distance_function = dist await faiss_collection.ensure_collection_exists() record1 = {"id": "testid1", "content": "test content", "vector": [1.0, 1.0, 1.0, 1.0, 1.0]} record2 = {"id": "testid2", "content": "test content", "vector": [-1.0, -1.0, -1.0, -1.0, -1.0]} await faiss_collection.upsert([record1, record2]) results = await faiss_collection.search( vector=[0.9, 0.9, 0.9, 0.9, 0.9], vector_property_name="vector", include_total_count=True, include_vectors=True, ) assert results.total_count == 2 idx = 0 async for res in results.results: assert res.record == record1 if idx == 0 else record2 idx += 1 await faiss_collection.ensure_collection_deleted()