[Cosmos] [Embedding V0] AzureOpenAIEmbeddingGenerator provider implementation

[ananth7592]

[Cosmos] [Embedding V0] AzureOpenAIEmbeddingGenerator — provider implementation

Parent: #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).

Public API

# 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 [Cosmos] [Embedding V0] Emulator tests + sample for embedding generation #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 ([Cosmos] [Embedding V0] Public surface: EmbeddingGenerator protocols + embedding_generator on query_items #46730)