diff --git a/python/.cspell.json b/python/.cspell.json index 485789ae22a1..9fd255aed239 100644 --- a/python/.cspell.json +++ b/python/.cspell.json @@ -36,6 +36,7 @@ "dotenv", "endregion", "entra", + "faiss", "genai", "generativeai", "hnsw", diff --git a/python/pyproject.toml b/python/pyproject.toml index 34d5f03b7fb7..cf5aebf7fc59 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -57,7 +57,7 @@ autogen = [ ] azure = [ "azure-ai-inference >= 1.0.0b6", - "azure-ai-projects >= 1.0.0b5", + "azure-ai-projects == 1.0.0b5", "azure-core-tracing-opentelemetry >= 1.0.0b11", "azure-search-documents >= 11.6.0b4", "azure-identity ~= 1.13", @@ -70,6 +70,9 @@ google = [ "google-cloud-aiplatform == 1.82.0", "google-generativeai ~= 0.8" ] +faiss = [ + "faiss-cpu>=1.10.0" +] hugging_face = [ "transformers[torch] ~= 4.28", "sentence-transformers >= 2.2,< 4.0", diff --git a/python/samples/concepts/memory/complex_memory.py b/python/samples/concepts/memory/complex_memory.py index 423508388e9c..b13b60b075a8 100644 --- a/python/samples/concepts/memory/complex_memory.py +++ b/python/samples/concepts/memory/complex_memory.py @@ -23,6 +23,7 @@ AzureCosmosDBNoSQLCollection, ) from semantic_kernel.connectors.memory.chroma import ChromaCollection +from semantic_kernel.connectors.memory.faiss import FaissCollection from semantic_kernel.connectors.memory.in_memory import InMemoryVectorCollection from semantic_kernel.connectors.memory.postgres import PostgresCollection from semantic_kernel.connectors.memory.qdrant import QdrantCollection @@ -121,7 +122,7 @@ class DataModelList: # since not all combinations are supported by all databases. # The values below might need to be changed for your collection to work. distance_function = DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE -index_kind = IndexKind.HNSW +index_kind = IndexKind.FLAT DataModel = get_data_model("array", index_kind, distance_function) # A list of VectorStoreRecordCollection that can be used. @@ -190,6 +191,10 @@ class DataModelList: collection_name=collection_name, ), "chroma": lambda: ChromaCollection(data_model_type=DataModel, collection_name=collection_name), + "faiss": lambda: FaissCollection[str, DataModel]( + collection_name=collection_name, + data_model_type=DataModel, + ), } @@ -202,7 +207,6 @@ async def main(collection: str, use_azure_openai: bool): kernel.add_service(embedder) async with collections[collection]() as record_collection: print_with_color(f"Creating {collection} collection!", Colors.CGREY) - await record_collection.delete_collection() await record_collection.create_collection_if_not_exists() record1 = DataModel( diff --git a/python/samples/concepts/memory/simple_memory.py b/python/samples/concepts/memory/simple_memory.py index 941b5f59baa7..a66ab5153fec 100644 --- a/python/samples/concepts/memory/simple_memory.py +++ b/python/samples/concepts/memory/simple_memory.py @@ -104,7 +104,7 @@ async def main(): # we also use the async with to open and close the connection # for the in memory collection, this is just a no-op # but for other collections, like Azure AI Search, this will open and close the connection - async with InMemoryVectorCollection[DataModel]( + async with InMemoryVectorCollection[str, DataModel]( collection_name="test", data_model_type=DataModel, ) as record_collection: diff --git a/python/semantic_kernel/connectors/memory/faiss.py b/python/semantic_kernel/connectors/memory/faiss.py new file mode 100644 index 000000000000..d21161adeb13 --- /dev/null +++ b/python/semantic_kernel/connectors/memory/faiss.py @@ -0,0 +1,248 @@ +# Copyright (c) Microsoft. All rights reserved. +import logging +import sys +from collections.abc import MutableMapping, Sequence +from typing import TYPE_CHECKING, Any, Generic + +if sys.version_info >= (3, 12): + from typing import override # pragma: no cover +else: + from typing_extensions import override # pragma: no cover + +import faiss +import numpy as np +from pydantic import Field + +from semantic_kernel.connectors.memory.in_memory.in_memory_collection import ( + IN_MEMORY_SCORE_KEY, + InMemoryVectorCollection, + TKey, + TModel, +) +from semantic_kernel.data.const import DistanceFunction, IndexKind +from semantic_kernel.data.kernel_search_results import KernelSearchResults +from semantic_kernel.data.record_definition.vector_store_model_definition import VectorStoreRecordDefinition +from semantic_kernel.data.record_definition.vector_store_record_fields import VectorStoreRecordVectorField +from semantic_kernel.data.vector_search.vector_search_options import VectorSearchOptions +from semantic_kernel.data.vector_search.vector_search_result import VectorSearchResult +from semantic_kernel.data.vector_storage.vector_store import VectorStore +from semantic_kernel.exceptions import ( + VectorStoreInitializationException, + VectorStoreOperationException, +) + +if TYPE_CHECKING: + from semantic_kernel.data.vector_storage.vector_store_record_collection import VectorStoreRecordCollection + +logger = logging.getLogger(__name__) + + +class FaissCollection(InMemoryVectorCollection[TKey, TModel], Generic[TKey, TModel]): + """Create a Faiss collection. + + The Faiss Collection builds on the InMemoryVectorCollection, + it maintains indexes and mappings for each vector field. + """ + + indexes: MutableMapping[str, faiss.Index] = Field(default_factory=dict) + indexes_key_map: MutableMapping[str, MutableMapping[TKey, int]] = Field(default_factory=dict) + + def __init__( + self, + collection_name: str, + data_model_type: type[TModel], + data_model_definition: VectorStoreRecordDefinition | None = None, + **kwargs: Any, + ): + """Create a Faiss Collection. + + To allow more complex index setups, you can pass them in here: + ```python + import faiss + + index = faiss.IndexFlatL2(128) + FaissCollection(..., indexes={"vector_field_name": index}) + ``` + + or you can manually add them to the indexes field of the collection. + + Args: + collection_name: The name of the collection. + data_model_type: The type of the data model. + data_model_definition: The definition of the data model. + kwargs: Additional arguments. + """ + super().__init__( + data_model_type=data_model_type, + data_model_definition=data_model_definition, + collection_name=collection_name, + **kwargs, + ) + + def _create_indexes(self, index: faiss.Index | None = None, indexes: dict[str, faiss.Index] | None = None) -> None: + """Create Faiss indexes for each vector field. + + 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. + """ + if len(self.data_model_definition.vector_fields) == 1 and index is not None: + if not isinstance(index, faiss.Index): + raise VectorStoreInitializationException("Index must be a subtype of faiss.Index") + if not index.is_trained: + raise VectorStoreInitializationException("Index must be trained before using.") + self.indexes[self.data_model_definition.vector_fields[0].name] = index + return + for vector_field in self.data_model_definition.vector_fields: + if indexes and vector_field.name in indexes: + if not isinstance(indexes[vector_field.name], faiss.Index): + raise VectorStoreInitializationException( + f"Index for {vector_field.name} must be a subtype of faiss.Index" + ) + if not indexes[vector_field.name].is_trained: + raise VectorStoreInitializationException( + f"Index for {vector_field.name} must be trained before using." + ) + self.indexes[vector_field.name] = indexes[vector_field.name] + if vector_field.name not in self.indexes_key_map: + self.indexes_key_map.setdefault(vector_field.name, {}) + continue + if vector_field.name not in self.indexes: + index = self._create_index(vector_field) + self.indexes[vector_field.name] = index + if vector_field.name not in self.indexes_key_map: + self.indexes_key_map.setdefault(vector_field.name, {}) + + def _create_index(self, field: VectorStoreRecordVectorField) -> faiss.Index: + """Create a Faiss index.""" + index_kind = field.index_kind or IndexKind.FLAT + distance_function = field.distance_function or DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE + match index_kind: + case IndexKind.FLAT: + match distance_function: + case DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE: + return faiss.IndexFlatL2(field.dimensions) + case DistanceFunction.DOT_PROD: + return faiss.IndexFlatIP(field.dimensions) + case _: + raise VectorStoreInitializationException( + f"Distance function {distance_function} is not supported for index kind {index_kind}." + ) + case _: + raise VectorStoreInitializationException(f"Index with {index_kind} is not supported.") + + @override + async def create_collection( + self, index: faiss.Index | None = None, indexes: dict[str, faiss.Index] | None = None, **kwargs: Any + ) -> None: + """Create a collection. + + Considering the complexity of different faiss indexes, we support a limited set. + For more advanced scenario's you can create your own indexes and pass them in here. + 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. + + 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. + kwargs: Additional arguments. + """ + self._create_indexes(index=index, indexes=indexes) + + @override + async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: + """Upsert records.""" + for vector_field in self.data_model_definition.vector_field_names: + vectors_to_add = [record.get(vector_field) for record in records] + vectors = np.array(vectors_to_add, dtype=np.float32) + if not self.indexes[vector_field].is_trained: + raise VectorStoreOperationException( + f"This index (of type {type(self.indexes[vector_field])}) requires training, " + "which is not supported. To train the index, " + f"use .indexes[{vector_field}].train, " + "see faiss docs for more details." + ) + self.indexes[vector_field].add(vectors) + start = len(self.indexes_key_map[vector_field]) + for i, record in enumerate(records): + key = record[self.data_model_definition.key_field.name] + self.indexes_key_map[vector_field][key] = start + i + return await super()._inner_upsert(records, **kwargs) + + @override + async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None: + for key in keys: + for vector_field in self.data_model_definition.vector_field_names: + if key in self.indexes_key_map[vector_field]: + vector_index = self.indexes_key_map[vector_field][key] + self.indexes[vector_field].remove_ids(np.array([vector_index])) + self.indexes_key_map[vector_field].pop(key, None) + await super()._inner_delete(keys, **kwargs) + + @override + async def delete_collection(self, **kwargs: Any) -> None: + for vector_field in self.data_model_definition.vector_field_names: + if vector_field in self.indexes: + del self.indexes[vector_field] + if vector_field in self.indexes_key_map: + del self.indexes_key_map[vector_field] + await super().delete_collection(**kwargs) + + @override + async def does_collection_exist(self, **kwargs: Any) -> bool: + return bool(self.indexes) + + async def _inner_search_vectorized( + self, + vector: list[float | int], + options: VectorSearchOptions, + **kwargs: Any, + ) -> KernelSearchResults[VectorSearchResult[TModel]]: + field = options.vector_field_name or self.data_model_definition.vector_field_names[0] + assert isinstance(self.data_model_definition.fields.get(field), VectorStoreRecordVectorField) # nosec + if vector and field: + return_list = [] + # since the vector index works independently of the record index, + # we will need to get all records that adhere to the filter first + filtered_records = self._get_filtered_records(options) + np_vector = np.array(vector, dtype=np.float32).reshape(1, -1) + # then do the actual vector search + distances, indexes = self.indexes[field].search(np_vector, min(options.top, self.indexes[field].ntotal)) + # we then iterate through the results, the order is the order of relevance + # (less or most distance, dependant on distance metric used) + for i, index in enumerate(indexes[0]): + key = list(self.indexes_key_map[field].keys())[index] + # if the key is not in the filtered records, we ignore it + if key not in filtered_records: + continue + filtered_records[key][IN_MEMORY_SCORE_KEY] = distances[0][i] + # so we return the list in the order of the search, with the record from the inner_storage. + return_list.append(filtered_records[key]) + return KernelSearchResults( + results=self._get_vector_search_results_from_results(return_list, options), + total_count=len(return_list) if options and options.include_total_count else None, + ) + + +class FaissStore(VectorStore): + """Create a Faiss store.""" + + @override + async def list_collection_names(self, **kwargs) -> Sequence[str]: + return list(self.vector_record_collections.keys()) + + @override + def get_collection( + self, + collection_name: str, + data_model_type: type[object], + data_model_definition=None, + **kwargs, + ) -> "VectorStoreRecordCollection": + self.vector_record_collections[collection_name] = FaissCollection( + collection_name=collection_name, + data_model_type=data_model_type, + data_model_definition=data_model_definition, + **kwargs, + ) + return self.vector_record_collections[collection_name] diff --git a/python/semantic_kernel/connectors/memory/in_memory/in_memory_collection.py b/python/semantic_kernel/connectors/memory/in_memory/in_memory_collection.py index 95e2e537f9ee..d00774ce9509 100644 --- a/python/semantic_kernel/connectors/memory/in_memory/in_memory_collection.py +++ b/python/semantic_kernel/connectors/memory/in_memory/in_memory_collection.py @@ -2,7 +2,7 @@ import sys from collections.abc import AsyncIterable, Callable, Mapping, Sequence -from typing import Any, ClassVar, TypeVar +from typing import Any, ClassVar, Generic, TypeVar if sys.version_info >= (3, 12): from typing import override # pragma: no cover @@ -30,7 +30,7 @@ from semantic_kernel.kernel_types import OneOrMany from semantic_kernel.utils.list_handler import empty_generator -KEY_TYPES = str | int | float +TKey = TypeVar("TKey", bound=str | int | float) TModel = TypeVar("TModel") @@ -38,11 +38,14 @@ class InMemoryVectorCollection( - VectorSearchBase[KEY_TYPES, TModel], VectorTextSearchMixin[TModel], VectorizedSearchMixin[TModel] + VectorSearchBase[TKey, TModel], + VectorTextSearchMixin[TModel], + VectorizedSearchMixin[TModel], + Generic[TKey, TModel], ): """In Memory Collection.""" - inner_storage: dict[KEY_TYPES, dict] = Field(default_factory=dict) + inner_storage: dict[TKey, dict] = Field(default_factory=dict) supported_key_types: ClassVar[list[str] | None] = ["str", "int", "float"] def __init__( @@ -50,12 +53,14 @@ def __init__( collection_name: str, data_model_type: type[TModel], data_model_definition: VectorStoreRecordDefinition | None = None, + **kwargs: Any, ): """Create a In Memory Collection.""" super().__init__( data_model_type=data_model_type, data_model_definition=data_model_definition, collection_name=collection_name, + **kwargs, ) def _validate_data_model(self): @@ -65,16 +70,16 @@ def _validate_data_model(self): raise VectorStoreModelValidationError(f"Field name '{IN_MEMORY_SCORE_KEY}' is reserved for internal use.") @override - async def _inner_delete(self, keys: Sequence[KEY_TYPES], **kwargs: Any) -> None: + async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None: for key in keys: self.inner_storage.pop(key, None) @override - async def _inner_get(self, keys: Sequence[KEY_TYPES], **kwargs: Any) -> Any | OneOrMany[TModel] | None: + async def _inner_get(self, keys: Sequence[TKey], **kwargs: Any) -> Any | OneOrMany[TModel] | None: return [self.inner_storage[key] for key in keys if key in self.inner_storage] @override - async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[KEY_TYPES]: + async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: updated_keys = [] for record in records: key = record[self._key_field_name] if isinstance(record, Mapping) else getattr(record, self._key_field_name) @@ -125,7 +130,7 @@ async def _inner_search_text( **kwargs: Any, ) -> KernelSearchResults[VectorSearchResult[TModel]]: """Inner search method.""" - return_records: dict[KEY_TYPES, float] = {} + return_records: dict[TKey, float] = {} for key, record in self._get_filtered_records(options).items(): if self._should_add_text_search(search_text, record): return_records[key] = 1.0 @@ -144,10 +149,8 @@ async def _inner_search_vectorized( options: VectorSearchOptions, **kwargs: Any, ) -> KernelSearchResults[VectorSearchResult[TModel]]: - return_records: dict[KEY_TYPES, float] = {} - if not options.vector_field_name: - raise ValueError("Vector field name must be provided in options for vector search.") - field = options.vector_field_name + return_records: dict[TKey, float] = {} + field = options.vector_field_name or self.data_model_definition.vector_field_names[0] assert isinstance(self.data_model_definition.fields.get(field), VectorStoreRecordVectorField) # nosec distance_metric = ( self.data_model_definition.fields.get(field).distance_function # type: ignore @@ -180,7 +183,7 @@ async def _inner_search_vectorized( return KernelSearchResults(results=empty_generator()) async def _generate_return_list( - self, return_records: dict[KEY_TYPES, float], options: VectorSearchOptions | None + self, return_records: dict[TKey, float], options: VectorSearchOptions | None ) -> AsyncIterable[dict]: top = 3 if not options else options.top skip = 0 if not options else options.skip @@ -194,7 +197,7 @@ async def _generate_return_list( if returned >= top: break - def _get_filtered_records(self, options: VectorSearchOptions | None) -> dict[KEY_TYPES, dict]: + def _get_filtered_records(self, options: VectorSearchOptions | None) -> dict[TKey, dict]: if options and options.filter: for filter in options.filter.filters: return {key: record for key, record in self.inner_storage.items() if self._apply_filter(record, filter)} diff --git a/python/semantic_kernel/data/vector_storage/vector_store.py b/python/semantic_kernel/data/vector_storage/vector_store.py index d0e24a0bf5da..3fba08ee2bd2 100644 --- a/python/semantic_kernel/data/vector_storage/vector_store.py +++ b/python/semantic_kernel/data/vector_storage/vector_store.py @@ -27,7 +27,7 @@ def get_collection( data_model_type: type[object], data_model_definition: VectorStoreRecordDefinition | None = None, **kwargs: Any, - ) -> VectorStoreRecordCollection: + ) -> "VectorStoreRecordCollection": """Get a vector record store.""" ... # pragma: no cover diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 60bb1bda97da..9a3ec33d3697 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -42,7 +42,7 @@ def pytest_configure(config): @fixture(scope="function") def kernel() -> "Kernel": - from semantic_kernel.kernel import Kernel + from semantic_kernel import Kernel return Kernel() @@ -590,7 +590,7 @@ class MyDataModel: @fixture def data_model_definition( index_kind: str, distance_function: str, vector_property_type: str, dimensions: int -) -> object: +) -> VectorStoreRecordDefinition: return VectorStoreRecordDefinition( fields={ "id": VectorStoreRecordKeyField(), diff --git a/python/tests/unit/connectors/memory/faiss/test_faiss.py b/python/tests/unit/connectors/memory/faiss/test_faiss.py new file mode 100644 index 000000000000..17941c7b6d62 --- /dev/null +++ b/python/tests/unit/connectors/memory/faiss/test_faiss.py @@ -0,0 +1,206 @@ +# Copyright (c) Microsoft. All rights reserved. + +import faiss +from pytest import fixture, mark, raises + +from semantic_kernel.connectors.memory.faiss import FaissCollection, FaissStore +from semantic_kernel.data import ( + DistanceFunction, + VectorSearchFilter, + VectorSearchOptions, + VectorStoreRecordDataField, + VectorStoreRecordDefinition, + VectorStoreRecordKeyField, + VectorStoreRecordVectorField, +) +from semantic_kernel.exceptions import VectorStoreInitializationException + + +@fixture(scope="function") +def data_model_def() -> VectorStoreRecordDefinition: + return VectorStoreRecordDefinition( + fields={ + "id": VectorStoreRecordKeyField(), + "content": VectorStoreRecordDataField( + has_embedding=True, + embedding_property_name="vector", + ), + "vector": VectorStoreRecordVectorField( + dimensions=5, + index_kind="flat", + distance_function="dot_prod", + property_type="float", + ), + } + ) + + +@fixture(scope="function") +def faiss_collection(data_model_def): + return FaissCollection("test", dict, data_model_def) + + +def test_store_init(): + store = FaissStore() + assert store.vector_record_collections == {} + + +async def test_store_get_collection(data_model_def): + store = FaissStore() + collection = store.get_collection("test", dict, data_model_def) + assert collection.collection_name == "test" + assert collection.data_model_type is dict + assert collection.data_model_definition == data_model_def + assert collection.inner_storage == {} + assert (await store.list_collection_names()) == ["test"] + + +@mark.parametrize( + "dist", + [ + DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE, + DistanceFunction.DOT_PROD, + ], +) +async def test_create_collection(data_model_def, dist): + store = FaissStore() + data_model_def.fields["vector"].distance_function = dist + collection = store.get_collection("test", dict, data_model_def) + await collection.create_collection() + assert collection.inner_storage == {} + assert collection.indexes + assert collection.indexes["vector"] is not None + + +async def test_create_collection_incompatible_dist(data_model_def): + store = FaissStore() + data_model_def.fields["vector"].distance_function = "cosine_distance" + collection = store.get_collection("test", dict, data_model_def) + with raises(VectorStoreInitializationException): + await collection.create_collection() + + +async def test_create_collection_custom(data_model_def): + index = faiss.IndexFlat(5) + store = FaissStore() + collection = store.get_collection("test", dict, data_model_def) + await collection.create_collection(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.delete_collection() + + +async def test_create_collection_custom_untrained(data_model_def): + index = faiss.IndexIVFFlat(faiss.IndexFlat(5), 5, 10) + store = FaissStore() + collection = store.get_collection("test", dict, data_model_def) + with raises(VectorStoreInitializationException): + await collection.create_collection(index=index) + del index + + +async def test_create_collection_custom_dict(data_model_def): + index = faiss.IndexFlat(5) + store = FaissStore() + collection = store.get_collection("test", dict, data_model_def) + await collection.create_collection(indexes={"vector": index}) + assert collection.inner_storage == {} + assert collection.indexes + assert collection.indexes["vector"] is not None + assert collection.indexes["vector"] == index + await collection.delete_collection() + + +async def test_upsert(faiss_collection): + await faiss_collection.create_collection() + 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.delete_collection() + + +async def test_get(faiss_collection): + await faiss_collection.create_collection() + 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 == record + await faiss_collection.delete_collection() + + +async def test_get_missing(faiss_collection): + await faiss_collection.create_collection() + result = await faiss_collection.get("testid") + assert result is None + await faiss_collection.delete_collection() + + +async def test_delete(faiss_collection): + await faiss_collection.create_collection() + 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.delete_collection() + + +async def test_does_collection_exist(faiss_collection): + assert await faiss_collection.does_collection_exist() is False + await faiss_collection.create_collection() + assert await faiss_collection.does_collection_exist() is True + await faiss_collection.delete_collection() + + +async def test_delete_collection(faiss_collection): + await faiss_collection.create_collection() + 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.delete_collection() + assert faiss_collection.inner_storage == {} + + +async def test_text_search(faiss_collection): + await faiss_collection.create_collection() + record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]} + await faiss_collection.upsert(record) + results = await faiss_collection.text_search(search_text="content") + assert len([res async for res in results.results]) == 1 + await faiss_collection.delete_collection() + + +async def test_text_search_with_filter(faiss_collection): + await faiss_collection.create_collection() + record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]} + await faiss_collection.upsert(record) + results = await faiss_collection.text_search( + search_text="content", + options=VectorSearchOptions( + filter=VectorSearchFilter.any_tag_equal_to("vector", 0.1).equal_to("content", "content") + ), + ) + assert len([res async for res in results.results]) == 1 + await faiss_collection.delete_collection() + + +@mark.parametrize("dist", [DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE, DistanceFunction.DOT_PROD]) +async def test_create_collection_and_search(faiss_collection, dist): + faiss_collection.data_model_definition.fields["vector"].distance_function = dist + await faiss_collection.create_collection() + 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_batch([record1, record2]) + results = await faiss_collection.vectorized_search( + vector=[0.9, 0.9, 0.9, 0.9, 0.9], + options=VectorSearchOptions(vector_field_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.delete_collection() diff --git a/python/uv.lock b/python/uv.lock index 9ca0655aa040..583ccf6a6af4 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -343,30 +343,30 @@ wheels = [ [[package]] name = "azure-ai-inference" -version = "1.0.0b9" +version = "1.0.0b6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4e/6a/ed85592e5c64e08c291992f58b1a94dab6869f28fb0f40fd753dced73ba6/azure_ai_inference-1.0.0b9.tar.gz", hash = "sha256:1feb496bd84b01ee2691befc04358fa25d7c344d8288e99364438859ad7cd5a4", size = 182408 } +sdist = { url = "https://files.pythonhosted.org/packages/2a/c9/264ae0ef0460dbd7c7efe1d3a093ad6a00fb2823d341ac457459396df2d6/azure_ai_inference-1.0.0b6.tar.gz", hash = "sha256:b8ac941de1e69151bad464191e18856d4e74f962ae03235da137a9a326143676", size = 145414 } wheels = [ - { url = "https://files.pythonhosted.org/packages/4f/0f/27520da74769db6e58327d96c98e7b9a07ce686dff582c9a5ec60b03f9dd/azure_ai_inference-1.0.0b9-py3-none-any.whl", hash = "sha256:49823732e674092dad83bb8b0d1b65aa73111fab924d61349eb2a8cdc0493990", size = 124885 }, + { url = "https://files.pythonhosted.org/packages/5a/aa/47459ab2e67c55ff98dbb9694c47cf98e484ce1ae1acb244d28b25a8c1c1/azure_ai_inference-1.0.0b6-py3-none-any.whl", hash = "sha256:5699ad78d70ec2d227a5eff2c1bafc845018f6624edc5b03589dfff861c54958", size = 115312 }, ] [[package]] name = "azure-ai-projects" -version = "1.0.0b6" +version = "1.0.0b5" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1f/72/5f9a78c913af66c55222ff912227b494707c4adfbdca27a78c3687a1b8ba/azure_ai_projects-1.0.0b6.tar.gz", hash = "sha256:ce6cfb2403eeb1a80e5dd84193fb2864953cd95a351f3d4572a5451bbb4c30d2", size = 298737 } +sdist = { url = "https://files.pythonhosted.org/packages/98/0f/47600c1c3bb1b92913350c9878c517ccbd2630f4c5393cd4054dc71d5fbd/azure_ai_projects-1.0.0b5.tar.gz", hash = "sha256:7bb068b38a8810c5a59a628e18f196c76af2a06cb774e725b97bfc8b61b3aaf0", size = 303846 } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/d9/14b31fc773072b63493d55a1a5b60e656f11aeea2b603fef2eb567686d96/azure_ai_projects-1.0.0b6-py3-none-any.whl", hash = "sha256:b0689825065648b54b4405e9edd22b1de3ea0dfc1ca3baf99db5173fd6208692", size = 187221 }, + { url = "https://files.pythonhosted.org/packages/2f/3b/5ef0717685e4d3e3b537d7538e430199ab4c16665766ab3de90d5fedc8ef/azure_ai_projects-1.0.0b5-py3-none-any.whl", hash = "sha256:7ed7a44e8d76bf108be7de6f33b9c54a63fc47193cd259770006a6862ce25ccd", size = 193481 }, ] [[package]] @@ -1195,6 +1195,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/8f/c4d9bafc34ad7ad5d8dc16dd1347ee0e507a52c3adb6bfa8887e1c6a26ba/executing-2.2.0-py2.py3-none-any.whl", hash = "sha256:11387150cad388d62750327a53d3339fad4888b39a6fe233c3afbb54ecffd3aa", size = 26702 }, ] +[[package]] +name = "faiss-cpu" +version = "1.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "packaging", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/1b/6fe5dbe5be0240cfd82b52bd7c186655c578d935c0ce2e713c100e6f8cce/faiss_cpu-1.10.0.tar.gz", hash = "sha256:5bdca555f24bc036f4d67f8a5a4d6cc91b8d2126d4e78de496ca23ccd46e479d", size = 69159 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/56/87eb506d8634f08fc7c63d1ca5631aeec7d6b9afbfabedf2cb7a2a804b13/faiss_cpu-1.10.0-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:6693474be296a7142ade1051ea18e7d85cedbfdee4b7eac9c52f83fed0467855", size = 7693034 }, + { url = "https://files.pythonhosted.org/packages/51/46/f4d9de34ed1b06300b1a75b824d4857963216f5826de33f291af78088e39/faiss_cpu-1.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:70ebe60a560414dc8dd6cfe8fed105c8f002c0d11f765f5adfe8d63d42c0467f", size = 3234656 }, + { url = "https://files.pythonhosted.org/packages/74/3a/e146861019d9290e0198b3470b8d13a658c3b5f228abefc3658ce0afd63d/faiss_cpu-1.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:74c5712d4890f15c661ab7b1b75867812e9596e1469759956fad900999bedbb5", size = 3663789 }, + { url = "https://files.pythonhosted.org/packages/aa/40/624f0002bb777e37aac1aadfadec1eb4391be6ad05b7fcfbf66049b99a48/faiss_cpu-1.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:473d158fbd638d6ad5fb64469ba79a9f09d3494b5f4e8dfb4f40ce2fc335dca4", size = 30673545 }, + { url = "https://files.pythonhosted.org/packages/d6/39/298ffcbefd899e84a43e63df217a6dc800d52bca37ebe0d1155ff367886a/faiss_cpu-1.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:dcd0cb2ec84698cbe3df9ed247d2392f09bda041ad34b92d38fa916cd019ad4b", size = 13684176 }, + { url = "https://files.pythonhosted.org/packages/78/93/81800f41cb2c719c199d3eb534fcc154853123261d841e37482e8e468619/faiss_cpu-1.10.0-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:8ff6924b0f00df278afe70940ae86302066466580724c2f3238860039e9946f1", size = 7693037 }, + { url = "https://files.pythonhosted.org/packages/8d/83/fc9028f6d6aec2c2f219f53a5d4a2b279434715643242e59a2e9755b1ce0/faiss_cpu-1.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:cb80b530a9ded44a7d4031a7355a237aaa0ff1f150c1176df050e0254ea5f6f6", size = 3234657 }, + { url = "https://files.pythonhosted.org/packages/af/45/588a02e60daa73f6052611334fbbdffcedf37122320f1c91cb90f3e69b96/faiss_cpu-1.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:7a9fef4039ed877d40e41d5563417b154c7f8cd57621487dad13c4eb4f32515f", size = 3663710 }, + { url = "https://files.pythonhosted.org/packages/cb/cf/9caa08ca4e21ab935f82be0713e5d60566140414c3fff7932d9427c8fd72/faiss_cpu-1.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:49b6647aa9e159a2c4603cbff2e1b313becd98ad6e851737ab325c74fe8e0278", size = 30673629 }, + { url = "https://files.pythonhosted.org/packages/2c/2d/d2a4171a9cca9a7c04cd9d6f9441a37f1e0558724b90bf7fc7db08553601/faiss_cpu-1.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:6f8c0ef8b615c12c7bf612bd1fc51cffa49c1ddaa6207c6981f01ab6782e6b3b", size = 13683966 }, + { url = "https://files.pythonhosted.org/packages/bd/cc/f6aa1288dbb40b2a4f101d16900885e056541f37d8d08ec70462e92cf277/faiss_cpu-1.10.0-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:2aca486fe2d680ea64a18d356206c91ff85db99fd34c19a757298c67c23262b1", size = 7720242 }, + { url = "https://files.pythonhosted.org/packages/be/56/40901306324a17fbc1eee8a6e86ba67bd99a67e768ce9908f271e648e9e0/faiss_cpu-1.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c1108a4059c66c37c403183e566ca1ed0974a6af7557c92d49207639aab661bc", size = 3239223 }, + { url = "https://files.pythonhosted.org/packages/2e/34/5b1463c450c9a6de3109caf8f38fbf0c329ef940ed1973fcf8c8ec7fa27e/faiss_cpu-1.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:449f3eb778d6d937e01a16a3170de4bb8aabfe87c7cb479b458fb790276310c5", size = 3671461 }, + { url = "https://files.pythonhosted.org/packages/78/d9/0b78c474289f23b31283d8fb64c8e6a522a7fa47b131a3c6c141c8e6639d/faiss_cpu-1.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:9899c340f92bd94071d6faf4bef0ccb5362843daea42144d4ba857a2a1f67511", size = 30663859 }, + { url = "https://files.pythonhosted.org/packages/17/f0/194727b9e6e282e2877bc001ba886228f6af52e9a6730bbdb223e38591c3/faiss_cpu-1.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:345a52dbfa980d24b93c94410eadf82d1eef359c6a42e5e0768cca96539f1c3c", size = 13687087 }, + { url = "https://files.pythonhosted.org/packages/93/25/23239a83142faa319c4f8c025e25fec6cccc7418995eba3515218a57a45b/faiss_cpu-1.10.0-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:cb8473d69c3964c1bf3f8eb3e04287bb3275f536e6d9635ef32242b5f506b45d", size = 7720240 }, + { url = "https://files.pythonhosted.org/packages/18/f1/0e979277831af337739dbacf386d8a359a05eef9642df23d36e6c7d1b1a9/faiss_cpu-1.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:82ca5098de694e7b8495c1a8770e2c08df6e834922546dad0ae1284ff519ced6", size = 3239224 }, + { url = "https://files.pythonhosted.org/packages/bd/fa/c2ad85b017a5754f6cdb09c179f8c4f4198d2a264046a8daa7a4d080521f/faiss_cpu-1.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:035e4d797e2db7fc0d0c90531d4a655d089ad5d1382b7a49358c1f2307b3a309", size = 3671236 }, + { url = "https://files.pythonhosted.org/packages/4f/9b/759962f2c34800058f6a76457df3b0ab93b24f383650ea1ef0231acd322c/faiss_cpu-1.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:e02af3696a6b9e1f9072e502f48095a305de2163c42ceb1f6f6b1db9e7ffe574", size = 30663948 }, + { url = "https://files.pythonhosted.org/packages/2c/9a/6c496e0189897761978653177386452d62f4060579413d109bff05f458f2/faiss_cpu-1.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:e71f7e24d5b02d3a51df47b77bd10f394a1b48a8331d5c817e71e9e27a8a75ac", size = 13687212 }, +] + [[package]] name = "fastapi" version = "0.115.11" @@ -5059,6 +5091,9 @@ dapr = [ { name = "dapr-ext-fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "flask-dapr", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] +faiss = [ + { name = "faiss-cpu", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] google = [ { name = "google-cloud-aiplatform", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "google-generativeai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -5142,7 +5177,7 @@ requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = "~=0.32" }, { name = "autogen-agentchat", marker = "extra == 'autogen'", specifier = ">=0.2,<0.4" }, { name = "azure-ai-inference", marker = "extra == 'azure'", specifier = ">=1.0.0b6" }, - { name = "azure-ai-projects", marker = "extra == 'azure'", specifier = ">=1.0.0b5" }, + { name = "azure-ai-projects", marker = "extra == 'azure'", specifier = "==1.0.0b5" }, { name = "azure-core-tracing-opentelemetry", marker = "extra == 'azure'", specifier = ">=1.0.0b11" }, { name = "azure-cosmos", marker = "extra == 'azure'", specifier = "~=4.7" }, { name = "azure-identity", specifier = "~=1.13" }, @@ -5154,6 +5189,7 @@ requires-dist = [ { name = "dapr", marker = "extra == 'dapr'", specifier = ">=1.14.0" }, { name = "dapr-ext-fastapi", marker = "extra == 'dapr'", specifier = ">=1.14.0" }, { name = "defusedxml", specifier = "~=0.7" }, + { name = "faiss-cpu", marker = "extra == 'faiss'", specifier = ">=1.10.0" }, { name = "flask-dapr", marker = "extra == 'dapr'", specifier = ">=1.14.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'google'", specifier = "==1.82.0" }, { name = "google-generativeai", marker = "extra == 'google'", specifier = "~=0.8" }, @@ -5194,7 +5230,7 @@ requires-dist = [ { name = "websockets", specifier = ">=13,<15" }, { name = "websockets", marker = "extra == 'realtime'", specifier = ">=13,<15" }, ] -provides-extras = ["anthropic", "autogen", "aws", "azure", "chroma", "dapr", "google", "hugging-face", "milvus", "mistralai", "mongo", "notebooks", "ollama", "onnx", "pandas", "pinecone", "postgres", "qdrant", "realtime", "redis", "usearch", "weaviate"] +provides-extras = ["anthropic", "autogen", "aws", "azure", "chroma", "dapr", "faiss", "google", "hugging-face", "milvus", "mistralai", "mongo", "notebooks", "ollama", "onnx", "pandas", "pinecone", "postgres", "qdrant", "realtime", "redis", "usearch", "weaviate"] [package.metadata.requires-dev] dev = [