diff --git a/python/semantic_kernel/connectors/faiss.py b/python/semantic_kernel/connectors/faiss.py index ec9ba8e22dfa..1f4c92e70226 100644 --- a/python/semantic_kernel/connectors/faiss.py +++ b/python/semantic_kernel/connectors/faiss.py @@ -122,6 +122,7 @@ def _create_indexes(self, index: faiss.Index | None = None, indexes: dict[str, f if not index.is_trained: raise VectorStoreInitializationException("Index must be trained before using.") self.indexes[self.definition.vector_fields[0].name] = index + self.indexes_key_map.setdefault(self.definition.vector_fields[0].name, {}) return for vector_field in self.definition.vector_fields: if indexes and vector_field.name in indexes: @@ -153,6 +154,9 @@ async def ensure_collection_exists( This includes indexes that need training, or GPU-based indexes, since you would also need to build the faiss package for use with GPU's yourself. + For a new collection, supplied indexes must contain no vectors. Insert records + through this collection so the vectors, records, and key mappings stay aligned. + Args: index: The index to use, this can be used when there is only one vector field. indexes: A dictionary of indexes, the key is the name of the vector field. diff --git a/python/tests/unit/connectors/memory/test_faiss.py b/python/tests/unit/connectors/memory/test_faiss.py index 07c19f98aa10..f0f4e7375ffc 100644 --- a/python/tests/unit/connectors/memory/test_faiss.py +++ b/python/tests/unit/connectors/memory/test_faiss.py @@ -102,6 +102,48 @@ async def test_ensure_collection_exists_custom_dict(store, data_model_def): await collection.ensure_collection_deleted() +@mark.parametrize("use_index_dict", [False, True], ids=["single-index", "index-dict"]) +@mark.parametrize("index_kind", ["flat-l2", "flat-ip", "hnsw"]) +async def test_custom_index_upsert_and_search(store, data_model_def, use_index_dict, index_kind): + collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def) + records = [ + {"id": "first", "content": "first record", "vector": [1.0, 0.0, 0.0, 0.0, 0.0]}, + {"id": "second", "content": "second record", "vector": [-1.0, 0.0, 0.0, 0.0, 0.0]}, + ] + + for _ in range(2): + if index_kind == "flat-l2": + index = faiss.IndexFlatL2(5) + elif index_kind == "flat-ip": + index = faiss.IndexFlatIP(5) + else: + index = faiss.IndexHNSWFlat(5, 4) + kwargs = {"indexes": {"vector": index}} if use_index_dict else {"index": index} + await collection.ensure_collection_exists(**kwargs) + assert collection.indexes["vector"] is index + + assert await collection.upsert(records[0]) == "first" + key_map = collection.indexes_key_map["vector"] + await collection.ensure_collection_exists(**kwargs) + assert collection.indexes_key_map["vector"] is key_map + assert key_map == {"first": 0} + assert await collection.upsert(records[1:]) == ["second"] + assert key_map == {"first": 0, "second": 1} + + results = await collection.search( + vector=records[0]["vector"], top=2, include_total_count=True, include_vectors=True + ) + assert results.total_count == 2 + assert [result.record async for result in results.results] == records + assert await collection.get("first", include_vectors=True) == records[0] + assert index.ntotal == 2 + + await collection.ensure_collection_deleted() + assert collection.indexes == {} + assert collection.indexes_key_map == {} + assert collection.inner_storage == {} + + 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]}