Azure / Azure/azure-sdk-for-python

[Cosmos] [Embedding V0] AzureOpenAIEmbeddingGenerator provider implementation

Open
#46,766 0 comments 0 reactions 0 assignees View on GitHub
Cosmos
Dominant language
Python
Stars
5.6k
Forks
3.4k
Avg merge
1d 21h
Merged PRs (30d)
193

Description

# [Cosmos] [Embedding V0] `AzureOpenAIEmbeddingGenerator` — provider implementation

Parent: [#46729](https://github.com/Azure/azure-sdk-for-python/issues/46729)

## Background

This is the concrete implementation of `EmbeddingGenerator` / `AsyncEmbeddingGenerator` that lives inside the `azure-cosmos-embedding` package. It wraps the `openai.AzureOpenAI` (sync) and `openai.AsyncAzureOpenAI` (async) clients to generate text embeddings, and can be constructed from a `VectorEmbeddingPolicy` `EmbeddingSource` or from explicit parameters.

Reference: the .NET sibling is `AzureOpenAIEmbeddingGenerator` in `Microsoft.Azure.Cosmos.Embedding` (see [Azure/azure-cosmos-dotnet-v3#5842](https://github.com/Azure/azure-cosmos-dotnet-v3/issues/5842)).

## Public API

```python
# azure/cosmos/embedding/_azure_openai_embedding_generator.py

from azure.cosmos import EmbeddingGenerator, AsyncEmbeddingGenerator
from azure.cosmos._models import EmbeddingSource

class AzureOpenAIEmbeddingGenerator:
"""
Implements both EmbeddingGenerator (sync) and AsyncEmbeddingGenerator (async).

Can be constructed directly or from a container's VectorEmbeddingPolicy EmbeddingSource.
"""

@classmethod
def from_embedding_source(
cls,
source: EmbeddingSource,
api_key: str | None = None,
credential=None, # azure.core.credentials.TokenCredential
) -> "AzureOpenAIEmbeddingGenerator": ...

def __init__(
self,
endpoint: str,
deployment_name: str,
dimensions: int,
api_key: str | None = None,
credential=None, # azure.core.credentials.TokenCredential for Entra auth
max_batch_size: int = 100,
) -> None: ...

def generate_embeddings(
self, texts: Sequence[str]
) -> Sequence[Sequence[float]]: ... # sync — EmbeddingGenerator Protocol

async def generate_embeddings_async(
self, texts: Sequence[str]
) -> Sequence[Sequence[float]]: ... # async — AsyncEmbeddingGenerator Protocol
```

## Behavioral requirements

1. **Auth**: ApiKey via `openai.AzureOpenAI(api_key=api_key)` or Entra via `azure.identity` `get_token` injected as `AzureOpenAI(azure_ad_token_provider=...)`.
2. **Batching**: if `len(texts) > max_batch_size`, split into chunks and concatenate results in order.
3. **Dimensions**: pass `dimensions=self.dimensions` to the embeddings API call so output size matches the container policy.
4. **Return type**: `list[list[float]]` — each inner list is one embedding vector.
5. **Error handling**: surface `openai.APIError` / `openai.APIConnectionError` as `azure.cosmos.exceptions.CosmosHttpResponseError` (or a new `CosmosEmbeddingError`) with the original exception as `__cause__`.
6. **`from_embedding_source`**: reads `endpoint`, `deploymentName`, `dimensions`, and `authType` from the `EmbeddingSource` TypedDict; `api_key` or `credential` must be supplied by the caller (not stored in the policy for security reasons).
7. **Thread safety**: `AzureOpenAI` client is thread-safe; concurrent sync calls are supported.
8. **Event loop**: async variant uses `AsyncAzureOpenAI`; do not mix sync client into an async context.

## Acceptance criteria

- Unit tests (in package's `tests/` folder) cover:
- Happy path sync: single batch, returns correct shape `list[list[float]]`.
- Happy path async: same.
- Multi-batch: input > `max_batch_size`, two API calls made, results concatenated in order.
- `APIError` → wrapped as `CosmosHttpResponseError` (or `CosmosEmbeddingError`).
- `from_embedding_source` constructs correctly from a typed `EmbeddingSource` dict.
- `AzureOpenAIEmbeddingGenerator` is exported from `azure.cosmos.embedding`.
- `isinstance(gen, EmbeddingGenerator)` returns `True` (Protocol `runtime_checkable`).
- Works end-to-end when passed as `embedding_generator` to `container.query_items` in an emulator test (covered by #46736).

## Files likely touched

- `sdk/cosmos/azure-cosmos-embedding/azure/cosmos/embedding/_azure_openai_embedding_generator.py` (new)
- `sdk/cosmos/azure-cosmos-embedding/azure/cosmos/embedding/__init__.py`
- `sdk/cosmos/azure-cosmos-embedding/tests/test_azure_openai_embedding_generator.py` (new)

## Dependencies

- Package scaffold (previous sub-issue)
- `EmbeddingSource` TypedDict (#NEW-A — VectorEmbeddingPolicy issue)
- `EmbeddingGenerator` / `AsyncEmbeddingGenerator` protocols (#46730)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.