Python: Azure Cosmos DB NoSQL Vector Store & Collection implementation

.github/workflows/python-integration-tests.yml

@@ -119,6 +119,12 @@ jobs:
         run: docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest
       - name: Setup Weaviate docker deployment
         run: docker run -d -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.26.6
+      - name: Start Azure Cosmos DB emulator
+        if: matrix.os == 'windows-latest'
+        run: |
+          Write-Host "Launching Cosmos DB Emulator"
+          Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator"
+          Start-CosmosDbEmulator
       - name: Azure CLI Login
         if: github.event_name != 'pull_request'
         uses: azure/login@v2
@@ -170,6 +176,8 @@ jobs:
           VERTEX_AI_GEMINI_MODEL_ID: ${{ vars.VERTEX_AI_GEMINI_MODEL_ID }}
           VERTEX_AI_EMBEDDING_MODEL_ID: ${{ vars.VERTEX_AI_EMBEDDING_MODEL_ID }}
           REDIS_CONNECTION_STRING: ${{ vars.REDIS_CONNECTION_STRING }}
+          COSMOS_DB_NOSQL_URL: ${{ vars.COSMOS_DB_NOSQL_URL }}
+          COSMOS_DB_NOSQL_KEY: ${{ secrets.COSMOS_DB_NOSQL_KEY }}
         run: |
           uv run pytest -n logical --dist loadfile --dist worksteal ./tests/integration ./tests/samples -v --junitxml=pytest.xml
       - name: Surface failing tests
@@ -259,6 +267,12 @@ jobs:
         run: docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest
       - name: Setup Weaviate docker deployment
         run: docker run -d -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.26.6
+      - name: Start Azure Cosmos DB emulator
+        if: matrix.os == 'windows-latest'
+        run: |
+          Write-Host "Launching Cosmos DB Emulator"
+          Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator"
+          Start-CosmosDbEmulator
       - name: Azure CLI Login
         if: github.event_name != 'pull_request'
         uses: azure/login@v2
@@ -310,6 +324,8 @@ jobs:
           VERTEX_AI_GEMINI_MODEL_ID: ${{ vars.VERTEX_AI_GEMINI_MODEL_ID }}
           VERTEX_AI_EMBEDDING_MODEL_ID: ${{ vars.VERTEX_AI_EMBEDDING_MODEL_ID }}
           REDIS_CONNECTION_STRING: ${{ vars.REDIS_CONNECTION_STRING }}
+          COSMOS_DB_NOSQL_URL: ${{ vars.COSMOS_DB_NOSQL_URL }}
+          COSMOS_DB_NOSQL_KEY: ${{ secrets.COSMOS_DB_NOSQL_KEY }}
         run: |
           uv run pytest -n logical --dist loadfile --dist worksteal ./tests/integration ./tests/samples -v --junitxml=pytest.xml
       - name: Surface failing tests

python/.coveragerc

@@ -4,7 +4,7 @@ omit =
     semantic_kernel/connectors/memory/astradb/*
     semantic_kernel/connectors/memory/azure_cognitive_search/*
     semantic_kernel/connectors/memory/azure_cosmosdb/*
-    semantic_kernel/connectors/memory/azure_cosmosdb_no_sql/*
+    semantic_kernel/connectors/memory/azure_cosmosdb_no_sql/azure_cosmosdb_no_sql_memory_store.py
     semantic_kernel/connectors/memory/chroma/*
     semantic_kernel/connectors/memory/milvus/*
     semantic_kernel/connectors/memory/mongodb_atlas/*

python/.cspell.json

@@ -45,6 +45,7 @@
         "mongocluster",
         "ndarray",
         "nopep",
+        "NOSQL",
         "ollama",
         "onyourdatatest",
         "OPENAI",

python/samples/concepts/memory/new_memory.py

@@ -12,6 +12,9 @@
 from semantic_kernel.connectors.ai.open_ai import OpenAIEmbeddingPromptExecutionSettings, OpenAITextEmbedding
 from semantic_kernel.connectors.ai.open_ai.services.azure_text_embedding import AzureTextEmbedding
 from semantic_kernel.connectors.memory.azure_ai_search import AzureAISearchCollection
+from semantic_kernel.connectors.memory.azure_cosmosdb_no_sql.azure_cosmos_db_no_sql_collection import (
+    AzureCosmosDBNoSQLCollection,
+)
 from semantic_kernel.connectors.memory.postgres.postgres_collection import PostgresCollection
 from semantic_kernel.connectors.memory.qdrant import QdrantCollection
 from semantic_kernel.connectors.memory.redis import RedisHashsetCollection, RedisJsonCollection
@@ -25,55 +28,64 @@
     VectorStoreRecordVectorField,
     vectorstoremodel,
 )
-
-
-@vectorstoremodel
-@dataclass
-class MyDataModelArray:
-    vector: Annotated[
-        np.ndarray | None,
-        VectorStoreRecordVectorField(
-            embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
-            index_kind="hnsw",
-            dimensions=1536,
-            distance_function="cosine",
-            property_type="float",
-            serialize_function=np.ndarray.tolist,
-            deserialize_function=np.array,
-        ),
-    ] = None
-    other: str | None = None
-    id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
-    content: Annotated[
-        str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
-    ] = "content1"
-
-
-@vectorstoremodel
-@dataclass
-class MyDataModelList:
-    vector: Annotated[
-        list[float] | None,
-        VectorStoreRecordVectorField(
-            embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
-            index_kind="hnsw",
-            dimensions=1536,
-            distance_function="cosine",
-            property_type="float",
-        ),
-    ] = None
-    other: str | None = None
-    id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
-    content: Annotated[
-        str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
-    ] = "content1"
+from semantic_kernel.data.const import DistanceFunction, IndexKind
+
+
+def get_data_model_array(index_kind: IndexKind, distance_function: DistanceFunction) -> type:
+    @vectorstoremodel
+    @dataclass
+    class DataModelArray:
+        vector: Annotated[
+            np.ndarray | None,
+            VectorStoreRecordVectorField(
+                embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
+                index_kind=index_kind,
+                dimensions=1536,
+                distance_function=distance_function,
+                property_type="float",
+                serialize_function=np.ndarray.tolist,
+                deserialize_function=np.array,
+            ),
+        ] = None
+        other: str | None = None
+        id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+        content: Annotated[
+            str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+        ] = "content1"
+
+    return DataModelArray
+
+
+def get_data_model_list(index_kind: IndexKind, distance_function: DistanceFunction) -> type:
+    @vectorstoremodel
+    @dataclass
+    class DataModelList:
+        vector: Annotated[
+            list[float] | None,
+            VectorStoreRecordVectorField(
+                embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
+                index_kind=index_kind,
+                dimensions=1536,
+                distance_function=distance_function,
+                property_type="float",
+            ),
+        ] = None
+        other: str | None = None
+        id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+        content: Annotated[
+            str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+        ] = "content1"
+
+    return DataModelList
 
 
 collection_name = "test"
-MyDataModel = MyDataModelArray
+# Depending on the vector database, the index kind and distance function may need to be adjusted,
+# since not all combinations are supported by all databases.
+DataModel = get_data_model_array(IndexKind.HNSW, DistanceFunction.COSINE)
 
 # A list of VectorStoreRecordCollection that can be used.
-# Available stores are:
+# Available collections are:
 # - ai_search: Azure AI Search
 # - postgres: PostgreSQL
 # - redis_json: Redis JSON
@@ -83,63 +95,72 @@ class MyDataModelList:
 # - weaviate: Weaviate
 #   Please either configure the weaviate settings via environment variables or provide them through the constructor.
 #   Note that embed mode is not supported on Windows: https://github.com/weaviate/weaviate/issues/3315
-#
-# This is represented as a mapping from the store name to a
-# function which returns the store.
-# Using a function allows for lazy initialization of the store,
-# so that settings for unused stores do not cause validation errors.
-stores: dict[str, Callable[[], VectorStoreRecordCollection]] = {
-    "ai_search": lambda: AzureAISearchCollection[MyDataModel](
-        data_model_type=MyDataModel,
+# - azure_cosmos_nosql: Azure Cosmos NoSQL
+#   https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-develop-emulator?tabs=windows%2Cpython&pivots=api-nosql
+#   Please see the link above to learn how to set up the Azure Cosmos NoSQL emulator on your machine.
+#   For this sample to work with Azure Cosmos NoSQL, please adjust the index_kind of the data model to QUANTIZED_FLAT.
+# This is represented as a mapping from the collection name to a
+# function which returns the collection.
+# Using a function allows for lazy initialization of the collection,
+# so that settings for unused collections do not cause validation errors.
+collections: dict[str, Callable[[], VectorStoreRecordCollection]] = {
+    "ai_search": lambda: AzureAISearchCollection[DataModel](
+        data_model_type=DataModel,
     ),
-    "postgres": lambda: PostgresCollection[str, MyDataModel](
-        data_model_type=MyDataModel,
+    "postgres": lambda: PostgresCollection[str, DataModel](
+        data_model_type=DataModel,
         collection_name=collection_name,
     ),
-    "redis_json": lambda: RedisJsonCollection[MyDataModel](
-        data_model_type=MyDataModel,
+    "redis_json": lambda: RedisJsonCollection[DataModel](
+        data_model_type=DataModel,
         collection_name=collection_name,
         prefix_collection_name_to_key_names=True,
     ),
-    "redis_hashset": lambda: RedisHashsetCollection[MyDataModel](
-        data_model_type=MyDataModel,
+    "redis_hashset": lambda: RedisHashsetCollection[DataModel](
+        data_model_type=DataModel,
         collection_name=collection_name,
         prefix_collection_name_to_key_names=True,
     ),
-    "qdrant": lambda: QdrantCollection[MyDataModel](
-        data_model_type=MyDataModel, collection_name=collection_name, prefer_grpc=True, named_vectors=False
+    "qdrant": lambda: QdrantCollection[DataModel](
+        data_model_type=DataModel, collection_name=collection_name, prefer_grpc=True, named_vectors=False
+    ),
+    "volatile": lambda: VolatileCollection[DataModel](
+        data_model_type=DataModel,
+        collection_name=collection_name,
     ),
-    "volatile": lambda: VolatileCollection[MyDataModel](
-        data_model_type=MyDataModel,
+    "weaviate": lambda: WeaviateCollection[DataModel](
+        data_model_type=DataModel,
         collection_name=collection_name,
     ),
-    "weaviate": lambda: WeaviateCollection[MyDataModel](
-        data_model_type=MyDataModel,
+    "azure_cosmos_nosql": lambda: AzureCosmosDBNoSQLCollection(
+        data_model_type=DataModel,
+        database_name="sample_database",
         collection_name=collection_name,
+        create_database=True,
     ),
 }
 
 
-async def main(store: str, use_azure_openai: bool, embedding_model: str):
+async def main(collection: str, use_azure_openai: bool, embedding_model: str):
     kernel = Kernel()
     service_id = "embedding"
     if use_azure_openai:
         kernel.add_service(AzureTextEmbedding(service_id=service_id, deployment_name=embedding_model))
     else:
         kernel.add_service(OpenAITextEmbedding(service_id=service_id, ai_model_id=embedding_model))
-    async with stores[store]() as record_store:
-        await record_store.create_collection_if_not_exists()
+    async with collections[collection]() as record_collection:
+        await record_collection.create_collection_if_not_exists()
 
-        record1 = MyDataModel(content="My text", id="e6103c03-487f-4d7d-9c23-4723651c17f4")
-        record2 = MyDataModel(content="My other text", id="09caec77-f7e1-466a-bcec-f1d51c5b15be")
+        record1 = DataModel(content="My text", id="e6103c03-487f-4d7d-9c23-4723651c17f4")
+        record2 = DataModel(content="My other text", id="09caec77-f7e1-466a-bcec-f1d51c5b15be")
 
         records = await VectorStoreRecordUtils(kernel).add_vector_to_records(
-            [record1, record2], data_model_type=MyDataModel
+            [record1, record2], data_model_type=DataModel
         )
-        keys = await record_store.upsert_batch(records)
+        keys = await record_collection.upsert_batch(records)
         print(f"upserted {keys=}")
 
-        results = await record_store.get_batch([record1.id, record2.id])
+        results = await record_collection.get_batch([record1.id, record2.id])
         if results:
             for result in results:
                 print(f"found {result.id=}")
@@ -156,7 +177,7 @@ async def main(store: str, use_azure_openai: bool, embedding_model: str):
     argparse.ArgumentParser()
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("--store", default="volatile", choices=stores.keys(), help="What store to use.")
+    parser.add_argument("--collection", default="volatile", choices=collections.keys(), help="What collection to use.")
     # Option of whether to use OpenAI or Azure OpenAI.
     parser.add_argument("--use-azure-openai", action="store_true", help="Use Azure OpenAI instead of OpenAI.")
     # Model
@@ -165,4 +186,4 @@ async def main(store: str, use_azure_openai: bool, embedding_model: str):
     )
     args = parser.parse_args()
 
-    asyncio.run(main(store=args.store, use_azure_openai=args.use_azure_openai, embedding_model=args.model))
+    asyncio.run(main(collection=args.collection, use_azure_openai=args.use_azure_openai, embedding_model=args.model))

python/semantic_kernel/connectors/memory/azure_cosmosdb_no_sql/azure_cosmos_db_no_sql_base.py

@@ -0,0 +1,67 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from azure.cosmos.aio import ContainerProxy, CosmosClient, DatabaseProxy
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+from azure.identity import DefaultAzureCredential
+
+from semantic_kernel.connectors.memory.azure_cosmosdb_no_sql.azure_cosmos_db_no_sql_settings import (
+    AzureCosmosDBNoSQLSettings,
+)
+from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorResourceNotFound
+from semantic_kernel.kernel_pydantic import KernelBaseModel
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+
+@experimental_class
+class AzureCosmosDBNoSQLBase(KernelBaseModel):
+    """An Azure Cosmos DB NoSQL collection stores documents in a Azure Cosmos DB NoSQL account."""
+
+    cosmos_db_nosql_settings: AzureCosmosDBNoSQLSettings
+    database_name: str
+    # If create_database is True, the database will be created
+    # if it does not exist when an operation requires a database.
+    create_database: bool
+
+    def _get_cosmos_client(self) -> CosmosClient:
+        """Gets the Cosmos client.
+
+        We cannot cache the Cosmos client because it is only good for one context.
+        https://github.com/Azure/azure-sdk-for-python/issues/25640
+        """
+        if not self.cosmos_db_nosql_settings.key:
+            return CosmosClient(str(self.cosmos_db_nosql_settings.url), credential=DefaultAzureCredential())
+
+        return CosmosClient(
+            str(self.cosmos_db_nosql_settings.url),
+            credential=self.cosmos_db_nosql_settings.key.get_secret_value(),
+        )
+
+    async def _does_database_exist(self, cosmos_client: CosmosClient) -> bool:
+        """Checks if the database exists."""
+        try:
+            await cosmos_client.get_database_client(self.database_name).read()
+            return True
+        except CosmosResourceNotFoundError:
+            return False
+        except Exception as e:
+            raise MemoryConnectorResourceNotFound(f"Failed to check if database '{self.database_name}' exists.") from e
+
+    async def _get_database_proxy(self, cosmos_client: CosmosClient) -> DatabaseProxy:
+        """Gets the database proxy."""
+        try:
+            if await self._does_database_exist(cosmos_client):
+                return cosmos_client.get_database_client(self.database_name)
+
+            if self.create_database:
+                return await cosmos_client.create_database(self.database_name)
+            raise MemoryConnectorResourceNotFound(f"Database '{self.database_name}' does not exist.")
+        except Exception as e:
+            raise MemoryConnectorResourceNotFound(f"Failed to get database proxy for '{id}'.") from e
+
+    async def _get_container_proxy(self, container_name: str, cosmos_client: CosmosClient) -> ContainerProxy:
+        """Gets the container proxy."""
+        try:
+            database_proxy = await self._get_database_proxy(cosmos_client)
+            return database_proxy.get_container_client(container_name)
+        except Exception as e:
+            raise MemoryConnectorResourceNotFound(f"Failed to get container proxy for '{container_name}'.") from e