microsoft--semantic-kernel/python/semantic_kernel/connectors/azure_cosmos_db.py

# Copyright (c) Microsoft. All rights reserved.

import ast
import asyncio
import sys
from collections.abc import Sequence
from importlib import metadata
from typing import Any, ClassVar, Final, Generic, TypeVar

from azure.core.credentials_async import AsyncTokenCredential
from azure.cosmos.aio import ContainerProxy, CosmosClient, DatabaseProxy
from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceNotFoundError
from azure.cosmos.partition_key import PartitionKey
from pydantic import HttpUrl, SecretStr, ValidationError
from pymongo import AsyncMongoClient
from pymongo.driver_info import DriverInfo
from typing_extensions import override

from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
from semantic_kernel.connectors.mongodb import (
    DEFAULT_DB_NAME,
    MONGODB_SCORE_FIELD,
    MongoDBAtlasCollection,
    MongoDBAtlasStore,
)
from semantic_kernel.data.vector import (
    DistanceFunction,
    FieldTypes,
    GetFilteredRecordOptions,
    IndexKind,
    KernelSearchResults,
    SearchType,
    TModel,
    VectorSearch,
    VectorSearchOptions,
    VectorSearchResult,
    VectorStore,
    VectorStoreCollection,
    VectorStoreCollectionDefinition,
    _get_collection_name_from_model,
)
from semantic_kernel.exceptions import (
    VectorSearchExecutionException,
    VectorStoreInitializationException,
    VectorStoreModelDeserializationException,
    VectorStoreModelException,
    VectorStoreOperationException,
)
from semantic_kernel.kernel_pydantic import KernelBaseModel, KernelBaseSettings
from semantic_kernel.utils.feature_stage_decorator import release_candidate
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT

if sys.version_info >= (3, 12):
    from typing import override  # pragma: no cover
else:
    from typing_extensions import override  # pragma: no cover


@release_candidate
class CosmosNoSqlCompositeKey(KernelBaseModel):
    """CosmosDB NoSQL composite key."""

    partition_key: str
    key: str


TKey = TypeVar("TKey", bound=str)
TNoSQLKey = TypeVar("TNoSQLKey", str, CosmosNoSqlCompositeKey)

# region: Constants

COSMOS_ITEM_ID_PROPERTY_NAME: Final[str] = "id"
NOSQL_SCORE_PROPERTY_NAME: Final[str] = "distance"
INDEX_KIND_MAP_NOSQL: Final[dict[IndexKind, str]] = {
    IndexKind.FLAT: "flat",
    IndexKind.QUANTIZED_FLAT: "quantizedFlat",
    IndexKind.DISK_ANN: "diskANN",
    IndexKind.DEFAULT: "flat",
}
INDEX_KIND_MAP_MONGODB: Final[dict[IndexKind, str]] = {
    IndexKind.IVF_FLAT: "vector-ivf",
    IndexKind.HNSW: "vector-hnsw",
    IndexKind.DISK_ANN: "vector-diskann",
    IndexKind.DEFAULT: "vector-ivf",
}
DISTANCE_FUNCTION_MAP_NOSQL: Final[dict[DistanceFunction, str]] = {
    DistanceFunction.COSINE_SIMILARITY: "cosine",
    DistanceFunction.DOT_PROD: "dotproduct",
    DistanceFunction.EUCLIDEAN_DISTANCE: "euclidean",
    DistanceFunction.DEFAULT: "cosine",
}
DISTANCE_FUNCTION_MAP_MONGODB: Final[dict[DistanceFunction, str]] = {
    DistanceFunction.COSINE_SIMILARITY: "COS",
    DistanceFunction.DOT_PROD: "IP",
    DistanceFunction.EUCLIDEAN_DISTANCE: "L2",
    DistanceFunction.DEFAULT: "COS",
}
VECTOR_DATATYPES_MAP: Final[dict[str, str]] = {
    "default": "float32",
    "float": "float32",
    "list[float]": "float32",
    "int": "int32",
    "list[int]": "int32",
}

# region: Helpers


def _create_default_indexing_policy_nosql(definition: VectorStoreCollectionDefinition) -> dict[str, Any]:
    """Creates a default indexing policy for the Azure Cosmos DB NoSQL container.

    A default indexing policy is created based on the data model definition and has an automatic indexing policy.

    Args:
        definition (VectorStoreRecordDefinition): The definition of the data model.

    Returns:
        dict[str, Any]: The indexing policy.

    Raises:
        VectorStoreModelException: If the field is not full text searchable and not filterable.
    """
    indexing_policy: dict[str, Any] = {
        "automatic": True,
        "includedPaths": [
            {
                "path": "/*",
            }
        ],
        "excludedPaths": [
            {
                "path": '/"_etag"/?',
            }
        ],
        "vectorIndexes": [],
    }

    for field in definition.fields:
        if field.field_type == FieldTypes.DATA and (not field.is_full_text_indexed and not field.is_indexed):
            indexing_policy["excludedPaths"].append({"path": f'/"{field.storage_name or field.name}"/*'})

        if field.field_type == FieldTypes.VECTOR:
            if field.index_kind not in INDEX_KIND_MAP_NOSQL:
                raise VectorStoreModelException(
                    f"Index kind '{field.index_kind}' is not supported by Azure Cosmos DB NoSQL container."
                )
            indexing_policy["vectorIndexes"].append({
                "path": f'/"{field.storage_name or field.name}"',
                "type": INDEX_KIND_MAP_NOSQL[field.index_kind],
            })
            # Exclude the vector field from the index for performance optimization.
            indexing_policy["excludedPaths"].append({"path": f'/"{field.storage_name or field.name}"/*'})

    return indexing_policy


def _create_default_vector_embedding_policy(definition: VectorStoreCollectionDefinition) -> dict[str, Any]:
    """Creates a default vector embedding policy for the Azure Cosmos DB NoSQL container.

    A default vector embedding policy is created based on the data model definition.

    Args:
        definition (VectorStoreRecordDefinition): The definition of the data model.

    Returns:
        dict[str, Any]: The vector embedding policy.

    Raises:
        VectorStoreModelException: If the datatype or distance function is not supported by Azure Cosmos DB NoSQL.

    """
    vector_embedding_policy: dict[str, Any] = {"vectorEmbeddings": []}

    for field in definition.fields:
        if field.field_type == FieldTypes.VECTOR:
            if field.distance_function not in DISTANCE_FUNCTION_MAP_NOSQL:
                raise VectorStoreModelException(
                    f"Distance function '{field.distance_function}' is not supported by Azure Cosmos DB NoSQL."
                )
            if field.type_ and field.type_ not in VECTOR_DATATYPES_MAP:
                raise VectorStoreModelException(
                    f"Vector property type '{field.type_}' is not supported by Azure Cosmos DB NoSQL."
                )

            vector_embedding_policy["vectorEmbeddings"].append({
                "path": f'/"{field.storage_name or field.name}"',
                "dataType": VECTOR_DATATYPES_MAP[field.type_ or "default"],
                "distanceFunction": DISTANCE_FUNCTION_MAP_NOSQL[field.distance_function],
                "dimensions": field.dimensions,
            })

    return vector_embedding_policy


def _get_key(key: str | CosmosNoSqlCompositeKey) -> str:
    """Gets the key value from the key."""
    if isinstance(key, CosmosNoSqlCompositeKey):
        return key.key
    return key


def _get_partition_key(key: str | CosmosNoSqlCompositeKey) -> str:
    """Gets the partition key value from the key."""
    if isinstance(key, CosmosNoSqlCompositeKey):
        return key.partition_key
    return key


# region: Settings


@release_candidate
class CosmosMongoSettings(KernelBaseSettings):
    """Azure CosmosDB for MongoDB settings.

    The settings are first loaded from environment variables with
    the prefix 'AZURE_COSMOS_DB_MONGODB_'.
    If the environment variables are not found, the settings can
    be loaded from a .env file with the encoding 'utf-8'.
    If the settings are not found in the .env file, the settings
    are ignored; however, validation will fail alerting that the
    settings are missing.

    Required settings for prefix 'AZURE_COSMOS_DB_MONGODB_':
    - connection_string:  The connection string of the Azure CosmosDB for MongoDB account.
       This value can be found in the Keys & Endpoint section when examining
       your resource from the Azure portal.
       (Env var name: AZURE_COSMOS_DB_MONGODB_CONNECTION_STRING)
    - database_name: str - The name of the database. Please refer to this documentation
       on Azure CosmosDB NoSQL resource model:
       https://learn.microsoft.com/en-us/azure/cosmos-db/resource-model
       (Env var name: AZURE_COSMOS_DB_MONGODB_DATABASE_NAME)
    """

    env_prefix: ClassVar[str] = "AZURE_COSMOS_DB_MONGODB_"

    connection_string: SecretStr
    database_name: str = DEFAULT_DB_NAME


@release_candidate
class CosmosNoSqlSettings(KernelBaseSettings):
    """Azure CosmosDB NoSQL settings.

    The settings are first loaded from environment variables with
    the prefix 'AZURE_COSMOS_DB_NO_SQL_'.
    If the environment variables are not found, the settings can
    be loaded from a .env file with the encoding 'utf-8'.
    If the settings are not found in the .env file, the settings
    are ignored; however, validation will fail alerting that the
    settings are missing.

    Required settings for prefix 'AZURE_COSMOS_DB_NO_SQL_':
    - url: HttpsUrl - The uri of the Azure CosmosDB NoSQL account.
       This value can be found in the Keys & Endpoint section when examining
       your resource from the Azure portal.
       (Env var name: AZURE_COSMOS_DB_NO_SQL_URL)

    Optional settings for prefix 'AZURE_COSMOS_DB_NO_SQL_':
    - key: SecretStr - The primary key of the Azure CosmosDB NoSQL account.
       This value can be found in the Keys & Endpoint section when examining
       your resource from the Azure portal.
       (Env var name: AZURE_COSMOS_DB_NO_SQL_KEY)
    - database_name: str - The name of the database. Please refer to this documentation
       on Azure CosmosDB NoSQL resource model:
       https://learn.microsoft.com/en-us/azure/cosmos-db/resource-model
       (Env var name: AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME)
    """

    env_prefix: ClassVar[str] = "AZURE_COSMOS_DB_NO_SQL_"

    url: HttpUrl
    key: SecretStr | None = None
    database_name: str | None = None


# region: Mongo Collection


@release_candidate
class CosmosMongoCollection(MongoDBAtlasCollection[TKey, TModel], Generic[TKey, TModel]):
    """Azure Cosmos DB for MongoDB collection."""

    def __init__(
        self,
        record_type: type[TModel],
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        mongo_client: AsyncMongoClient | None = None,
        connection_string: str | None = None,
        database_name: str | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Initializes a new instance of the CosmosMongoCollection class.

        Args:
            record_type: The type of the data model.
            definition: The model definition, optional.
            collection_name: The name of the collection, optional.
            embedding_generator: The embedding generator to use for generating embeddings.
            mongo_client: The MongoDB client for interacting with Azure CosmosDB for MongoDB,
                used for creating and deleting collections.
            connection_string: The connection string for MongoDB Atlas, optional.
            Can be read from environment variables.
            database_name: The name of the database, will be filled from the env when this is not set.
            connection_string: str | None = None,
            env_file_path: str | None = None,
            env_file_encoding: str | None = None
            **kwargs: Additional keyword arguments

        """
        if not collection_name:
            collection_name = _get_collection_name_from_model(record_type, definition)
        if not collection_name:
            raise VectorStoreInitializationException(
                "The collection name is required, can be passed directly or through the data model."
            )
        managed_client = not mongo_client
        if mongo_client:
            super().__init__(
                record_type=record_type,
                definition=definition,
                mongo_client=mongo_client,
                collection_name=collection_name,
                database_name=database_name or DEFAULT_DB_NAME,
                managed_client=managed_client,
                embedding_generator=embedding_generator,
            )
            return

        try:
            settings = CosmosMongoSettings(
                env_file_path=env_file_path,
                env_file_encoding=env_file_encoding,
                connection_string=connection_string,
                database_name=database_name,
            )
        except ValidationError as exc:
            raise VectorStoreInitializationException("Failed to create Azure CosmosDB for MongoDB settings.") from exc

        mongo_client = AsyncMongoClient(
            settings.connection_string.get_secret_value(),
            driver=DriverInfo(SEMANTIC_KERNEL_USER_AGENT, metadata.version("semantic-kernel")),
        )

        super().__init__(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            mongo_client=mongo_client,
            managed_client=managed_client,
            database_name=settings.database_name,
            embedding_generator=embedding_generator,
        )

    @override
    async def ensure_collection_exists(self, **kwargs) -> None:
        """Create a new collection in Azure CosmosDB for MongoDB.

        This first creates a collection, with the kwargs.
        Then creates a search index based on the data model definition.

        By the naming convection of MongoDB indexes are created by using the field name
        with a underscore.

        Args:
            **kwargs: Additional keyword arguments.
                These are the additional keyword arguments for creating
                vector indexes in Azure Cosmos DB for MongoDB.
                And they depend on the kind of index you are creating.
                See https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/vcore/vector-search
                for more information.
                Other kwargs are passed to the create_collection method.
        """
        await self._get_database().create_collection(self.collection_name, **kwargs)
        await self._get_database().command(command=self._get_index_definitions(**kwargs))

    def _get_index_definitions(self, **kwargs: Any) -> dict[str, Any]:
        """Creates index definitions for the collection."""
        indexes = [
            {
                "name": f"{field.storage_name or field.name}_",
                FieldTypes.KEY: {field.storage_name or field.name: 1},
            }
            for field in self.definition.data_fields
            if field.is_indexed or field.is_full_text_indexed
        ]
        for field in self.definition.vector_fields:
            if field.index_kind not in INDEX_KIND_MAP_MONGODB:
                raise VectorStoreModelException(
                    f"Index kind '{field.index_kind}' is not supported by Azure Cosmos DB for MongoDB."
                )
            if field.distance_function not in DISTANCE_FUNCTION_MAP_MONGODB:
                raise VectorStoreModelException(
                    f"Distance function '{field.distance_function}' is not supported by Azure Cosmos DB for MongoDB."
                )
            index_name = f"{field.storage_name or field.name}_"
            index_kind = DISTANCE_FUNCTION_MAP_MONGODB[field.distance_function]
            index: dict[str, Any] = {
                "name": index_name,
                FieldTypes.KEY: {field.storage_name or field.name: "cosmosSearch"},
                "cosmosSearchOptions": {
                    "kind": index_kind,
                    "similarity": DISTANCE_FUNCTION_MAP_MONGODB[field.distance_function],
                    "dimensions": field.dimensions,
                },
            }
            match index_kind:
                case "vector-diskann":
                    if "maxDegree" in kwargs:
                        index["cosmosSearchOptions"]["maxDegree"] = kwargs["maxDegree"]
                    if "lBuild" in kwargs:
                        index["cosmosSearchOptions"]["lBuild"] = kwargs["lBuild"]
                case "vector-hnsw":
                    if "m" in kwargs:
                        index["cosmosSearchOptions"]["m"] = kwargs["m"]
                    if "efConstruction" in kwargs:
                        index["cosmosSearchOptions"]["efConstruction"] = kwargs["efConstruction"]
                case "vector-ivf":
                    if "numList" in kwargs:
                        index["cosmosSearchOptions"]["numList"] = kwargs["numList"]
            indexes.append(index)

        return {"createIndexes": self.collection_name, "indexes": indexes}

    @override
    async def _inner_vector_search(
        self,
        options: VectorSearchOptions,
        values: Any | None = None,
        vector: Sequence[float | int] | None = None,
        **kwargs: Any,
    ) -> KernelSearchResults[VectorSearchResult[TModel]]:
        collection = self._get_collection()
        vector_field = self.definition.try_get_vector_field(options.vector_property_name)
        if not vector_field:
            raise VectorStoreModelException(
                f"Vector field '{options.vector_property_name}' not found in the data model definition."
            )
        if not vector:
            vector = await self._generate_vector_from_values(values, options)
        vector_search_query: dict[str, Any] = {
            "k": options.top + options.skip,
            "index": f"{vector_field.storage_name or vector_field.name}_",
            "vector": vector,
            "path": vector_field.storage_name or vector_field.name,
        }
        if filter := self._build_filter(options.filter):  # type: ignore
            vector_search_query["filter"] = filter if isinstance(filter, dict) else {"$and": filter}

        projection_query: dict[str, int | dict] = {
            field: 1
            for field in self.definition.get_names(
                include_vector_fields=options.include_vectors,
                include_key_field=False,  # _id is always included
            )
        }
        projection_query[MONGODB_SCORE_FIELD] = {"$meta": "searchScore"}
        try:
            raw_results = await collection.aggregate([
                {"$search": {"cosmosSearch": vector_search_query}},
                {"$project": projection_query},
            ])
        except Exception as exc:
            raise VectorSearchExecutionException("Failed to search the collection.") from exc
        return KernelSearchResults(
            results=self._get_vector_search_results_from_results(raw_results, options),
            total_count=None,  # no way to get a count before looping through the result cursor
        )


# region: Mongo Store


@release_candidate
class CosmosMongoStore(MongoDBAtlasStore):
    """Azure Cosmos DB for MongoDB store."""

    def __init__(
        self,
        connection_string: str | None = None,
        database_name: str | None = None,
        mongo_client: AsyncMongoClient | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
    ) -> None:
        """Initializes a new instance of the CosmosMongoStore class."""
        managed_client: bool = not mongo_client
        if mongo_client:
            super().__init__(
                mongo_client=mongo_client,
                managed_client=managed_client,
                database_name=database_name,
                embedding_generator=embedding_generator,
            )
            return
        try:
            settings = CosmosMongoSettings(
                env_file_path=env_file_path,
                env_file_encoding=env_file_encoding,
                connection_string=connection_string,
                database_name=database_name,
            )
        except ValidationError as exc:
            raise VectorStoreInitializationException("Failed to create Azure CosmosDB for MongoDB settings.") from exc

        mongo_client = AsyncMongoClient(
            settings.connection_string.get_secret_value(),
            driver=DriverInfo(SEMANTIC_KERNEL_USER_AGENT, metadata.version("semantic-kernel")),
        )

        super().__init__(
            mongo_client=mongo_client,
            managed_client=managed_client,
            database_name=settings.database_name,
            embedding_generator=embedding_generator,
        )

    @override
    def get_collection(
        self,
        record_type: type[TModel],
        *,
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        **kwargs: Any,
    ) -> CosmosMongoCollection:
        return CosmosMongoCollection(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            mongo_client=self.mongo_client,
            managed_client=False,
            database_name=self.database_name,
            embedding_generator=embedding_generator or self.embedding_generator,
            **kwargs,
        )


# region: NoSQL Base


@release_candidate
class CosmosNoSqlBase(KernelBaseModel):
    """An Azure Cosmos DB NoSQL collection stores documents in a Azure Cosmos DB NoSQL account."""

    cosmos_client: CosmosClient
    database_name: str
    cosmos_db_nosql_settings: CosmosNoSqlSettings
    # 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 __init__(
        self,
        url: str | None = None,
        key: str | None = None,
        database_name: str | None = None,
        cosmos_client: CosmosClient | None = None,
        create_database: bool = False,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        credential: AsyncTokenCredential | None = None,
        **kwargs,
    ):
        """Initialize the CosmosNoSqlBase.

        Args:
            url (str): The URL of the Azure Cosmos DB NoSQL account. Defaults to None.
            key (str): The key of the Azure Cosmos DB NoSQL account. Defaults to None.
            database_name (str): The name of the database. The database may not exist yet. If it does not exist,
                                 it will be created when the first collection is created. Defaults to None.
            cosmos_client (CosmosClient): The custom Azure Cosmos DB NoSQL client whose lifetime is managed by the user.
                                          Defaults to None.
            create_database (bool): If True, the database will be created if it does not exist.
                                    Defaults to False.
            env_file_path (str): The path to the .env file. Defaults to None.
            env_file_encoding (str): The encoding of the .env file. Defaults to None.
            credential: The credential to use for authentication to Azure Cosmos DB NoSQL.
            kwargs: Additional keyword arguments.
        """
        try:
            cosmos_db_nosql_settings = CosmosNoSqlSettings(
                url=url,
                key=key,
                database_name=database_name,
                env_file_path=env_file_path,
                env_file_encoding=env_file_encoding,
            )
        except ValidationError as e:
            raise VectorStoreInitializationException("Failed to validate Azure Cosmos DB NoSQL settings.") from e

        if cosmos_db_nosql_settings.database_name is None:
            raise VectorStoreInitializationException("The name of the Azure Cosmos DB NoSQL database is missing.")

        if cosmos_client is None:
            if cosmos_db_nosql_settings.key is not None:
                cosmos_client = CosmosClient(
                    str(cosmos_db_nosql_settings.url), credential=cosmos_db_nosql_settings.key.get_secret_value()
                )
            else:
                if credential is None:
                    raise VectorStoreInitializationException(
                        "The 'credential' parameter is required for authentication."
                    )
                cosmos_client = CosmosClient(str(cosmos_db_nosql_settings.url), credential=credential)

        super().__init__(
            cosmos_client=cosmos_client,
            database_name=cosmos_db_nosql_settings.database_name,
            cosmos_db_nosql_settings=cosmos_db_nosql_settings,
            create_database=create_database,
            **kwargs,
        )

    async def _does_database_exist(self) -> bool:
        """Checks if the database exists."""
        try:
            await self.cosmos_client.get_database_client(self.database_name).read()
            return True
        except CosmosResourceNotFoundError:
            return False
        except Exception as e:
            raise VectorStoreOperationException(
                f"Failed to check if database '{self.database_name}' exists, with message {e}"
            ) from e

    async def _get_database_proxy(self, **kwargs) -> DatabaseProxy:
        """Gets the database proxy."""
        try:
            if await self._does_database_exist():
                return self.cosmos_client.get_database_client(self.database_name)

            if self.create_database:
                return await self.cosmos_client.create_database(self.database_name, **kwargs)
            raise VectorStoreOperationException(f"Database '{self.database_name}' does not exist.")
        except Exception as e:
            raise VectorStoreOperationException(f"Failed to get database proxy for '{id}'.") from e

    async def _get_container_proxy(self, container_name: str, **kwargs) -> ContainerProxy:
        """Gets the container proxy."""
        try:
            database_proxy = await self._get_database_proxy(**kwargs)
            return database_proxy.get_container_client(container_name)
        except Exception as e:
            raise VectorStoreOperationException(f"Failed to get container proxy for '{container_name}'.") from e


# region: NoSQL Collection


@release_candidate
class CosmosNoSqlCollection(
    CosmosNoSqlBase,
    VectorStoreCollection[TNoSQLKey, TModel],
    VectorSearch[TNoSQLKey, TModel],
    Generic[TNoSQLKey, TModel],
):
    """An Azure Cosmos DB NoSQL collection stores documents in a Azure Cosmos DB NoSQL account."""

    partition_key: PartitionKey
    supported_key_types: ClassVar[set[str] | None] = {"str", "CosmosNoSqlCompositeKey"}
    supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR, SearchType.KEYWORD_HYBRID}

    def __init__(
        self,
        record_type: type[TModel],
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        database_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        url: str | None = None,
        key: str | None = None,
        cosmos_client: CosmosClient | None = None,
        partition_key: PartitionKey | str | None = None,
        create_database: bool = False,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        credential: AsyncTokenCredential | None = None,
    ):
        """Initializes a new instance of the CosmosNoSqlCollection class.

        Args:
            record_type: The type of the data model.
            collection_name: The name of the collection.
            database_name: The name of the database. Used to create a database proxy if not provided.
                                 Defaults to None.
            embedding_generator: The embedding generator to use for generating embeddings.
            definition: The definition of the data model. Defaults to None.
            url: The URL of the Azure Cosmos DB NoSQL account. Defaults to None.
            key: The key of the Azure Cosmos DB NoSQL account. Defaults to None.
            cosmos_client: The custom Azure Cosmos DB NoSQL client whose lifetime is managed by the user.
            partition_key: The partition key. Defaults to None. If not provided, the partition
                                                key will be based on the key field of the data model definition.
                                                https://learn.microsoft.com/en-us/azure/cosmos-db/partitioning-overview
            create_database: Indicates whether to create the database if it does not exist.
                                    Defaults to False.
            env_file_path: The path to the .env file. Defaults to None.
            env_file_encoding: The encoding of the .env file. Defaults to None.
            credential: The credential to use for authentication to Azure Cosmos DB NoSQL.
        """
        if not collection_name:
            collection_name = _get_collection_name_from_model(record_type, definition)
        if not partition_key:
            partition_key = PartitionKey(path=f"/{COSMOS_ITEM_ID_PROPERTY_NAME}")
        elif isinstance(partition_key, str):
            partition_key = PartitionKey(path=f"/{partition_key.strip('/')}")

        super().__init__(
            partition_key=partition_key,
            url=url,
            key=key,
            database_name=database_name,
            cosmos_client=cosmos_client,
            create_database=create_database,
            env_file_path=env_file_path,
            env_file_encoding=env_file_encoding,
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            managed_client=cosmos_client is None,
            embedding_generator=embedding_generator,
            credential=credential,
        )

    @override
    async def _inner_upsert(
        self,
        records: Sequence[Any],
        **kwargs: Any,
    ) -> Sequence[TNoSQLKey]:
        container_proxy = await self._get_container_proxy(self.collection_name, **kwargs)
        results = await asyncio.gather(*(container_proxy.upsert_item(record) for record in records))
        return [result[COSMOS_ITEM_ID_PROPERTY_NAME] for result in results]

    @override
    async def _inner_get(  # type: ignore
        self,
        keys: Sequence[TNoSQLKey] | None = None,
        options: GetFilteredRecordOptions | None = None,
        **kwargs: Any,
    ) -> Sequence[Any] | None:
        if not keys:
            if options is not None:
                raise NotImplementedError("Get without keys is not yet implemented.")
            return None
        include_vectors = kwargs.pop("include_vectors", False)
        query = (
            f"SELECT {self._build_select_clause(include_vectors)} FROM c WHERE "  # nosec: B608
            f"c.id IN ({', '.join([f'@id{i}' for i in range(len(keys))])})"  # nosec: B608
        )  # nosec: B608
        parameters: list[dict[str, Any]] = [{"name": f"@id{i}", "value": _get_key(key)} for i, key in enumerate(keys)]

        container_proxy = await self._get_container_proxy(self.collection_name, **kwargs)
        return [item async for item in container_proxy.query_items(query=query, parameters=parameters)]

    @override
    async def _inner_delete(self, keys: Sequence[TNoSQLKey], **kwargs: Any) -> None:  # type: ignore
        container_proxy = await self._get_container_proxy(self.collection_name, **kwargs)
        results = await asyncio.gather(
            *[container_proxy.delete_item(item=_get_key(key), partition_key=_get_partition_key(key)) for key in keys],
            return_exceptions=True,
        )
        exceptions = [result for result in results if isinstance(result, Exception)]
        if exceptions:
            raise VectorStoreOperationException("Failed to delete item(s).", exceptions)

    @override
    async def _inner_search(
        self,
        search_type: SearchType,
        options: VectorSearchOptions,
        values: Any | None = None,
        vector: Sequence[float | int] | None = None,
        **kwargs: Any,
    ) -> KernelSearchResults[VectorSearchResult[TModel]]:
        params = [{"name": "@top", "value": options.top}]
        vector_field = self.definition.try_get_vector_field(options.vector_property_name)
        if not vector_field:
            raise VectorStoreModelException(
                f"Vector field '{options.vector_property_name}' not found in the data model definition."
            )
        if not vector:
            vector = await self._generate_vector_from_values(values, options)

        if where_clauses := self._build_filter(options.filter):  # type: ignore
            where_clauses = (
                f"WHERE {where_clauses} "
                if isinstance(where_clauses, str)
                else f"WHERE ({' AND '.join(where_clauses)}) "
            )
        else:
            where_clauses = ""  # Empty string instead of None
        vector_field_name = vector_field.storage_name or vector_field.name
        select_clause = self._build_select_clause(options.include_vectors)
        params.append({"name": "@vector", "value": vector})
        if vector_field.distance_function not in DISTANCE_FUNCTION_MAP_NOSQL:
            raise VectorStoreModelException(
                f"Distance function '{vector_field.distance_function}' is not supported by Azure Cosmos DB NoSQL."
            )
        # Cosmos DB VectorDistance function only accepts 2 parameters: field and vector
        # Distance function is configured in the vector index, not in the query
        if search_type == SearchType.VECTOR:
            distance_clause = f"VectorDistance(c.{vector_field_name}, @vector)"
        elif search_type == SearchType.KEYWORD_HYBRID:
            # Hybrid search: requires both a vector and keywords
            params.append({"name": "@keywords", "value": values})
            text_field = options.additional_property_name
            if not text_field:
                raise VectorStoreModelException("Hybrid search requires 'keyword_field_name' in options.")
            distance_clause = (
                f"RRF(VectorDistance(c.{vector_field_name}, @vector), FullTextScore(c.{text_field}, @keywords))"
            )
        else:
            raise VectorStoreModelException(f"Search type '{search_type}' is not supported.")
        query = (
            f"SELECT TOP @top {select_clause}, "  # nosec: B608
            f"{distance_clause} as {NOSQL_SCORE_PROPERTY_NAME} "  # nosec: B608
            "FROM c "
            f"{where_clauses}"  # nosec: B608
            f"ORDER BY {distance_clause}"  # nosec: B608
        )

        container_proxy = await self._get_container_proxy(self.collection_name, **kwargs)
        try:
            results = container_proxy.query_items(query, parameters=params)
        except Exception as exc:
            raise VectorSearchExecutionException("Failed to search items.") from exc
        return KernelSearchResults(
            results=self._get_vector_search_results_from_results(results, options),
            total_count=None,
        )

    def _build_select_clause(self, include_vectors: bool) -> str:
        """Create the select clause for a CosmosDB query."""
        included_fields = [field for field in self.definition.get_storage_names(include_vector_fields=include_vectors)]
        if self.definition.key_name != COSMOS_ITEM_ID_PROPERTY_NAME:
            # Replace the key field name with the Cosmos item id property name
            included_fields = [
                field if field != self.definition.key_name else COSMOS_ITEM_ID_PROPERTY_NAME
                for field in included_fields
            ]

        return ", ".join(f"c.{field}" for field in included_fields)

    @override
    def _lambda_parser(self, node: ast.AST) -> Any:
        match node:
            case ast.Compare():
                if len(node.ops) > 1:
                    # Chain comparisons (e.g., a < b < c) become AND-ed comparisons
                    values = []
                    for idx in range(len(node.ops)):
                        if idx == 0:
                            values.append(
                                ast.Compare(
                                    left=node.left,
                                    ops=[node.ops[idx]],
                                    comparators=[node.comparators[idx]],
                                )
                            )
                        else:
                            values.append(
                                ast.Compare(
                                    left=node.comparators[idx - 1],
                                    ops=[node.ops[idx]],
                                    comparators=[node.comparators[idx]],
                                )
                            )
                    return "(" + " AND ".join([self._lambda_parser(v) for v in values]) + ")"
                left = self._lambda_parser(node.left)
                right = self._lambda_parser(node.comparators[0])
                op = node.ops[0]
                match op:
                    case ast.In():
                        # Cosmos DB: ARRAY_CONTAINS(right, left)
                        return f"ARRAY_CONTAINS({right}, {left})"
                    case ast.NotIn():
                        return f"NOT ARRAY_CONTAINS({right}, {left})"
                    case ast.Eq():
                        return f"{left} = {right}"
                    case ast.NotEq():
                        return f"{left} != {right}"
                    case ast.Gt():
                        return f"{left} > {right}"
                    case ast.GtE():
                        return f"{left} >= {right}"
                    case ast.Lt():
                        return f"{left} < {right}"
                    case ast.LtE():
                        return f"{left} <= {right}"
                raise NotImplementedError(f"Unsupported operator: {type(op)}")
            case ast.BoolOp():
                op_str = "AND" if isinstance(node.op, ast.And) else "OR"
                return "(" + f" {op_str} ".join([self._lambda_parser(v) for v in node.values]) + ")"
            case ast.UnaryOp():
                match node.op:
                    case ast.Not():
                        return f"NOT ({self._lambda_parser(node.operand)})"
                    case ast.UAdd():
                        return f"+{self._lambda_parser(node.operand)}"
                    case ast.USub():
                        return f"-{self._lambda_parser(node.operand)}"
                    case ast.Invert():
                        raise NotImplementedError("Invert operation is not supported.")
                raise NotImplementedError(f"Unsupported unary operator: {type(node.op)}")
            case ast.Attribute():
                # Cosmos DB: c.field_name
                if node.attr not in self.definition.storage_names:
                    raise VectorStoreOperationException(
                        f"Field '{node.attr}' not in data model (storage property names are used)."
                    )
                return f"c.{node.attr}"
            case ast.Name():
                # Could be a variable or constant; not supported
                raise NotImplementedError("Constants or variables are not supported, use a value or attribute.")
            case ast.Constant():
                # Quote strings, leave numbers as is
                if isinstance(node.value, str):
                    return "'" + node.value.replace("'", "''") + "'"
                if isinstance(node.value, (float, int)):
                    return str(node.value)
                if node.value is None:
                    return "null"
                raise NotImplementedError(f"Unsupported constant type: {type(node.value)}")
        raise NotImplementedError(f"Unsupported AST node: {type(node)}")

    @override
    def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]:
        return result

    @override
    def _get_score_from_result(self, result: dict[str, Any]) -> float | None:
        return result.get(NOSQL_SCORE_PROPERTY_NAME)

    @override
    def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]:
        serialized_records = []

        key_field_name = self.definition.key_name
        for record in records:
            serialized_record = {**record, COSMOS_ITEM_ID_PROPERTY_NAME: record[key_field_name]}
            if key_field_name != COSMOS_ITEM_ID_PROPERTY_NAME:
                # Remove the key field from the serialized record
                serialized_record.pop(key_field_name, None)

            serialized_records.append(serialized_record)

        return serialized_records

    @override
    def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]:
        deserialized_records = []

        key_field_name = self.definition.key_name
        for record in records:
            if COSMOS_ITEM_ID_PROPERTY_NAME not in record:
                raise VectorStoreModelDeserializationException(
                    f"The record does not have the {COSMOS_ITEM_ID_PROPERTY_NAME} property."
                )

            deserialized_record = {**record, key_field_name: record[COSMOS_ITEM_ID_PROPERTY_NAME]}
            if key_field_name != COSMOS_ITEM_ID_PROPERTY_NAME:
                # Remove the id property from the deserialized record
                deserialized_record.pop(COSMOS_ITEM_ID_PROPERTY_NAME, None)

            deserialized_records.append(deserialized_record)

        return deserialized_records

    @override
    async def ensure_collection_exists(self, **kwargs) -> None:
        indexing_policy = kwargs.pop("indexing_policy", _create_default_indexing_policy_nosql(self.definition))
        vector_embedding_policy = kwargs.pop(
            "vector_embedding_policy", _create_default_vector_embedding_policy(self.definition)
        )
        database_proxy = await self._get_database_proxy(**kwargs)
        try:
            await database_proxy.create_container_if_not_exists(
                id=self.collection_name,
                partition_key=self.partition_key,
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy,
                **kwargs,
            )
        except CosmosHttpResponseError as e:
            raise VectorStoreOperationException("Failed to create container.") from e

    @override
    async def collection_exists(self, **kwargs) -> bool:
        container_proxy = await self._get_container_proxy(self.collection_name, **kwargs)
        try:
            await container_proxy.read(**kwargs)
            return True
        except CosmosHttpResponseError:
            return False

    @override
    async def ensure_collection_deleted(self, **kwargs) -> None:
        database_proxy = await self._get_database_proxy(**kwargs)
        try:
            await database_proxy.delete_container(self.collection_name)
        except Exception as e:
            raise VectorStoreOperationException("Container could not be deleted.") from e

    @override
    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        """Exit the context manager."""
        if self.managed_client:
            await self.cosmos_client.close()


# region: NoSQL Store


@release_candidate
class CosmosNoSqlStore(CosmosNoSqlBase, VectorStore):
    """A VectorStore implementation that uses Azure CosmosDB NoSQL as the backend storage."""

    def __init__(
        self,
        url: str | None = None,
        key: str | None = None,
        database_name: str | None = None,
        cosmos_client: CosmosClient | None = None,
        create_database: bool = False,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        credential: AsyncTokenCredential | None = None,
    ):
        """Initialize the CosmosNoSqlStore.

        Args:
            url: The URL of the Azure Cosmos DB NoSQL account. Defaults to None.
            key: The key of the Azure Cosmos DB NoSQL account. Defaults to None.
            database_name: The name of the database. The database may not exist yet. If it does not exist,
                                 it will be created when the first collection is created. Defaults to None.
            cosmos_client: The custom Azure Cosmos DB NoSQL client whose lifetime is managed by the user.
                                          Defaults to None.
            create_database: If True, the database will be created if it does not exist.
                                    Defaults to False.
            embedding_generator: The embedding generator to use for generating embeddings.
            env_file_path: The path to the .env file. Defaults to None.
            env_file_encoding: The encoding of the .env file. Defaults to None.
            credential: The credential to use for authentication to Azure Cosmos DB NoSQL.
        """
        super().__init__(
            url=url,
            key=key,
            database_name=database_name,
            cosmos_client=cosmos_client,
            create_database=create_database,
            embedding_generator=embedding_generator,
            env_file_path=env_file_path,
            env_file_encoding=env_file_encoding,
            credential=credential,
        )

    @override
    def get_collection(
        self,
        record_type: type[TModel],
        *,
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        **kwargs: Any,
    ) -> CosmosNoSqlCollection:
        return CosmosNoSqlCollection(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            database_name=self.database_name,
            embedding_generator=embedding_generator or self.embedding_generator,
            url=str(self.cosmos_db_nosql_settings.url),
            key=self.cosmos_db_nosql_settings.key.get_secret_value() if self.cosmos_db_nosql_settings.key else None,
            cosmos_client=self.cosmos_client,
            partition_key=None,
            create_database=self.create_database,
            env_file_path=None,
            env_file_encoding=None,
            **kwargs,
        )

    @override
    async def list_collection_names(self, **kwargs) -> Sequence[str]:
        try:
            database = await self._get_database_proxy()
            containers = database.list_containers()
            return [container["id"] async for container in containers]
        except Exception as e:
            raise VectorStoreOperationException("Failed to list collection names.") from e

    @override
    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        """Exit the context manager."""
        if self.managed_client:
            await self.cosmos_client.close()