[WIP][EGS]: Add azure-cosmos-ai companion package with Azure OpenAI embedding provider

sdk/cosmos/azure-cosmos-ai/CHANGELOG.md

@@ -0,0 +1,14 @@
+# Release History
+
+## 1.0.0b1 (Unreleased)
+
+### Features Added
+
+- Initial preview release of the `azure-cosmos-ai` package, a companion to `azure-cosmos`.
+- Added `azure.cosmos.ai.AzureOpenAIEmbeddingProvider` (sync) and `azure.cosmos.ai.aio.AzureOpenAIEmbeddingProvider` (async): the default Azure OpenAI implementation of the `EmbeddingProvider` Protocol introduced in `azure-cosmos`.
+
+### Breaking Changes
+
+### Bugs Fixed
+
+### Other Changes

sdk/cosmos/azure-cosmos-ai/LICENSE

@@ -0,0 +1,21 @@
+Copyright (c) Microsoft Corporation.
+
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sdk/cosmos/azure-cosmos-ai/MANIFEST.in

@@ -0,0 +1,4 @@
+recursive-include tests *.py
+include *.md
+include LICENSE
+include azure/cosmos/ai/py.typed

sdk/cosmos/azure-cosmos-ai/README.md

@@ -0,0 +1,83 @@
+# Azure Cosmos DB AI extensions for Python
+
+`azure-cosmos-ai` is a companion package to [`azure-cosmos`](https://pypi.org/project/azure-cosmos/) that provides AI-related extensions for the Azure Cosmos DB SDK.
+
+It ships the default Azure OpenAI implementation of the `EmbeddingProvider` Protocol introduced in `azure-cosmos` 4.16.0b3, used by the SDK to generate vector embeddings for `GenerateEmbeddings(...)` query expressions.
+
+## Getting started
+
+### Install the package
+
+```bash
+pip install azure-cosmos-ai
+```
+
+### Prerequisites
+
+- Python 3.9 or later
+- An Azure subscription
+- An existing Azure Cosmos DB for NoSQL account
+- An Azure OpenAI resource with an embeddings deployment (e.g. `text-embedding-3-small`)
+
+## Key concepts
+
+The provider stores **only the credential**. Endpoint, deployment name, and dimensions are read from the container's `vectorEmbeddingPolicy.embeddingSource` and forwarded to the provider by the Cosmos SDK at query time. This keeps the policy as the single source of truth.
+
+## Examples
+
+### API key (sync)
+
+```python
+from azure.cosmos import CosmosClient
+from azure.cosmos.ai import AzureOpenAIEmbeddingProvider
+
+provider = AzureOpenAIEmbeddingProvider(credential="<aoai-api-key>")
+
+client = CosmosClient(
+    url="https://my-cosmos.documents.azure.com:443/",
+    credential="<cosmos-key>",
+    embedding_provider=provider,
+)
+```
+
+### Entra — shared credential (recommended)
+
+Pass the same `TokenCredential` to `CosmosClient` (for Cosmos RBAC) and to the
+provider (for Azure OpenAI). One identity covers both services.
+
+```python
+from azure.cosmos.aio import CosmosClient
+from azure.cosmos.ai.aio import AzureOpenAIEmbeddingProvider
+from azure.identity.aio import DefaultAzureCredential
+
+async with DefaultAzureCredential() as cred:
+    async with AzureOpenAIEmbeddingProvider(credential=cred) as provider:
+        async with CosmosClient(
+            url="https://my-cosmos.documents.azure.com:443/",
+            credential=cred,
+            embedding_provider=provider,
+        ) as client:
+            ...
+```
+
+### Supported credential types
+
+| Type                                           | Auth mode |
+|------------------------------------------------|-----------|
+| `str`                                          | Azure OpenAI API key |
+| `azure.core.credentials.AzureKeyCredential`    | Azure OpenAI API key |
+| `azure.core.credentials.TokenCredential` (sync) / `azure.core.credentials_async.AsyncTokenCredential` (async) | Entra (RBAC) |
+
+## Troubleshooting
+
+The provider deliberately does not wrap exceptions thrown by the underlying
+[`openai`](https://pypi.org/project/openai/) client (e.g. `openai.BadRequestError`,
+`openai.AuthenticationError`, `openai.RateLimitError`, `openai.APIConnectionError`).
+Inputs that exceed the model's context length surface as `openai.BadRequestError`
+with code `context_length_exceeded`.
+
+Retries are handled by the `openai` SDK; this provider adds no extra retry policy.
+
+## Contributing
+
+This project welcomes contributions and suggestions. See the [Azure SDK for Python contributing guide](https://github.com/Azure/azure-sdk-for-python/blob/main/CONTRIBUTING.md) for details.

sdk/cosmos/azure-cosmos-ai/azure/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore

sdk/cosmos/azure-cosmos-ai/azure/cosmos/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore

sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py

@@ -0,0 +1,27 @@
+# The MIT License (MIT)
+# Copyright (c) 2023 Microsoft Corporation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from ._version import VERSION
+from ._azure_openai_provider import AzureOpenAIEmbeddingProvider
+
+
+__version__ = VERSION
+__all__ = ["AzureOpenAIEmbeddingProvider"]

sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py

@@ -0,0 +1,198 @@
+# The MIT License (MIT)
+# Copyright (c) 2023 Microsoft Corporation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Synchronous Azure OpenAI implementation of the EmbeddingProvider Protocol."""
+
+import inspect
+import time
+from typing import Any, Dict, Mapping, Optional, Sequence, Union
+
+from azure.core.credentials import AzureKeyCredential, TokenCredential
+from azure.cosmos import EmbeddingResult
+from azure.identity import get_bearer_token_provider
+from openai import AzureOpenAI
+
+_AZURE_OPENAI_API_VERSION = "2024-10-21"
+_COGNITIVE_SERVICES_SCOPE = "https://cognitiveservices.azure.com/.default"
+
+
+class AzureOpenAIEmbeddingProvider:
+    """Default Azure OpenAI implementation of the
+    :class:`azure.cosmos.EmbeddingProvider` Protocol.
+
+    The provider only stores the credential. Endpoint, deployment name, and
+    dimensions are read from the container's ``vectorEmbeddingPolicy`` and
+    forwarded to :meth:`generate_embeddings` by the Cosmos SDK at query time.
+
+    :param credential: One of:
+
+        * ``str`` – Azure OpenAI API key.
+        * :class:`~azure.core.credentials.AzureKeyCredential` – Azure OpenAI API key.
+        * :class:`~azure.core.credentials.TokenCredential` – Entra (RBAC). Pass the
+          same credential you use with :class:`~azure.cosmos.CosmosClient` to share
+          one identity across both services.
+    :type credential: str or
+        ~azure.core.credentials.AzureKeyCredential or
+        ~azure.core.credentials.TokenCredential
+    :keyword str api_version: Azure OpenAI REST API version. Defaults to
+        ``"2024-10-21"`` (the GA version when this package shipped). Override to
+        access newer model features without waiting for a new release of this
+        package.
+    :keyword openai_client_kwargs: Additional keyword arguments forwarded
+        verbatim to :class:`openai.AzureOpenAI` (e.g. ``timeout``, ``max_retries``,
+        ``http_client``, ``default_headers``, ``user``). Keys that this provider
+        controls (``azure_endpoint``, ``api_version``, ``api_key``,
+        ``azure_ad_token_provider``) are not overridable through this mapping.
+    :paramtype openai_client_kwargs: ~typing.Mapping[str, ~typing.Any] or None
+    """
+
+    def __init__(
+        self,
+        credential: Union[str, AzureKeyCredential, TokenCredential],
+        *,
+        api_version: str = _AZURE_OPENAI_API_VERSION,
+        openai_client_kwargs: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        if isinstance(credential, (str, AzureKeyCredential)):
+            pass
+        elif _is_token_credential(credential):
+            pass
+        elif _is_async_token_credential(credential):
+            raise TypeError(
+                "Synchronous AzureOpenAIEmbeddingProvider received an async "
+                f"credential ({type(credential).__name__}). Either use "
+                "azure.cosmos.ai.aio.AzureOpenAIEmbeddingProvider instead, or "
+                "pass a synchronous TokenCredential such as "
+                "azure.identity.DefaultAzureCredential."
+            )
+        else:
+            raise TypeError(
+                "credential must be a str, AzureKeyCredential, or synchronous "
+                f"TokenCredential; got {type(credential).__name__}"
+            )
+
+        self._credential = credential
+        self._api_version = api_version
+        self._openai_client_kwargs: Dict[str, Any] = dict(openai_client_kwargs or {})
+        self._clients: Dict[str, AzureOpenAI] = {}
+
+    def generate_embeddings(
+        self,
+        texts: Sequence[str],
+        *,
+        endpoint: str,
+        deployment_name: str,
+        dimensions: int,
+        **kwargs: Any,
+    ) -> EmbeddingResult:
+        """Generate embeddings for ``texts`` using Azure OpenAI.
+
+        :param texts: Input strings.
+        :type texts: ~typing.Sequence[str]
+        :keyword str endpoint: Azure OpenAI endpoint
+            (from ``vectorEmbeddingPolicy.embeddingSource.endpoint``).
+        :keyword str deployment_name: Azure OpenAI deployment name
+            (from ``vectorEmbeddingPolicy.embeddingSource.deploymentName``).
+        :keyword int dimensions: Embedding dimensions
+            (from ``vectorEmbeddingPolicy.dimensions``).
+        :keyword Any kwargs: Reserved for forward compatibility with future
+            Cosmos SDK additions. Currently, no per-call kwargs are forwarded to
+            the underlying ``openai`` call; use ``openai_client_kwargs`` on the
+            constructor (e.g. ``timeout``, ``max_retries``) to configure the
+            underlying client.
+        :returns: Vectors in the same order as ``texts``, plus token usage and
+            measured latency.
+        :rtype: ~azure.cosmos.EmbeddingResult
+        """
+        if not texts:
+            return EmbeddingResult(vectors=[], total_tokens=0, latency=None)
+
+        client = self._get_or_create_client(endpoint)
+        start = time.perf_counter()
+        response = client.embeddings.create(
+            input=list(texts),
+            model=deployment_name,
+            dimensions=dimensions,
+        )
+        latency = time.perf_counter() - start
+        total_tokens: Optional[int] = response.usage.total_tokens if response.usage else None
+        return EmbeddingResult(
+            vectors=[item.embedding for item in response.data],
+            total_tokens=total_tokens,
+            latency=latency,
+        )
+
+    def close(self) -> None:
+        """Close every cached underlying Azure OpenAI client and clear the cache."""
+        clients = list(self._clients.values())
+        self._clients.clear()
+        for client in clients:
+            try:
+                client.close()
+            except Exception:  # pylint: disable=broad-except
+                pass
+
+    def __enter__(self) -> "AzureOpenAIEmbeddingProvider":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+    def _get_or_create_client(self, endpoint: str) -> AzureOpenAI:
+        key = endpoint.rstrip("/")
+        client = self._clients.get(key)
+        if client is None:
+            client = self._build_client(key)
+            self._clients[key] = client
+        return client
+
+    def _build_client(self, endpoint: str) -> AzureOpenAI:
+        # User-supplied kwargs go first so our explicit args win on collision.
+        common: Dict[str, Any] = dict(self._openai_client_kwargs)
+        common.update(azure_endpoint=endpoint, api_version=self._api_version)
+        if isinstance(self._credential, str):
+            return AzureOpenAI(api_key=self._credential, **common)
+        if isinstance(self._credential, AzureKeyCredential):
+            return AzureOpenAI(api_key=self._credential.key, **common)
+        token_provider = get_bearer_token_provider(self._credential, _COGNITIVE_SERVICES_SCOPE)
+        return AzureOpenAI(azure_ad_token_provider=token_provider, **common)
+
+
+def _is_token_credential(obj: Any) -> bool:
+    """Duck-type check for a *synchronous* TokenCredential.
+
+    Accepts any object that exposes a non-coroutine, callable ``get_token``.
+    Async credentials (where ``get_token`` is a coroutine function) are rejected
+    so the mismatch is caught at ``__init__`` instead of failing deep inside
+    ``openai`` with a confusing ``coroutine`` error.
+    """
+    get_token = getattr(obj, "get_token", None)
+    return callable(get_token) and not inspect.iscoroutinefunction(get_token)
+
+
+def _is_async_token_credential(obj: Any) -> bool:
+    """Duck-type check for an *asynchronous* TokenCredential.
+
+    Used only to produce an actionable error message when an async credential is
+    accidentally passed to the sync provider.
+    """
+    get_token = getattr(obj, "get_token", None)
+    return callable(get_token) and inspect.iscoroutinefunction(get_token)

sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py

@@ -0,0 +1,22 @@
+# The MIT License (MIT)
+# Copyright (c) 2023 Microsoft Corporation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+VERSION = "1.0.0b1"