From b4b41d7f0e4ba3d57c10a56b2059e2fa7fad83db Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Thu, 14 May 2026 15:21:10 -0700 Subject: [PATCH 1/7] Adding the new packlage scaffoling --- sdk/cosmos/azure-cosmos-ai/CHANGELOG.md | 13 +++ sdk/cosmos/azure-cosmos-ai/LICENSE | 21 +++++ sdk/cosmos/azure-cosmos-ai/MANIFEST.in | 7 ++ sdk/cosmos/azure-cosmos-ai/README.md | 24 ++++++ sdk/cosmos/azure-cosmos-ai/azure/__init__.py | 1 + .../azure-cosmos-ai/azure/cosmos/__init__.py | 1 + .../azure/cosmos/ai/__init__.py | 26 ++++++ .../azure/cosmos/ai/_version.py | 22 +++++ .../azure-cosmos-ai/azure/cosmos/ai/py.typed | 0 .../azure-cosmos-ai/dev_requirements.txt | 4 + sdk/cosmos/azure-cosmos-ai/pyproject.toml | 7 ++ sdk/cosmos/azure-cosmos-ai/sdk_packaging.toml | 2 + sdk/cosmos/azure-cosmos-ai/setup.py | 83 +++++++++++++++++++ sdk/cosmos/ci.yml | 2 + 14 files changed, 213 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos-ai/CHANGELOG.md create mode 100644 sdk/cosmos/azure-cosmos-ai/LICENSE create mode 100644 sdk/cosmos/azure-cosmos-ai/MANIFEST.in create mode 100644 sdk/cosmos/azure-cosmos-ai/README.md create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/__init__.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/__init__.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/py.typed create mode 100644 sdk/cosmos/azure-cosmos-ai/dev_requirements.txt create mode 100644 sdk/cosmos/azure-cosmos-ai/pyproject.toml create mode 100644 sdk/cosmos/azure-cosmos-ai/sdk_packaging.toml create mode 100644 sdk/cosmos/azure-cosmos-ai/setup.py diff --git a/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md b/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md new file mode 100644 index 000000000000..dff65ab45dc1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md @@ -0,0 +1,13 @@ +# Release History + +## 1.0.0b1 (Unreleased) + +### Features Added + +- Initial preview release of the `azure-cosmos-ai` package, a companion to `azure-cosmos` that will host AI-related extensions including the default Azure OpenAI implementation of the `EmbeddingProvider` Protocol. + +### Breaking Changes + +### Bugs Fixed + +### Other Changes diff --git a/sdk/cosmos/azure-cosmos-ai/LICENSE b/sdk/cosmos/azure-cosmos-ai/LICENSE new file mode 100644 index 000000000000..63447fd8bbbf --- /dev/null +++ b/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. \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-ai/MANIFEST.in b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in new file mode 100644 index 000000000000..f2097fcd6842 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in @@ -0,0 +1,7 @@ +recursive-include samples *.py +recursive-include tests *.py +include *.md +include LICENSE +include azure/__init__.py +include azure/cosmos/__init__.py +include azure/cosmos/ai/py.typed diff --git a/sdk/cosmos/azure-cosmos-ai/README.md b/sdk/cosmos/azure-cosmos-ai/README.md new file mode 100644 index 000000000000..830e0f26afad --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/README.md @@ -0,0 +1,24 @@ +# 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 will host 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 + +## 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. diff --git a/sdk/cosmos/azure-cosmos-ai/azure/__init__.py b/sdk/cosmos/azure-cosmos-ai/azure/__init__.py new file mode 100644 index 000000000000..d55ccad1f573 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/__init__.py @@ -0,0 +1 @@ +__path__ = __import__("pkgutil").extend_path(__path__, __name__) # type: ignore diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/__init__.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/__init__.py new file mode 100644 index 000000000000..d55ccad1f573 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/__init__.py @@ -0,0 +1 @@ +__path__ = __import__("pkgutil").extend_path(__path__, __name__) # type: ignore diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py new file mode 100644 index 000000000000..9d0d8e7bcc93 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py @@ -0,0 +1,26 @@ +# The MIT License (MIT) +# Copyright (c) 2014 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 + + +__version__ = VERSION +__all__: list = [] diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py new file mode 100644 index 000000000000..0d405533a861 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py @@ -0,0 +1,22 @@ +# The MIT License (MIT) +# Copyright (c) 2014 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" diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/py.typed b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/py.typed new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt b/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt new file mode 100644 index 000000000000..5357c3beb6a7 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt @@ -0,0 +1,4 @@ +../azure-cosmos +../../core/azure-core +openai>=1.0.0 +-e ../../../eng/tools/azure-sdk-tools diff --git a/sdk/cosmos/azure-cosmos-ai/pyproject.toml b/sdk/cosmos/azure-cosmos-ai/pyproject.toml new file mode 100644 index 000000000000..52c35fbc61c2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/pyproject.toml @@ -0,0 +1,7 @@ +[tool.azure-sdk-build] +mypy = true +pyright = false +pylint = true + +[tool.azure-sdk-conda] +in_bundle = false diff --git a/sdk/cosmos/azure-cosmos-ai/sdk_packaging.toml b/sdk/cosmos/azure-cosmos-ai/sdk_packaging.toml new file mode 100644 index 000000000000..901bc8ccbfa6 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/sdk_packaging.toml @@ -0,0 +1,2 @@ +[packaging] +auto_update = false diff --git a/sdk/cosmos/azure-cosmos-ai/setup.py b/sdk/cosmos/azure-cosmos-ai/setup.py new file mode 100644 index 000000000000..47c903385f1f --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/setup.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python + +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +# pylint:disable=missing-docstring + +import re +import os +from io import open +from setuptools import find_packages, setup + +# Change the PACKAGE_NAME only to change folder and different name +PACKAGE_NAME = "azure-cosmos-ai" +PACKAGE_PPRINT_NAME = "Cosmos AI" + +# a-b-c => a/b/c +PACKAGE_FOLDER_PATH = PACKAGE_NAME.replace("-", "/") +# a-b-c => a.b.c +NAMESPACE_NAME = PACKAGE_NAME.replace("-", ".") + +# Version extraction inspired from 'requests' +with open(os.path.join(PACKAGE_FOLDER_PATH, '_version.py'), 'r') as fd: + version = re.search(r'^VERSION\s*=\s*[\'"]([^\'"]*)[\'"]', + fd.read(), re.MULTILINE).group(1) + +if not version: + raise RuntimeError("Cannot find version information") + +with open("README.md", encoding="utf-8") as f: + readme = f.read() +with open("CHANGELOG.md", encoding="utf-8") as f: + changelog = f.read() + +exclude_packages = [ + "tests", + "samples", + # Exclude packages that will be covered by PEP420 or nspkg + "azure", +] + +setup( + name=PACKAGE_NAME, + version=version, + include_package_data=True, + description="Microsoft Azure {} Extensions for Python".format(PACKAGE_PPRINT_NAME), + long_description=readme + "\n\n" + changelog, + long_description_content_type="text/markdown", + license="MIT License", + author="Microsoft Corporation", + author_email="askdocdb@microsoft.com", + maintainer="Microsoft", + maintainer_email="askdocdb@microsoft.com", + url="https://github.com/Azure/azure-sdk-for-python", + keywords="azure, azure sdk", + classifiers=[ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Natural Language :: English", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "License :: OSI Approved :: MIT License", + ], + zip_safe=False, + packages=find_packages(exclude=exclude_packages), + package_data={ + "azure.cosmos.ai": ["py.typed"], + }, + python_requires=">=3.9", + install_requires=[ + "azure-cosmos>=4.16.0b3", + "azure-core>=1.30.0", + "openai>=1.0.0", + ], +) diff --git a/sdk/cosmos/ci.yml b/sdk/cosmos/ci.yml index 9565a7aeb76e..6104ead306c7 100644 --- a/sdk/cosmos/ci.yml +++ b/sdk/cosmos/ci.yml @@ -31,5 +31,7 @@ extends: Artifacts: - name: azure-cosmos safeName: azurecosmos + - name: azure-cosmos-ai + safeName: azurecosmosai - name: azure-mgmt-cosmosdb safeName: azuremgmtcosmosdb From 3b8c28accde07f5b26f96407d35411ff29f063a9 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Sat, 16 May 2026 16:12:39 -0700 Subject: [PATCH 2/7] Adding the AOAI provider implementation --- sdk/cosmos/azure-cosmos-ai/CHANGELOG.md | 3 +- sdk/cosmos/azure-cosmos-ai/README.md | 63 ++++- .../azure/cosmos/ai/__init__.py | 5 +- .../azure/cosmos/ai/_azure_openai_provider.py | 156 +++++++++++++ .../azure/cosmos/ai/_version.py | 2 +- .../azure/cosmos/ai/aio/__init__.py | 24 ++ .../cosmos/ai/aio/_azure_openai_provider.py | 156 +++++++++++++ .../azure-cosmos-ai/dev_requirements.txt | 1 + .../samples/sample_embedding_provider.py | 90 ++++++++ .../sample_embedding_provider_async.py | 87 +++++++ sdk/cosmos/azure-cosmos-ai/setup.py | 26 ++- sdk/cosmos/azure-cosmos-ai/tests/conftest.py | 43 ++++ .../tests/test_azure_openai_provider.py | 203 +++++++++++++++++ .../tests/test_azure_openai_provider_async.py | 215 ++++++++++++++++++ 14 files changed, 1063 insertions(+), 11 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/__init__.py create mode 100644 sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py create mode 100644 sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py create mode 100644 sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py create mode 100644 sdk/cosmos/azure-cosmos-ai/tests/conftest.py create mode 100644 sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py create mode 100644 sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py diff --git a/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md b/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md index dff65ab45dc1..938f5025b033 100644 --- a/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-ai/CHANGELOG.md @@ -4,7 +4,8 @@ ### Features Added -- Initial preview release of the `azure-cosmos-ai` package, a companion to `azure-cosmos` that will host AI-related extensions including the default Azure OpenAI implementation of the `EmbeddingProvider` Protocol. +- 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 diff --git a/sdk/cosmos/azure-cosmos-ai/README.md b/sdk/cosmos/azure-cosmos-ai/README.md index 830e0f26afad..c854583dab3d 100644 --- a/sdk/cosmos/azure-cosmos-ai/README.md +++ b/sdk/cosmos/azure-cosmos-ai/README.md @@ -2,7 +2,7 @@ `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 will host 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. +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 @@ -17,7 +17,66 @@ pip install azure-cosmos-ai - Python 3.9 or later - An Azure subscription - An existing Azure Cosmos DB for NoSQL account -- An Azure OpenAI resource with an embeddings deployment +- 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="") + +client = CosmosClient( + url="https://my-cosmos.documents.azure.com:443/", + credential="", + 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 diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py index 9d0d8e7bcc93..d4612a0e579c 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/__init__.py @@ -1,5 +1,5 @@ # The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation +# 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 @@ -20,7 +20,8 @@ # SOFTWARE. from ._version import VERSION +from ._azure_openai_provider import AzureOpenAIEmbeddingProvider __version__ = VERSION -__all__: list = [] +__all__ = ["AzureOpenAIEmbeddingProvider"] diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py new file mode 100644 index 000000000000..8ad4c795d7c1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py @@ -0,0 +1,156 @@ +# 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.""" + +from typing import Any, Dict, 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 + """ + + def __init__( + self, + credential: Union[str, AzureKeyCredential, TokenCredential], + **kwargs: Any, # pylint: disable=unused-argument + ) -> None: + if not isinstance(credential, (str, AzureKeyCredential)) and not _is_token_credential(credential): + raise TypeError( + "credential must be a str, AzureKeyCredential, or TokenCredential; " + f"got {type(credential).__name__}" + ) + self._credential = credential + self._api_version = _AZURE_OPENAI_API_VERSION + self._clients: Dict[str, AzureOpenAI] = {} + + def generate_embeddings( + self, + texts: Sequence[str], + *, + endpoint: str, + deployment_name: str, + dimensions: int, + **kwargs: Any, # pylint: disable=unused-argument + ) -> 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``). + :returns: Vectors in the same order as ``texts``, plus token usage. + :rtype: ~azure.cosmos.EmbeddingResult + """ + if not texts: + return EmbeddingResult(vectors=[], total_tokens=0) + + client = self._get_or_create_client(endpoint) + response = client.embeddings.create( + input=list(texts), + model=deployment_name, + dimensions=dimensions, + ) + 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, + ) + + def close(self) -> None: + """Close every cached underlying Azure OpenAI client and clear the cache.""" + for client in self._clients.values(): + try: + client.close() + except Exception: # pylint: disable=broad-except + pass + self._clients.clear() + + 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: + if isinstance(self._credential, str): + return AzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + api_key=self._credential, + ) + if isinstance(self._credential, AzureKeyCredential): + return AzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + api_key=self._credential.key, + ) + token_provider = get_bearer_token_provider(self._credential, _COGNITIVE_SERVICES_SCOPE) + return AzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + azure_ad_token_provider=token_provider, + ) + + +def _is_token_credential(obj: Any) -> bool: + """Duck-type check for a sync TokenCredential. + + Avoids ``isinstance(obj, TokenCredential)`` because ``TokenCredential`` is a + ``Protocol`` in some azure-core versions and not always runtime_checkable. + """ + return callable(getattr(obj, "get_token", None)) diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py index 0d405533a861..36cd1059c7bb 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_version.py @@ -1,5 +1,5 @@ # The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation +# 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 diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/__init__.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/__init__.py new file mode 100644 index 000000000000..db33676aeb88 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/__init__.py @@ -0,0 +1,24 @@ +# 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 ._azure_openai_provider import AzureOpenAIEmbeddingProvider + +__all__ = ["AzureOpenAIEmbeddingProvider"] diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py new file mode 100644 index 000000000000..d8833308c404 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py @@ -0,0 +1,156 @@ +# 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. + +"""Asynchronous Azure OpenAI implementation of the EmbeddingProvider Protocol.""" + +from typing import Any, Dict, Optional, Sequence, Union + +from azure.core.credentials import AzureKeyCredential +from azure.core.credentials_async import AsyncTokenCredential +from azure.cosmos import EmbeddingResult +from azure.identity.aio import get_bearer_token_provider +from openai import AsyncAzureOpenAI + +_AZURE_OPENAI_API_VERSION = "2024-10-21" +_COGNITIVE_SERVICES_SCOPE = "https://cognitiveservices.azure.com/.default" + + +class AzureOpenAIEmbeddingProvider: + """Async default Azure OpenAI implementation of the + :class:`azure.cosmos.aio.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_async.AsyncTokenCredential` – Entra (RBAC). + Pass the same credential you use with + :class:`~azure.cosmos.aio.CosmosClient` to share one identity across both + services. + :type credential: str or + ~azure.core.credentials.AzureKeyCredential or + ~azure.core.credentials_async.AsyncTokenCredential + """ + + def __init__( + self, + credential: Union[str, AzureKeyCredential, AsyncTokenCredential], + **kwargs: Any, # pylint: disable=unused-argument + ) -> None: + if not isinstance(credential, (str, AzureKeyCredential)) and not _is_async_token_credential(credential): + raise TypeError( + "credential must be a str, AzureKeyCredential, or AsyncTokenCredential; " + f"got {type(credential).__name__}" + ) + self._credential = credential + self._api_version = _AZURE_OPENAI_API_VERSION + self._clients: Dict[str, AsyncAzureOpenAI] = {} + + async def generate_embeddings( + self, + texts: Sequence[str], + *, + endpoint: str, + deployment_name: str, + dimensions: int, + **kwargs: Any, # pylint: disable=unused-argument + ) -> EmbeddingResult: + """Generate embeddings for ``texts`` using Azure OpenAI. + + Safe to call concurrently from multiple coroutines. + + :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``). + :returns: Vectors in the same order as ``texts``, plus token usage. + :rtype: ~azure.cosmos.EmbeddingResult + """ + if not texts: + return EmbeddingResult(vectors=[], total_tokens=0) + + client = self._get_or_create_client(endpoint) + response = await client.embeddings.create( + input=list(texts), + model=deployment_name, + dimensions=dimensions, + ) + 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, + ) + + async def close(self) -> None: + """Close every cached underlying Azure OpenAI client and clear the cache.""" + for client in self._clients.values(): + try: + await client.close() + except Exception: # pylint: disable=broad-except + pass + self._clients.clear() + + async def __aenter__(self) -> "AzureOpenAIEmbeddingProvider": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + def _get_or_create_client(self, endpoint: str) -> AsyncAzureOpenAI: + 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) -> AsyncAzureOpenAI: + if isinstance(self._credential, str): + return AsyncAzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + api_key=self._credential, + ) + if isinstance(self._credential, AzureKeyCredential): + return AsyncAzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + api_key=self._credential.key, + ) + token_provider = get_bearer_token_provider(self._credential, _COGNITIVE_SERVICES_SCOPE) + return AsyncAzureOpenAI( + azure_endpoint=endpoint, + api_version=self._api_version, + azure_ad_token_provider=token_provider, + ) + + +def _is_async_token_credential(obj: Any) -> bool: + """Duck-type check for an async TokenCredential.""" + return callable(getattr(obj, "get_token", None)) diff --git a/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt b/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt index 5357c3beb6a7..fc31dc33fba7 100644 --- a/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt +++ b/sdk/cosmos/azure-cosmos-ai/dev_requirements.txt @@ -1,4 +1,5 @@ ../azure-cosmos ../../core/azure-core +../../identity/azure-identity openai>=1.0.0 -e ../../../eng/tools/azure-sdk-tools diff --git a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py new file mode 100644 index 000000000000..b53bebf2f3f8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py @@ -0,0 +1,90 @@ +# 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. + +"""Sample: use AzureOpenAIEmbeddingProvider with a synchronous CosmosClient. + +Demonstrates two credential modes: + +* Variant A — Azure OpenAI API key. +* Variant B — Entra (RBAC). The SAME ``DefaultAzureCredential`` is shared + between ``CosmosClient`` (Cosmos RBAC) and the embedding provider + (Azure OpenAI), so the user only needs one identity. + +Required environment variables: + +* ``COSMOS_ENDPOINT`` – e.g. ``https://my-cosmos.documents.azure.com:443/`` +* ``COSMOS_KEY`` – only for Variant A +* ``AOAI_API_KEY`` – only for Variant A + +Both samples assume a database ``samples-db`` with a container ``items`` +whose ``vectorEmbeddingPolicy.embeddingSource`` already points at the +Azure OpenAI endpoint and deployment you intend to use. +""" + +import os + +from azure.cosmos import CosmosClient +from azure.cosmos.ai import AzureOpenAIEmbeddingProvider +from azure.identity import DefaultAzureCredential + + +def variant_a_api_key() -> None: + cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] + cosmos_key = os.environ["COSMOS_KEY"] + aoai_api_key = os.environ["AOAI_API_KEY"] + + with AzureOpenAIEmbeddingProvider(credential=aoai_api_key) as provider: + client = CosmosClient( + url=cosmos_endpoint, + credential=cosmos_key, + embedding_provider=provider, + ) + _run_query(client) + + +def variant_b_shared_entra() -> None: + cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] + + cred = DefaultAzureCredential() + with AzureOpenAIEmbeddingProvider(credential=cred) as provider: + client = CosmosClient( + url=cosmos_endpoint, + credential=cred, + embedding_provider=provider, + ) + _run_query(client) + + +def _run_query(client: CosmosClient) -> None: + db = client.get_database_client("samples-db") + container = db.get_container_client("items") + + query = ( + "SELECT TOP 5 c.id, " + "VectorDistance(c.embedding, GenerateEmbeddings('healthcare research papers')) AS score " + "FROM c ORDER BY score" + ) + for item in container.query_items(query=query, enable_cross_partition_query=True): + print(item) + + +if __name__ == "__main__": + variant_a_api_key() diff --git a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py new file mode 100644 index 000000000000..5ed30e28685d --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py @@ -0,0 +1,87 @@ +# 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. + +"""Sample: use AzureOpenAIEmbeddingProvider with an async CosmosClient. + +Demonstrates two credential modes: + +* Variant A — Azure OpenAI API key. +* Variant B — Entra (RBAC). The SAME ``DefaultAzureCredential`` is shared + between ``CosmosClient`` (Cosmos RBAC) and the embedding provider + (Azure OpenAI), so the user only needs one identity. + +Required environment variables: + +* ``COSMOS_ENDPOINT`` – e.g. ``https://my-cosmos.documents.azure.com:443/`` +* ``COSMOS_KEY`` – only for Variant A +* ``AOAI_API_KEY`` – only for Variant A +""" + +import asyncio +import os + +from azure.cosmos.aio import CosmosClient +from azure.cosmos.ai.aio import AzureOpenAIEmbeddingProvider +from azure.identity.aio import DefaultAzureCredential + + +async def variant_a_api_key() -> None: + cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] + cosmos_key = os.environ["COSMOS_KEY"] + aoai_api_key = os.environ["AOAI_API_KEY"] + + async with AzureOpenAIEmbeddingProvider(credential=aoai_api_key) as provider: + async with CosmosClient( + url=cosmos_endpoint, + credential=cosmos_key, + embedding_provider=provider, + ) as client: + await _run_query(client) + + +async def variant_b_shared_entra() -> None: + cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] + + async with DefaultAzureCredential() as cred: + async with AzureOpenAIEmbeddingProvider(credential=cred) as provider: + async with CosmosClient( + url=cosmos_endpoint, + credential=cred, + embedding_provider=provider, + ) as client: + await _run_query(client) + + +async def _run_query(client: CosmosClient) -> None: + db = client.get_database_client("samples-db") + container = db.get_container_client("items") + + query = ( + "SELECT TOP 5 c.id, " + "VectorDistance(c.embedding, GenerateEmbeddings('healthcare research papers')) AS score " + "FROM c ORDER BY score" + ) + async for item in container.query_items(query=query): + print(item) + + +if __name__ == "__main__": + asyncio.run(variant_a_api_key()) diff --git a/sdk/cosmos/azure-cosmos-ai/setup.py b/sdk/cosmos/azure-cosmos-ai/setup.py index 47c903385f1f..f9bd9131c45b 100644 --- a/sdk/cosmos/azure-cosmos-ai/setup.py +++ b/sdk/cosmos/azure-cosmos-ai/setup.py @@ -1,10 +1,25 @@ #!/usr/bin/env python -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -# pylint:disable=missing-docstring +# 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. import re import os @@ -78,6 +93,7 @@ install_requires=[ "azure-cosmos>=4.16.0b3", "azure-core>=1.30.0", + "azure-identity>=1.19.0", "openai>=1.0.0", ], ) diff --git a/sdk/cosmos/azure-cosmos-ai/tests/conftest.py b/sdk/cosmos/azure-cosmos-ai/tests/conftest.py new file mode 100644 index 000000000000..c2ce42819bcb --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/tests/conftest.py @@ -0,0 +1,43 @@ +# 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. + +"""Test fixtures for azure-cosmos-ai. + +Until ``azure-cosmos`` 4.16.0b3 (PR #46902) is released, ``EmbeddingResult`` +isn't available from ``azure.cosmos`` in our checkout. Inject a minimal stub +on import so the provider modules — which ``from azure.cosmos import +EmbeddingResult`` at module load — work in CI/local dev. Once the dependency +ships, this stub becomes a no-op (the real class wins). +""" + +from dataclasses import dataclass +from typing import List, Optional + +import azure.cosmos as _cosmos + +if not hasattr(_cosmos, "EmbeddingResult"): + + @dataclass + class EmbeddingResult: # pylint: disable=too-few-public-methods + vectors: List[List[float]] + total_tokens: Optional[int] = None + + _cosmos.EmbeddingResult = EmbeddingResult # type: ignore[attr-defined] diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py new file mode 100644 index 000000000000..1b8cbea0a4b9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py @@ -0,0 +1,203 @@ +# 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. + +"""Unit tests for AzureOpenAIEmbeddingProvider (sync).""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from azure.core.credentials import AccessToken, AzureKeyCredential + +from azure.cosmos.ai import AzureOpenAIEmbeddingProvider + + +ENDPOINT = "https://example.com/" +ENDPOINT_KEY = "https://example.com" +DEPLOYMENT = "text-embedding-3-small" +DIMENSIONS = 1536 + + +class _FakeTokenCredential: + def __init__(self): + self.calls = [] + + def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + self.calls.append(scopes) + return AccessToken("fake-token", 9999999999) + + +def _fake_response(vectors, total_tokens=42): + return SimpleNamespace( + data=[SimpleNamespace(embedding=v) for v in vectors], + usage=SimpleNamespace(total_tokens=total_tokens) if total_tokens is not None else None, + ) + + +@pytest.fixture +def mock_aoai(): + """Patches AzureOpenAI inside the provider module.""" + with patch("azure.cosmos.ai._azure_openai_provider.AzureOpenAI") as cls: + instance = MagicMock(name="AzureOpenAIInstance") + cls.return_value = instance + instance.embeddings.create.return_value = _fake_response([[0.1, 0.2], [0.3, 0.4]]) + yield cls, instance + + +# ----- constructor / credential dispatch ----- + + +def test_init_accepts_str(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential="my-key") + + +def test_init_accepts_azure_key_credential(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) + + +def test_init_accepts_token_credential(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=_FakeTokenCredential()) + + +def test_init_rejects_unknown_credential(): + with pytest.raises(TypeError): + AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] + + +# ----- generate_embeddings ----- + + +def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai): + cls, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 + ) + + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( + ["a", "b", "c"], + endpoint=ENDPOINT, + deployment_name=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + + cls.assert_called_once() + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_key"] == "key" + assert ctor_kwargs["api_version"] == "2024-10-21" + + instance.embeddings.create.assert_called_once_with( + input=["a", "b", "c"], + model=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + + assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + assert result.total_tokens == 99 + + +def test_generate_embeddings_missing_usage_returns_none(mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) + + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.total_tokens is None + + +def test_empty_texts_short_circuits(mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( + [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.vectors == [] + assert result.total_tokens == 0 + cls.assert_not_called() + instance.embeddings.create.assert_not_called() + + +def test_exceptions_propagate(mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.side_effect = RuntimeError("boom") + + provider = AzureOpenAIEmbeddingProvider(credential="key") + with pytest.raises(RuntimeError, match="boom"): + provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + + +# ----- credential plumbing into AzureOpenAI ----- + + +def test_azure_key_credential_passed_as_api_key(mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_key"] == "aaa" + + +def test_token_credential_uses_bearer_token_provider(mock_aoai): + cls, _ = mock_aoai + cred = _FakeTokenCredential() + provider = AzureOpenAIEmbeddingProvider(credential=cred) + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert "api_key" not in ctor_kwargs + token_provider = ctor_kwargs["azure_ad_token_provider"] + token = token_provider() + assert token == "fake-token" + assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" + + +# ----- close / context manager ----- + + +def test_close_clears_cache(mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + provider.close() + instance.close.assert_called_once() + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert cls.call_count == 2 + + +def test_context_manager_closes(mock_aoai): + _, instance = mock_aoai + with AzureOpenAIEmbeddingProvider(credential="key") as provider: + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + instance.close.assert_called_once() diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py new file mode 100644 index 000000000000..d45cc59b5fcd --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py @@ -0,0 +1,215 @@ +# 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. + +"""Unit tests for AzureOpenAIEmbeddingProvider (async).""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from azure.core.credentials import AccessToken, AzureKeyCredential + +from azure.cosmos.ai.aio import AzureOpenAIEmbeddingProvider + + +ENDPOINT = "https://example.com/" +ENDPOINT_KEY = "https://example.com" +DEPLOYMENT = "text-embedding-3-small" +DIMENSIONS = 1536 + + +class _FakeAsyncTokenCredential: + def __init__(self): + self.calls = [] + + async def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + self.calls.append(scopes) + return AccessToken("fake-token", 9999999999) + + async def close(self): + pass + + +def _fake_response(vectors, total_tokens=42): + return SimpleNamespace( + data=[SimpleNamespace(embedding=v) for v in vectors], + usage=SimpleNamespace(total_tokens=total_tokens) if total_tokens is not None else None, + ) + + +@pytest.fixture +def mock_aoai(): + """Patches AsyncAzureOpenAI inside the async provider module.""" + with patch("azure.cosmos.ai.aio._azure_openai_provider.AsyncAzureOpenAI") as cls: + instance = MagicMock(name="AsyncAzureOpenAIInstance") + instance.embeddings.create = AsyncMock(return_value=_fake_response([[0.1, 0.2]])) + instance.close = AsyncMock() + cls.return_value = instance + yield cls, instance + + +# ----- constructor / credential dispatch ----- + + +def test_init_accepts_str(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential="my-key") + + +def test_init_accepts_azure_key_credential(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) + + +def test_init_accepts_async_token_credential(mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=_FakeAsyncTokenCredential()) + + +def test_init_rejects_unknown_credential(): + with pytest.raises(TypeError): + AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] + + +# ----- generate_embeddings ----- + + +@pytest.mark.asyncio +async def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai): + cls, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 + ) + + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( + ["a", "b", "c"], + endpoint=ENDPOINT, + deployment_name=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + + cls.assert_called_once() + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_key"] == "key" + assert ctor_kwargs["api_version"] == "2024-10-21" + + instance.embeddings.create.assert_awaited_once_with( + input=["a", "b", "c"], + model=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + + assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + assert result.total_tokens == 99 + + +@pytest.mark.asyncio +async def test_generate_embeddings_missing_usage_returns_none(mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) + + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.total_tokens is None + + +@pytest.mark.asyncio +async def test_empty_texts_short_circuits(mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( + [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.vectors == [] + assert result.total_tokens == 0 + cls.assert_not_called() + instance.embeddings.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_exceptions_propagate(mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.side_effect = RuntimeError("boom") + + provider = AzureOpenAIEmbeddingProvider(credential="key") + with pytest.raises(RuntimeError, match="boom"): + await provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + + +# ----- credential plumbing ----- + + +@pytest.mark.asyncio +async def test_azure_key_credential_passed_as_api_key(mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_key"] == "aaa" + + +@pytest.mark.asyncio +async def test_async_token_credential_uses_bearer_token_provider(mock_aoai): + cls, _ = mock_aoai + cred = _FakeAsyncTokenCredential() + provider = AzureOpenAIEmbeddingProvider(credential=cred) + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert "api_key" not in ctor_kwargs + token_provider = ctor_kwargs["azure_ad_token_provider"] + token = await token_provider() + assert token == "fake-token" + assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" + + +# ----- close / async context manager ----- + + +@pytest.mark.asyncio +async def test_close_clears_cache(mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + await provider.close() + instance.close.assert_awaited_once() + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert cls.call_count == 2 + + +@pytest.mark.asyncio +async def test_async_context_manager_closes(mock_aoai): + _, instance = mock_aoai + async with AzureOpenAIEmbeddingProvider(credential="key") as provider: + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + instance.close.assert_awaited_once() From 36fc41609679770a64e9fc793a22ed54a16f5b6a Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 19 May 2026 13:27:27 -0700 Subject: [PATCH 3/7] Adding open ai embedding provider implementation --- .../azure/cosmos/ai/_azure_openai_provider.py | 6 +++++- .../azure/cosmos/ai/aio/_azure_openai_provider.py | 6 +++++- sdk/cosmos/azure-cosmos-ai/tests/conftest.py | 1 + .../azure-cosmos-ai/tests/test_azure_openai_provider.py | 3 +++ .../tests/test_azure_openai_provider_async.py | 3 +++ 5 files changed, 17 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py index 8ad4c795d7c1..05fff3633c76 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py @@ -21,6 +21,7 @@ """Synchronous Azure OpenAI implementation of the EmbeddingProvider Protocol.""" +import time from typing import Any, Dict, Optional, Sequence, Union from azure.core.credentials import AzureKeyCredential, TokenCredential @@ -89,18 +90,21 @@ def generate_embeddings( :rtype: ~azure.cosmos.EmbeddingResult """ if not texts: - return EmbeddingResult(vectors=[], total_tokens=0) + return EmbeddingResult(vectors=[], total_tokens=0, latency=0.0) 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: diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py index d8833308c404..e67ca5728d1c 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py @@ -21,6 +21,7 @@ """Asynchronous Azure OpenAI implementation of the EmbeddingProvider Protocol.""" +import time from typing import Any, Dict, Optional, Sequence, Union from azure.core.credentials import AzureKeyCredential @@ -93,18 +94,21 @@ async def generate_embeddings( :rtype: ~azure.cosmos.EmbeddingResult """ if not texts: - return EmbeddingResult(vectors=[], total_tokens=0) + return EmbeddingResult(vectors=[], total_tokens=0, latency=0.0) client = self._get_or_create_client(endpoint) + start = time.perf_counter() response = await 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, ) async def close(self) -> None: diff --git a/sdk/cosmos/azure-cosmos-ai/tests/conftest.py b/sdk/cosmos/azure-cosmos-ai/tests/conftest.py index c2ce42819bcb..6f24c99ce5f7 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/conftest.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/conftest.py @@ -39,5 +39,6 @@ class EmbeddingResult: # pylint: disable=too-few-public-methods vectors: List[List[float]] total_tokens: Optional[int] = None + latency: Optional[float] = None _cosmos.EmbeddingResult = EmbeddingResult # type: ignore[attr-defined] diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py index 1b8cbea0a4b9..05b1128f4eb1 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py @@ -113,6 +113,8 @@ def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai): assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] assert result.total_tokens == 99 + assert isinstance(result.latency, float) + assert result.latency >= 0.0 def test_generate_embeddings_missing_usage_returns_none(mock_aoai): @@ -134,6 +136,7 @@ def test_empty_texts_short_circuits(mock_aoai): ) assert result.vectors == [] assert result.total_tokens == 0 + assert result.latency == 0.0 cls.assert_not_called() instance.embeddings.create.assert_not_called() diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py index d45cc59b5fcd..8be9ebc3e91a 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py @@ -118,6 +118,8 @@ async def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai) assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] assert result.total_tokens == 99 + assert isinstance(result.latency, float) + assert result.latency >= 0.0 @pytest.mark.asyncio @@ -141,6 +143,7 @@ async def test_empty_texts_short_circuits(mock_aoai): ) assert result.vectors == [] assert result.total_tokens == 0 + assert result.latency == 0.0 cls.assert_not_called() instance.embeddings.create.assert_not_called() From 784d526f68560d6baea7b05a7656c30c245b1543 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 19 May 2026 14:19:02 -0700 Subject: [PATCH 4/7] Adding test case --- .../tests/test_azure_openai_provider.py | 346 ++++++++++------ .../tests/test_azure_openai_provider_async.py | 386 ++++++++++++------ 2 files changed, 481 insertions(+), 251 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py index 05b1128f4eb1..7cd51b52057c 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py @@ -19,13 +19,29 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Unit tests for AzureOpenAIEmbeddingProvider (sync).""" +"""Tests for AzureOpenAIEmbeddingProvider (sync). +This module exposes two test classes: + +* ``TestAzureOpenAIProvider`` runs fully mocked unit tests and is always + collected. +* ``TestAzureOpenAIProviderLive`` runs opt-in live tests against a real + Azure OpenAI resource. Set ``COSMOS_AI_LIVE_TESTS=1`` and provide + connection settings via environment variables to enable it: + + * ``AZURE_OPENAI_ENDPOINT`` required (e.g. ``https://.openai.azure.com/``) + * ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT`` required + * ``AZURE_OPENAI_EMBEDDING_DIMENSIONS`` required (int) + * ``AZURE_OPENAI_API_KEY`` required for the API key tests +""" + +import os from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from azure.core.credentials import AccessToken, AzureKeyCredential +from azure.identity import DefaultAzureCredential from azure.cosmos.ai import AzureOpenAIEmbeddingProvider @@ -62,145 +78,239 @@ def mock_aoai(): yield cls, instance -# ----- constructor / credential dispatch ----- - - -def test_init_accepts_str(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential="my-key") - - -def test_init_accepts_azure_key_credential(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) - - -def test_init_accepts_token_credential(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential=_FakeTokenCredential()) - - -def test_init_rejects_unknown_credential(): - with pytest.raises(TypeError): - AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] - +class TestAzureOpenAIProvider: + """Unit tests with a mocked underlying ``AzureOpenAI`` client.""" -# ----- generate_embeddings ----- + # ----- constructor / credential dispatch ----- + def test_init_accepts_str(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential="my-key") -def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai): - cls, instance = mock_aoai - instance.embeddings.create.return_value = _fake_response( - [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 - ) - - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = provider.generate_embeddings( - ["a", "b", "c"], - endpoint=ENDPOINT, - deployment_name=DEPLOYMENT, - dimensions=DIMENSIONS, - ) + def test_init_accepts_azure_key_credential(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) - cls.assert_called_once() - _, ctor_kwargs = cls.call_args - assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY - assert ctor_kwargs["api_key"] == "key" - assert ctor_kwargs["api_version"] == "2024-10-21" + def test_init_accepts_token_credential(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=_FakeTokenCredential()) - instance.embeddings.create.assert_called_once_with( - input=["a", "b", "c"], - model=DEPLOYMENT, - dimensions=DIMENSIONS, - ) - - assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] - assert result.total_tokens == 99 - assert isinstance(result.latency, float) - assert result.latency >= 0.0 + def test_init_rejects_unknown_credential(self): + with pytest.raises(TypeError): + AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] + # ----- generate_embeddings ----- -def test_generate_embeddings_missing_usage_returns_none(mock_aoai): - _, instance = mock_aoai - instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) + def test_generate_embeddings_forwards_params_and_returns_result(self, mock_aoai): + cls, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 + ) - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = provider.generate_embeddings( - ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert result.total_tokens is None + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( + ["a", "b", "c"], + endpoint=ENDPOINT, + deployment_name=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + cls.assert_called_once() + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_key"] == "key" + assert ctor_kwargs["api_version"] == "2024-10-21" -def test_empty_texts_short_circuits(mock_aoai): - cls, instance = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = provider.generate_embeddings( - [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert result.vectors == [] - assert result.total_tokens == 0 - assert result.latency == 0.0 - cls.assert_not_called() - instance.embeddings.create.assert_not_called() + instance.embeddings.create.assert_called_once_with( + input=["a", "b", "c"], + model=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + assert result.total_tokens == 99 + assert isinstance(result.latency, float) + assert result.latency >= 0.0 -def test_exceptions_propagate(mock_aoai): - _, instance = mock_aoai - instance.embeddings.create.side_effect = RuntimeError("boom") + def test_generate_embeddings_missing_usage_returns_none(self, mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) - provider = AzureOpenAIEmbeddingProvider(credential="key") - with pytest.raises(RuntimeError, match="boom"): - provider.generate_embeddings( + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS ) + assert result.total_tokens is None + def test_empty_texts_short_circuits(self, mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = provider.generate_embeddings( + [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.vectors == [] + assert result.total_tokens == 0 + assert result.latency == 0.0 + cls.assert_not_called() + instance.embeddings.create.assert_not_called() + + def test_exceptions_propagate(self, mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.side_effect = RuntimeError("boom") + + provider = AzureOpenAIEmbeddingProvider(credential="key") + with pytest.raises(RuntimeError, match="boom"): + provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + + # ----- credential plumbing into AzureOpenAI ----- + + def test_azure_key_credential_passed_as_api_key(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_key"] == "aaa" -# ----- credential plumbing into AzureOpenAI ----- - - -def test_azure_key_credential_passed_as_api_key(mock_aoai): - cls, _ = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) - provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - _, ctor_kwargs = cls.call_args - assert ctor_kwargs["api_key"] == "aaa" + def test_token_credential_uses_bearer_token_provider(self, mock_aoai): + cls, _ = mock_aoai + cred = _FakeTokenCredential() + provider = AzureOpenAIEmbeddingProvider(credential=cred) + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert "api_key" not in ctor_kwargs + token_provider = ctor_kwargs["azure_ad_token_provider"] + token = token_provider() + assert token == "fake-token" + assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" + + # ----- close / context manager ----- + + def test_close_clears_cache(self, mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + provider.close() + instance.close.assert_called_once() + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert cls.call_count == 2 + def test_context_manager_closes(self, mock_aoai): + _, instance = mock_aoai + with AzureOpenAIEmbeddingProvider(credential="key") as provider: + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + instance.close.assert_called_once() -def test_token_credential_uses_bearer_token_provider(mock_aoai): - cls, _ = mock_aoai - cred = _FakeTokenCredential() - provider = AzureOpenAIEmbeddingProvider(credential=cred) - provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - _, ctor_kwargs = cls.call_args - assert "api_key" not in ctor_kwargs - token_provider = ctor_kwargs["azure_ad_token_provider"] - token = token_provider() - assert token == "fake-token" - assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" +# ----- live test config ----- -# ----- close / context manager ----- +_LIVE_ENABLED = os.getenv("COSMOS_AI_LIVE_TESTS") == "1" +_LIVE_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT", "") +_LIVE_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "") +_LIVE_DIMENSIONS = int(os.getenv("AZURE_OPENAI_EMBEDDING_DIMENSIONS") or "0") +_LIVE_API_KEY = os.getenv("AZURE_OPENAI_API_KEY", "") +_LIVE_SKIP_REASON = ( + "Set COSMOS_AI_LIVE_TESTS=1, AZURE_OPENAI_ENDPOINT, " + "AZURE_OPENAI_EMBEDDING_DEPLOYMENT and AZURE_OPENAI_EMBEDDING_DIMENSIONS " + "to run live tests." +) -def test_close_clears_cache(mock_aoai): - cls, instance = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential="key") - provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - provider.close() - instance.close.assert_called_once() - provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert cls.call_count == 2 +_LIVE_TEXTS = ["healthcare research papers", "azure cosmos vector search"] -def test_context_manager_closes(mock_aoai): - _, instance = mock_aoai - with AzureOpenAIEmbeddingProvider(credential="key") as provider: - provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - instance.close.assert_called_once() +def _assert_valid_live_result(result, expected_count): + assert len(result.vectors) == expected_count + for vec in result.vectors: + assert isinstance(vec, list) + assert len(vec) == _LIVE_DIMENSIONS + assert all(isinstance(v, float) for v in vec) + assert result.total_tokens is None or result.total_tokens > 0 + assert isinstance(result.latency, float) + assert result.latency > 0.0 + + +@pytest.mark.skipif( + not (_LIVE_ENABLED and _LIVE_ENDPOINT and _LIVE_DEPLOYMENT and _LIVE_DIMENSIONS), + reason=_LIVE_SKIP_REASON, +) +class TestAzureOpenAIProviderLive: + """Live tests against a real Azure OpenAI resource. Opt-in.""" + + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + def test_live_generate_embeddings_with_string_key(self): + with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY) as provider: + result = provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + def test_live_generate_embeddings_with_azure_key_credential(self): + with AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential(_LIVE_API_KEY)) as provider: + result = provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + + def test_live_generate_embeddings_with_default_azure_credential(self): + try: + credential = DefaultAzureCredential() + except Exception as exc: # pylint: disable=broad-except + pytest.skip(f"DefaultAzureCredential unavailable: {exc}") + with AzureOpenAIEmbeddingProvider(credential=credential) as provider: + try: + result = provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + except Exception as exc: # pylint: disable=broad-except + pytest.skip(f"Entra auth to Azure OpenAI failed (RBAC not granted?): {exc}") + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + + def test_live_empty_texts_short_circuits_no_network(self): + with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY or "unused") as provider: + result = provider.generate_embeddings( + [], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + assert result.vectors == [] + assert result.total_tokens == 0 + assert result.latency == 0.0 + + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + def test_live_underlying_client_is_cached_across_calls(self): + with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY) as provider: + provider.generate_embeddings( + _LIVE_TEXTS[:1], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + first_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + provider.generate_embeddings( + _LIVE_TEXTS[:1], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + second_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + assert first_client is second_client diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py index 8be9ebc3e91a..be06ad8ae011 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py @@ -19,13 +19,29 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Unit tests for AzureOpenAIEmbeddingProvider (async).""" +"""Tests for AzureOpenAIEmbeddingProvider (async). +This module exposes two test classes: + +* ``TestAzureOpenAIProviderAsync`` runs fully mocked unit tests and is always + collected. +* ``TestAzureOpenAIProviderLiveAsync`` runs opt-in live tests against a real + Azure OpenAI resource. Set ``COSMOS_AI_LIVE_TESTS=1`` and provide + connection settings via environment variables to enable it: + + * ``AZURE_OPENAI_ENDPOINT`` required (e.g. ``https://.openai.azure.com/``) + * ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT`` required + * ``AZURE_OPENAI_EMBEDDING_DIMENSIONS`` required (int) + * ``AZURE_OPENAI_API_KEY`` required for the API key tests +""" + +import os from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from azure.core.credentials import AccessToken, AzureKeyCredential +from azure.identity.aio import DefaultAzureCredential from azure.cosmos.ai.aio import AzureOpenAIEmbeddingProvider @@ -66,153 +82,257 @@ def mock_aoai(): yield cls, instance -# ----- constructor / credential dispatch ----- - - -def test_init_accepts_str(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential="my-key") - - -def test_init_accepts_azure_key_credential(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) - - -def test_init_accepts_async_token_credential(mock_aoai): # pylint: disable=unused-argument - AzureOpenAIEmbeddingProvider(credential=_FakeAsyncTokenCredential()) - - -def test_init_rejects_unknown_credential(): - with pytest.raises(TypeError): - AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] - - -# ----- generate_embeddings ----- +class TestAzureOpenAIProviderAsync: + """Unit tests with a mocked underlying ``AsyncAzureOpenAI`` client.""" + # ----- constructor / credential dispatch ----- -@pytest.mark.asyncio -async def test_generate_embeddings_forwards_params_and_returns_result(mock_aoai): - cls, instance = mock_aoai - instance.embeddings.create.return_value = _fake_response( - [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 - ) + def test_init_accepts_str(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential="my-key") - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = await provider.generate_embeddings( - ["a", "b", "c"], - endpoint=ENDPOINT, - deployment_name=DEPLOYMENT, - dimensions=DIMENSIONS, - ) + def test_init_accepts_azure_key_credential(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("my-key")) - cls.assert_called_once() - _, ctor_kwargs = cls.call_args - assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY - assert ctor_kwargs["api_key"] == "key" - assert ctor_kwargs["api_version"] == "2024-10-21" + def test_init_accepts_async_token_credential(self, mock_aoai): # pylint: disable=unused-argument + AzureOpenAIEmbeddingProvider(credential=_FakeAsyncTokenCredential()) - instance.embeddings.create.assert_awaited_once_with( - input=["a", "b", "c"], - model=DEPLOYMENT, - dimensions=DIMENSIONS, - ) + def test_init_rejects_unknown_credential(self): + with pytest.raises(TypeError): + AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] - assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] - assert result.total_tokens == 99 - assert isinstance(result.latency, float) - assert result.latency >= 0.0 + # ----- generate_embeddings ----- + @pytest.mark.asyncio + async def test_generate_embeddings_forwards_params_and_returns_result(self, mock_aoai): + cls, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], total_tokens=99 + ) -@pytest.mark.asyncio -async def test_generate_embeddings_missing_usage_returns_none(mock_aoai): - _, instance = mock_aoai - instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) - - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = await provider.generate_embeddings( - ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert result.total_tokens is None + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( + ["a", "b", "c"], + endpoint=ENDPOINT, + deployment_name=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + cls.assert_called_once() + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_key"] == "key" + assert ctor_kwargs["api_version"] == "2024-10-21" -@pytest.mark.asyncio -async def test_empty_texts_short_circuits(mock_aoai): - cls, instance = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential="key") - result = await provider.generate_embeddings( - [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert result.vectors == [] - assert result.total_tokens == 0 - assert result.latency == 0.0 - cls.assert_not_called() - instance.embeddings.create.assert_not_called() + instance.embeddings.create.assert_awaited_once_with( + input=["a", "b", "c"], + model=DEPLOYMENT, + dimensions=DIMENSIONS, + ) + assert result.vectors == [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + assert result.total_tokens == 99 + assert isinstance(result.latency, float) + assert result.latency >= 0.0 -@pytest.mark.asyncio -async def test_exceptions_propagate(mock_aoai): - _, instance = mock_aoai - instance.embeddings.create.side_effect = RuntimeError("boom") + @pytest.mark.asyncio + async def test_generate_embeddings_missing_usage_returns_none(self, mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.return_value = _fake_response([[1.0]], total_tokens=None) - provider = AzureOpenAIEmbeddingProvider(credential="key") - with pytest.raises(RuntimeError, match="boom"): - await provider.generate_embeddings( + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS ) - - -# ----- credential plumbing ----- - - -@pytest.mark.asyncio -async def test_azure_key_credential_passed_as_api_key(mock_aoai): - cls, _ = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) - await provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - _, ctor_kwargs = cls.call_args - assert ctor_kwargs["api_key"] == "aaa" - - -@pytest.mark.asyncio -async def test_async_token_credential_uses_bearer_token_provider(mock_aoai): - cls, _ = mock_aoai - cred = _FakeAsyncTokenCredential() - provider = AzureOpenAIEmbeddingProvider(credential=cred) - await provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - _, ctor_kwargs = cls.call_args - assert "api_key" not in ctor_kwargs - token_provider = ctor_kwargs["azure_ad_token_provider"] - token = await token_provider() - assert token == "fake-token" - assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" - - -# ----- close / async context manager ----- - - -@pytest.mark.asyncio -async def test_close_clears_cache(mock_aoai): - cls, instance = mock_aoai - provider = AzureOpenAIEmbeddingProvider(credential="key") - await provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - await provider.close() - instance.close.assert_awaited_once() - await provider.generate_embeddings( - ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS - ) - assert cls.call_count == 2 - - -@pytest.mark.asyncio -async def test_async_context_manager_closes(mock_aoai): - _, instance = mock_aoai - async with AzureOpenAIEmbeddingProvider(credential="key") as provider: + assert result.total_tokens is None + + @pytest.mark.asyncio + async def test_empty_texts_short_circuits(self, mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + result = await provider.generate_embeddings( + [], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + assert result.vectors == [] + assert result.total_tokens == 0 + assert result.latency == 0.0 + cls.assert_not_called() + instance.embeddings.create.assert_not_called() + + @pytest.mark.asyncio + async def test_exceptions_propagate(self, mock_aoai): + _, instance = mock_aoai + instance.embeddings.create.side_effect = RuntimeError("boom") + + provider = AzureOpenAIEmbeddingProvider(credential="key") + with pytest.raises(RuntimeError, match="boom"): + await provider.generate_embeddings( + ["a"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + + # ----- credential plumbing ----- + + @pytest.mark.asyncio + async def test_azure_key_credential_passed_as_api_key(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential=AzureKeyCredential("aaa")) + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_key"] == "aaa" + + @pytest.mark.asyncio + async def test_async_token_credential_uses_bearer_token_provider(self, mock_aoai): + cls, _ = mock_aoai + cred = _FakeAsyncTokenCredential() + provider = AzureOpenAIEmbeddingProvider(credential=cred) + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert "api_key" not in ctor_kwargs + token_provider = ctor_kwargs["azure_ad_token_provider"] + token = await token_provider() + assert token == "fake-token" + assert cred.calls and cred.calls[0][0] == "https://cognitiveservices.azure.com/.default" + + # ----- close / async context manager ----- + + @pytest.mark.asyncio + async def test_close_clears_cache(self, mock_aoai): + cls, instance = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key") + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + await provider.close() + instance.close.assert_awaited_once() await provider.generate_embeddings( ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS ) - instance.close.assert_awaited_once() + assert cls.call_count == 2 + + @pytest.mark.asyncio + async def test_async_context_manager_closes(self, mock_aoai): + _, instance = mock_aoai + async with AzureOpenAIEmbeddingProvider(credential="key") as provider: + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + instance.close.assert_awaited_once() + + +# ----- live test config ----- + +_LIVE_ENABLED = os.getenv("COSMOS_AI_LIVE_TESTS") == "1" +_LIVE_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT", "") +_LIVE_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "") +_LIVE_DIMENSIONS = int(os.getenv("AZURE_OPENAI_EMBEDDING_DIMENSIONS") or "0") +_LIVE_API_KEY = os.getenv("AZURE_OPENAI_API_KEY", "") + +_LIVE_SKIP_REASON = ( + "Set COSMOS_AI_LIVE_TESTS=1, AZURE_OPENAI_ENDPOINT, " + "AZURE_OPENAI_EMBEDDING_DEPLOYMENT and AZURE_OPENAI_EMBEDDING_DIMENSIONS " + "to run live tests." +) + +_LIVE_TEXTS = ["healthcare research papers", "azure cosmos vector search"] + + +def _assert_valid_live_result(result, expected_count): + assert len(result.vectors) == expected_count + for vec in result.vectors: + assert isinstance(vec, list) + assert len(vec) == _LIVE_DIMENSIONS + assert all(isinstance(v, float) for v in vec) + assert result.total_tokens is None or result.total_tokens > 0 + assert isinstance(result.latency, float) + assert result.latency > 0.0 + + +@pytest.mark.skipif( + not (_LIVE_ENABLED and _LIVE_ENDPOINT and _LIVE_DEPLOYMENT and _LIVE_DIMENSIONS), + reason=_LIVE_SKIP_REASON, +) +class TestAzureOpenAIProviderLiveAsync: + """Live tests against a real Azure OpenAI resource. Opt-in.""" + + @pytest.mark.asyncio + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + async def test_live_generate_embeddings_with_string_key(self): + async with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY) as provider: + result = await provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + + @pytest.mark.asyncio + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + async def test_live_generate_embeddings_with_azure_key_credential(self): + async with AzureOpenAIEmbeddingProvider( + credential=AzureKeyCredential(_LIVE_API_KEY) + ) as provider: + result = await provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + + @pytest.mark.asyncio + async def test_live_generate_embeddings_with_default_azure_credential(self): + try: + credential = DefaultAzureCredential() + except Exception as exc: # pylint: disable=broad-except + pytest.skip(f"Async DefaultAzureCredential unavailable: {exc}") + try: + async with AzureOpenAIEmbeddingProvider(credential=credential) as provider: + try: + result = await provider.generate_embeddings( + _LIVE_TEXTS, + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + except Exception as exc: # pylint: disable=broad-except + pytest.skip(f"Entra auth to Azure OpenAI failed (RBAC not granted?): {exc}") + _assert_valid_live_result(result, len(_LIVE_TEXTS)) + finally: + await credential.close() + + @pytest.mark.asyncio + async def test_live_empty_texts_short_circuits_no_network(self): + async with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY or "unused") as provider: + result = await provider.generate_embeddings( + [], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + assert result.vectors == [] + assert result.total_tokens == 0 + assert result.latency == 0.0 + + @pytest.mark.asyncio + @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") + async def test_live_underlying_client_is_cached_across_calls(self): + async with AzureOpenAIEmbeddingProvider(credential=_LIVE_API_KEY) as provider: + await provider.generate_embeddings( + _LIVE_TEXTS[:1], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + first_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + await provider.generate_embeddings( + _LIVE_TEXTS[:1], + endpoint=_LIVE_ENDPOINT, + deployment_name=_LIVE_DEPLOYMENT, + dimensions=_LIVE_DIMENSIONS, + ) + second_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + assert first_client is second_client From a09917727d5a94d78432baca9a113475947a760a Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 19 May 2026 16:22:52 -0700 Subject: [PATCH 5/7] Adding some code improvements --- sdk/cosmos/azure-cosmos-ai/MANIFEST.in | 2 - .../azure/cosmos/ai/_azure_openai_provider.py | 98 +++++++++----- .../cosmos/ai/aio/_azure_openai_provider.py | 127 +++++++++++++----- sdk/cosmos/azure-cosmos-ai/pytest.ini | 9 ++ .../samples/sample_embedding_provider.py | 16 +-- sdk/cosmos/azure-cosmos-ai/setup.py | 2 +- .../tests/test_azure_openai_provider.py | 59 +++++++- .../tests/test_azure_openai_provider_async.py | 62 ++++++++- 8 files changed, 290 insertions(+), 85 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-ai/pytest.ini diff --git a/sdk/cosmos/azure-cosmos-ai/MANIFEST.in b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in index f2097fcd6842..99ea955ebe45 100644 --- a/sdk/cosmos/azure-cosmos-ai/MANIFEST.in +++ b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in @@ -2,6 +2,4 @@ recursive-include samples *.py recursive-include tests *.py include *.md include LICENSE -include azure/__init__.py -include azure/cosmos/__init__.py include azure/cosmos/ai/py.typed diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py index 05fff3633c76..270f7ded5311 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/_azure_openai_provider.py @@ -21,8 +21,9 @@ """Synchronous Azure OpenAI implementation of the EmbeddingProvider Protocol.""" +import inspect import time -from typing import Any, Dict, Optional, Sequence, Union +from typing import Any, Dict, Mapping, Optional, Sequence, Union from azure.core.credentials import AzureKeyCredential, TokenCredential from azure.cosmos import EmbeddingResult @@ -51,20 +52,46 @@ class AzureOpenAIEmbeddingProvider: :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], - **kwargs: Any, # pylint: disable=unused-argument + *, + api_version: str = _AZURE_OPENAI_API_VERSION, + openai_client_kwargs: Optional[Mapping[str, Any]] = None, ) -> None: - if not isinstance(credential, (str, AzureKeyCredential)) and not _is_token_credential(credential): + 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 TokenCredential; " - f"got {type(credential).__name__}" + "credential must be a str, AzureKeyCredential, or synchronous " + f"TokenCredential; got {type(credential).__name__}" ) + self._credential = credential - self._api_version = _AZURE_OPENAI_API_VERSION + 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( @@ -74,7 +101,7 @@ def generate_embeddings( endpoint: str, deployment_name: str, dimensions: int, - **kwargs: Any, # pylint: disable=unused-argument + **kwargs: Any, ) -> EmbeddingResult: """Generate embeddings for ``texts`` using Azure OpenAI. @@ -86,11 +113,17 @@ def generate_embeddings( (from ``vectorEmbeddingPolicy.embeddingSource.deploymentName``). :keyword int dimensions: Embedding dimensions (from ``vectorEmbeddingPolicy.dimensions``). - :returns: Vectors in the same order as ``texts``, plus token usage. + :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=0.0) + return EmbeddingResult(vectors=[], total_tokens=0, latency=None) client = self._get_or_create_client(endpoint) start = time.perf_counter() @@ -109,12 +142,13 @@ def generate_embeddings( def close(self) -> None: """Close every cached underlying Azure OpenAI client and clear the cache.""" - for client in self._clients.values(): + clients = list(self._clients.values()) + self._clients.clear() + for client in clients: try: client.close() except Exception: # pylint: disable=broad-except pass - self._clients.clear() def __enter__(self) -> "AzureOpenAIEmbeddingProvider": return self @@ -131,30 +165,34 @@ def _get_or_create_client(self, endpoint: str) -> AzureOpenAI: 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( - azure_endpoint=endpoint, - api_version=self._api_version, - api_key=self._credential, - ) + return AzureOpenAI(api_key=self._credential, **common) if isinstance(self._credential, AzureKeyCredential): - return AzureOpenAI( - azure_endpoint=endpoint, - api_version=self._api_version, - api_key=self._credential.key, - ) + return AzureOpenAI(api_key=self._credential.key, **common) token_provider = get_bearer_token_provider(self._credential, _COGNITIVE_SERVICES_SCOPE) - return AzureOpenAI( - azure_endpoint=endpoint, - api_version=self._api_version, - azure_ad_token_provider=token_provider, - ) + return AzureOpenAI(azure_ad_token_provider=token_provider, **common) def _is_token_credential(obj: Any) -> bool: - """Duck-type check for a sync TokenCredential. + """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. - Avoids ``isinstance(obj, TokenCredential)`` because ``TokenCredential`` is a - ``Protocol`` in some azure-core versions and not always runtime_checkable. + Used only to produce an actionable error message when an async credential is + accidentally passed to the sync provider. """ - return callable(getattr(obj, "get_token", None)) + get_token = getattr(obj, "get_token", None) + return callable(get_token) and inspect.iscoroutinefunction(get_token) diff --git a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py index e67ca5728d1c..837f1824f8a4 100644 --- a/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/azure/cosmos/ai/aio/_azure_openai_provider.py @@ -21,8 +21,10 @@ """Asynchronous Azure OpenAI implementation of the EmbeddingProvider Protocol.""" +import asyncio +import inspect import time -from typing import Any, Dict, Optional, Sequence, Union +from typing import Any, Dict, Mapping, Optional, Sequence, Union from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential @@ -53,21 +55,50 @@ class AzureOpenAIEmbeddingProvider: :type credential: str or ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials_async.AsyncTokenCredential + :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.AsyncAzureOpenAI` (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, AsyncTokenCredential], - **kwargs: Any, # pylint: disable=unused-argument + *, + api_version: str = _AZURE_OPENAI_API_VERSION, + openai_client_kwargs: Optional[Mapping[str, Any]] = None, ) -> None: - if not isinstance(credential, (str, AzureKeyCredential)) and not _is_async_token_credential(credential): + if isinstance(credential, (str, AzureKeyCredential)): + pass + elif _is_async_token_credential(credential): + pass + elif _is_sync_token_credential(credential): + raise TypeError( + "Asynchronous AzureOpenAIEmbeddingProvider received a sync " + f"credential ({type(credential).__name__}). Either use " + "azure.cosmos.ai.AzureOpenAIEmbeddingProvider instead, or " + "pass an asynchronous TokenCredential such as " + "azure.identity.aio.DefaultAzureCredential." + ) + else: raise TypeError( - "credential must be a str, AzureKeyCredential, or AsyncTokenCredential; " - f"got {type(credential).__name__}" + "credential must be a str, AzureKeyCredential, or asynchronous " + f"TokenCredential; got {type(credential).__name__}" ) + self._credential = credential - self._api_version = _AZURE_OPENAI_API_VERSION + self._api_version = api_version + self._openai_client_kwargs: Dict[str, Any] = dict(openai_client_kwargs or {}) self._clients: Dict[str, AsyncAzureOpenAI] = {} + # Lazily created on first use so we don't require a running event loop at __init__. + self._clients_lock: Optional[asyncio.Lock] = None async def generate_embeddings( self, @@ -76,7 +107,7 @@ async def generate_embeddings( endpoint: str, deployment_name: str, dimensions: int, - **kwargs: Any, # pylint: disable=unused-argument + **kwargs: Any, ) -> EmbeddingResult: """Generate embeddings for ``texts`` using Azure OpenAI. @@ -90,13 +121,19 @@ async def generate_embeddings( (from ``vectorEmbeddingPolicy.embeddingSource.deploymentName``). :keyword int dimensions: Embedding dimensions (from ``vectorEmbeddingPolicy.dimensions``). - :returns: Vectors in the same order as ``texts``, plus token usage. + :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=0.0) + return EmbeddingResult(vectors=[], total_tokens=0, latency=None) - client = self._get_or_create_client(endpoint) + client = await self._get_or_create_client(endpoint) start = time.perf_counter() response = await client.embeddings.create( input=list(texts), @@ -112,13 +149,19 @@ async def generate_embeddings( ) async def close(self) -> None: - """Close every cached underlying Azure OpenAI client and clear the cache.""" - for client in self._clients.values(): + """Close every cached underlying Azure OpenAI client and clear the cache. + + Snapshots the cached clients and clears the dict *before* awaiting each + ``close()`` so that a concurrent :meth:`generate_embeddings` cannot + observe a half-closed client. + """ + clients = list(self._clients.values()) + self._clients.clear() + for client in clients: try: await client.close() except Exception: # pylint: disable=broad-except pass - self._clients.clear() async def __aenter__(self) -> "AzureOpenAIEmbeddingProvider": return self @@ -126,35 +169,51 @@ async def __aenter__(self) -> "AzureOpenAIEmbeddingProvider": async def __aexit__(self, *args: Any) -> None: await self.close() - def _get_or_create_client(self, endpoint: str) -> AsyncAzureOpenAI: + def _ensure_lock(self) -> asyncio.Lock: + if self._clients_lock is None: + self._clients_lock = asyncio.Lock() + return self._clients_lock + + async def _get_or_create_client(self, endpoint: str) -> AsyncAzureOpenAI: key = endpoint.rstrip("/") client = self._clients.get(key) - if client is None: - client = self._build_client(key) - self._clients[key] = client + if client is not None: + return client + async with self._ensure_lock(): + 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) -> AsyncAzureOpenAI: + # 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 AsyncAzureOpenAI( - azure_endpoint=endpoint, - api_version=self._api_version, - api_key=self._credential, - ) + return AsyncAzureOpenAI(api_key=self._credential, **common) if isinstance(self._credential, AzureKeyCredential): - return AsyncAzureOpenAI( - azure_endpoint=endpoint, - api_version=self._api_version, - api_key=self._credential.key, - ) + return AsyncAzureOpenAI(api_key=self._credential.key, **common) token_provider = get_bearer_token_provider(self._credential, _COGNITIVE_SERVICES_SCOPE) - return AsyncAzureOpenAI( - azure_endpoint=endpoint, - api_version=self._api_version, - azure_ad_token_provider=token_provider, - ) + return AsyncAzureOpenAI(azure_ad_token_provider=token_provider, **common) def _is_async_token_credential(obj: Any) -> bool: - """Duck-type check for an async TokenCredential.""" - return callable(getattr(obj, "get_token", None)) + """Duck-type check for an *asynchronous* TokenCredential. + + Accepts any object that exposes a coroutine ``get_token`` method. Sync + credentials are rejected so the mismatch is caught at ``__init__`` instead + of failing deep inside ``openai`` with a confusing error. + """ + get_token = getattr(obj, "get_token", None) + return callable(get_token) and inspect.iscoroutinefunction(get_token) + + +def _is_sync_token_credential(obj: Any) -> bool: + """Duck-type check for a *synchronous* TokenCredential. + + Used only to produce an actionable error message when a sync credential is + accidentally passed to the async provider. + """ + get_token = getattr(obj, "get_token", None) + return callable(get_token) and not inspect.iscoroutinefunction(get_token) diff --git a/sdk/cosmos/azure-cosmos-ai/pytest.ini b/sdk/cosmos/azure-cosmos-ai/pytest.ini new file mode 100644 index 000000000000..c3dee16689a6 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-ai/pytest.ini @@ -0,0 +1,9 @@ +[pytest] +# All tests under sdk/cosmos/ run through the shared cosmos-sdk-client pipeline +# template, which filters via `-m cosmosEmulator`. azure-cosmos-ai's unit tests +# are fully mocked and do not require the Cosmos emulator, but they still need +# to be collected by that filter — so we register the marker here (to silence +# pytest's unknown-marker warning) and apply it module-wide in each test file. +markers = + cosmosEmulator: marks tests as part of the cosmos service test umbrella (azure-cosmos-ai tests do not actually require the emulator). +asyncio_mode = auto diff --git a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py index b53bebf2f3f8..32646a87b318 100644 --- a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py @@ -63,14 +63,14 @@ def variant_a_api_key() -> None: def variant_b_shared_entra() -> None: cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] - cred = DefaultAzureCredential() - with AzureOpenAIEmbeddingProvider(credential=cred) as provider: - client = CosmosClient( - url=cosmos_endpoint, - credential=cred, - embedding_provider=provider, - ) - _run_query(client) + with DefaultAzureCredential() as cred: + with AzureOpenAIEmbeddingProvider(credential=cred) as provider: + client = CosmosClient( + url=cosmos_endpoint, + credential=cred, + embedding_provider=provider, + ) + _run_query(client) def _run_query(client: CosmosClient) -> None: diff --git a/sdk/cosmos/azure-cosmos-ai/setup.py b/sdk/cosmos/azure-cosmos-ai/setup.py index f9bd9131c45b..21eff3b7ab7e 100644 --- a/sdk/cosmos/azure-cosmos-ai/setup.py +++ b/sdk/cosmos/azure-cosmos-ai/setup.py @@ -51,8 +51,8 @@ exclude_packages = [ "tests", "samples", - # Exclude packages that will be covered by PEP420 or nspkg "azure", + "azure.cosmos", ] setup( diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py index 7cd51b52057c..254d62f7e159 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider.py @@ -52,6 +52,12 @@ DIMENSIONS = 1536 +# Apply the shared cosmos CI test marker at module scope. The cosmos-sdk-client +# pipeline template filters via ``-m cosmosEmulator``; without this declaration +# every test in this module would be silently deselected in CI. +pytestmark = pytest.mark.cosmosEmulator + + class _FakeTokenCredential: def __init__(self): self.calls = [] @@ -61,6 +67,17 @@ def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument return AccessToken("fake-token", 9999999999) +class _FakeAsyncTokenCredential: + """Stand-in for an azure.identity.aio credential without the import cost.""" + + def __init__(self): + self.calls = [] + + async def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + self.calls.append(scopes) + return AccessToken("fake-token", 9999999999) + + def _fake_response(vectors, total_tokens=42): return SimpleNamespace( data=[SimpleNamespace(embedding=v) for v in vectors], @@ -96,6 +113,35 @@ def test_init_rejects_unknown_credential(self): with pytest.raises(TypeError): AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] + def test_init_rejects_async_credential_with_actionable_message(self): + with pytest.raises(TypeError, match=r"(?i)async"): + AzureOpenAIEmbeddingProvider(credential=_FakeAsyncTokenCredential()) # type: ignore[arg-type] + + def test_init_accepts_api_version_override(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key", api_version="2024-12-01-preview") + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_version"] == "2024-12-01-preview" + + def test_init_accepts_openai_client_kwargs(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider( + credential="key", + openai_client_kwargs={"timeout": 30.0, "default_headers": {"x-test": "1"}}, + ) + provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["timeout"] == 30.0 + assert ctor_kwargs["default_headers"] == {"x-test": "1"} + # Explicit provider-controlled kwargs still win. + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_version"] == "2024-10-21" + # ----- generate_embeddings ----- def test_generate_embeddings_forwards_params_and_returns_result(self, mock_aoai): @@ -147,7 +193,7 @@ def test_empty_texts_short_circuits(self, mock_aoai): ) assert result.vectors == [] assert result.total_tokens == 0 - assert result.latency == 0.0 + assert result.latency is None cls.assert_not_called() instance.embeddings.create.assert_not_called() @@ -294,7 +340,7 @@ def test_live_empty_texts_short_circuits_no_network(self): ) assert result.vectors == [] assert result.total_tokens == 0 - assert result.latency == 0.0 + assert result.latency is None @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") def test_live_underlying_client_is_cached_across_calls(self): @@ -305,12 +351,15 @@ def test_live_underlying_client_is_cached_across_calls(self): deployment_name=_LIVE_DEPLOYMENT, dimensions=_LIVE_DIMENSIONS, ) - first_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + first_clients = dict(provider._clients) # pylint: disable=protected-access + assert len(first_clients) == 1 provider.generate_embeddings( _LIVE_TEXTS[:1], endpoint=_LIVE_ENDPOINT, deployment_name=_LIVE_DEPLOYMENT, dimensions=_LIVE_DIMENSIONS, ) - second_client = next(iter(provider._clients.values())) # pylint: disable=protected-access - assert first_client is second_client + second_clients = dict(provider._clients) # pylint: disable=protected-access + assert first_clients.keys() == second_clients.keys() + for key in first_clients: + assert first_clients[key] is second_clients[key] diff --git a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py index be06ad8ae011..743550feef03 100644 --- a/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py +++ b/sdk/cosmos/azure-cosmos-ai/tests/test_azure_openai_provider_async.py @@ -52,6 +52,12 @@ DIMENSIONS = 1536 +# Apply the shared cosmos CI test marker at module scope. The cosmos-sdk-client +# pipeline template filters via ``-m cosmosEmulator``; without this declaration +# every test in this module would be silently deselected in CI. +pytestmark = pytest.mark.cosmosEmulator + + class _FakeAsyncTokenCredential: def __init__(self): self.calls = [] @@ -64,6 +70,18 @@ async def close(self): pass +class _FakeSyncTokenCredential: + """Stand-in for an azure.identity (sync) credential. Used to verify the + async provider rejects sync credentials at __init__.""" + + def __init__(self): + self.calls = [] + + def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + self.calls.append(scopes) + return AccessToken("fake-token", 9999999999) + + def _fake_response(vectors, total_tokens=42): return SimpleNamespace( data=[SimpleNamespace(embedding=v) for v in vectors], @@ -100,6 +118,37 @@ def test_init_rejects_unknown_credential(self): with pytest.raises(TypeError): AzureOpenAIEmbeddingProvider(credential=12345) # type: ignore[arg-type] + def test_init_rejects_sync_credential_with_actionable_message(self): + with pytest.raises(TypeError, match=r"(?i)sync"): + AzureOpenAIEmbeddingProvider(credential=_FakeSyncTokenCredential()) # type: ignore[arg-type] + + @pytest.mark.asyncio + async def test_init_accepts_api_version_override(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider(credential="key", api_version="2024-12-01-preview") + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["api_version"] == "2024-12-01-preview" + + @pytest.mark.asyncio + async def test_init_accepts_openai_client_kwargs(self, mock_aoai): + cls, _ = mock_aoai + provider = AzureOpenAIEmbeddingProvider( + credential="key", + openai_client_kwargs={"timeout": 30.0, "default_headers": {"x-test": "1"}}, + ) + await provider.generate_embeddings( + ["x"], endpoint=ENDPOINT, deployment_name=DEPLOYMENT, dimensions=DIMENSIONS + ) + _, ctor_kwargs = cls.call_args + assert ctor_kwargs["timeout"] == 30.0 + assert ctor_kwargs["default_headers"] == {"x-test": "1"} + # Explicit provider-controlled kwargs still win. + assert ctor_kwargs["azure_endpoint"] == ENDPOINT_KEY + assert ctor_kwargs["api_version"] == "2024-10-21" + # ----- generate_embeddings ----- @pytest.mark.asyncio @@ -154,7 +203,7 @@ async def test_empty_texts_short_circuits(self, mock_aoai): ) assert result.vectors == [] assert result.total_tokens == 0 - assert result.latency == 0.0 + assert result.latency is None cls.assert_not_called() instance.embeddings.create.assert_not_called() @@ -315,7 +364,7 @@ async def test_live_empty_texts_short_circuits_no_network(self): ) assert result.vectors == [] assert result.total_tokens == 0 - assert result.latency == 0.0 + assert result.latency is None @pytest.mark.asyncio @pytest.mark.skipif(not _LIVE_API_KEY, reason="AZURE_OPENAI_API_KEY not set.") @@ -327,12 +376,15 @@ async def test_live_underlying_client_is_cached_across_calls(self): deployment_name=_LIVE_DEPLOYMENT, dimensions=_LIVE_DIMENSIONS, ) - first_client = next(iter(provider._clients.values())) # pylint: disable=protected-access + first_clients = dict(provider._clients) # pylint: disable=protected-access + assert len(first_clients) == 1 await provider.generate_embeddings( _LIVE_TEXTS[:1], endpoint=_LIVE_ENDPOINT, deployment_name=_LIVE_DEPLOYMENT, dimensions=_LIVE_DIMENSIONS, ) - second_client = next(iter(provider._clients.values())) # pylint: disable=protected-access - assert first_client is second_client + second_clients = dict(provider._clients) # pylint: disable=protected-access + assert first_clients.keys() == second_clients.keys() + for key in first_clients: + assert first_clients[key] is second_clients[key] From f1d3c71e08dc5cc9e4fb48ad2d2b9edfe55549e8 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 20 May 2026 11:46:30 -0700 Subject: [PATCH 6/7] Resolving comments --- .../samples/sample_embedding_provider.py | 90 ------------------- .../sample_embedding_provider_async.py | 87 ------------------ 2 files changed, 177 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py delete mode 100644 sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py diff --git a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py deleted file mode 100644 index 32646a87b318..000000000000 --- a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. - -"""Sample: use AzureOpenAIEmbeddingProvider with a synchronous CosmosClient. - -Demonstrates two credential modes: - -* Variant A — Azure OpenAI API key. -* Variant B — Entra (RBAC). The SAME ``DefaultAzureCredential`` is shared - between ``CosmosClient`` (Cosmos RBAC) and the embedding provider - (Azure OpenAI), so the user only needs one identity. - -Required environment variables: - -* ``COSMOS_ENDPOINT`` – e.g. ``https://my-cosmos.documents.azure.com:443/`` -* ``COSMOS_KEY`` – only for Variant A -* ``AOAI_API_KEY`` – only for Variant A - -Both samples assume a database ``samples-db`` with a container ``items`` -whose ``vectorEmbeddingPolicy.embeddingSource`` already points at the -Azure OpenAI endpoint and deployment you intend to use. -""" - -import os - -from azure.cosmos import CosmosClient -from azure.cosmos.ai import AzureOpenAIEmbeddingProvider -from azure.identity import DefaultAzureCredential - - -def variant_a_api_key() -> None: - cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] - cosmos_key = os.environ["COSMOS_KEY"] - aoai_api_key = os.environ["AOAI_API_KEY"] - - with AzureOpenAIEmbeddingProvider(credential=aoai_api_key) as provider: - client = CosmosClient( - url=cosmos_endpoint, - credential=cosmos_key, - embedding_provider=provider, - ) - _run_query(client) - - -def variant_b_shared_entra() -> None: - cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] - - with DefaultAzureCredential() as cred: - with AzureOpenAIEmbeddingProvider(credential=cred) as provider: - client = CosmosClient( - url=cosmos_endpoint, - credential=cred, - embedding_provider=provider, - ) - _run_query(client) - - -def _run_query(client: CosmosClient) -> None: - db = client.get_database_client("samples-db") - container = db.get_container_client("items") - - query = ( - "SELECT TOP 5 c.id, " - "VectorDistance(c.embedding, GenerateEmbeddings('healthcare research papers')) AS score " - "FROM c ORDER BY score" - ) - for item in container.query_items(query=query, enable_cross_partition_query=True): - print(item) - - -if __name__ == "__main__": - variant_a_api_key() diff --git a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py b/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py deleted file mode 100644 index 5ed30e28685d..000000000000 --- a/sdk/cosmos/azure-cosmos-ai/samples/sample_embedding_provider_async.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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. - -"""Sample: use AzureOpenAIEmbeddingProvider with an async CosmosClient. - -Demonstrates two credential modes: - -* Variant A — Azure OpenAI API key. -* Variant B — Entra (RBAC). The SAME ``DefaultAzureCredential`` is shared - between ``CosmosClient`` (Cosmos RBAC) and the embedding provider - (Azure OpenAI), so the user only needs one identity. - -Required environment variables: - -* ``COSMOS_ENDPOINT`` – e.g. ``https://my-cosmos.documents.azure.com:443/`` -* ``COSMOS_KEY`` – only for Variant A -* ``AOAI_API_KEY`` – only for Variant A -""" - -import asyncio -import os - -from azure.cosmos.aio import CosmosClient -from azure.cosmos.ai.aio import AzureOpenAIEmbeddingProvider -from azure.identity.aio import DefaultAzureCredential - - -async def variant_a_api_key() -> None: - cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] - cosmos_key = os.environ["COSMOS_KEY"] - aoai_api_key = os.environ["AOAI_API_KEY"] - - async with AzureOpenAIEmbeddingProvider(credential=aoai_api_key) as provider: - async with CosmosClient( - url=cosmos_endpoint, - credential=cosmos_key, - embedding_provider=provider, - ) as client: - await _run_query(client) - - -async def variant_b_shared_entra() -> None: - cosmos_endpoint = os.environ["COSMOS_ENDPOINT"] - - async with DefaultAzureCredential() as cred: - async with AzureOpenAIEmbeddingProvider(credential=cred) as provider: - async with CosmosClient( - url=cosmos_endpoint, - credential=cred, - embedding_provider=provider, - ) as client: - await _run_query(client) - - -async def _run_query(client: CosmosClient) -> None: - db = client.get_database_client("samples-db") - container = db.get_container_client("items") - - query = ( - "SELECT TOP 5 c.id, " - "VectorDistance(c.embedding, GenerateEmbeddings('healthcare research papers')) AS score " - "FROM c ORDER BY score" - ) - async for item in container.query_items(query=query): - print(item) - - -if __name__ == "__main__": - asyncio.run(variant_a_api_key()) From e664904e066667f5248c12366bf1b9bec4e7f06f Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 20 May 2026 11:47:58 -0700 Subject: [PATCH 7/7] Resolving comments --- sdk/cosmos/azure-cosmos-ai/MANIFEST.in | 1 - sdk/cosmos/azure-cosmos-ai/setup.py | 1 - 2 files changed, 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-ai/MANIFEST.in b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in index 99ea955ebe45..dfb0956250b6 100644 --- a/sdk/cosmos/azure-cosmos-ai/MANIFEST.in +++ b/sdk/cosmos/azure-cosmos-ai/MANIFEST.in @@ -1,4 +1,3 @@ -recursive-include samples *.py recursive-include tests *.py include *.md include LICENSE diff --git a/sdk/cosmos/azure-cosmos-ai/setup.py b/sdk/cosmos/azure-cosmos-ai/setup.py index 21eff3b7ab7e..704495d9ff58 100644 --- a/sdk/cosmos/azure-cosmos-ai/setup.py +++ b/sdk/cosmos/azure-cosmos-ai/setup.py @@ -50,7 +50,6 @@ exclude_packages = [ "tests", - "samples", "azure", "azure.cosmos", ]