Skip to content

[Cosmos] [Embedding V0] Public surface: EmbeddingGenerator protocols + embedding_generator on query_items #46730

@ananth7592

Description

@ananth7592

Public surface: EmbeddingGenerator protocols + embedding_generator keyword on query_items

Parent: 46729

Goal

Add the customer-facing API for embedding generation, with no behavioral change yet.

Scope

  1. New protocols (export from package root, e.g. azure.cosmos):
    from typing import Protocol, Sequence
    
    @runtime_checkable
    class EmbeddingGenerator(Protocol):
        def generate_embeddings(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: ...
    
    @runtime_checkable
    class AsyncEmbeddingGenerator(Protocol):
        async def generate_embeddings_async(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: ...
  2. New keyword-only argument on query_items:
    • azure/cosmos/container.py :: Container.query_itemsembedding_generator: Optional["EmbeddingGenerator"] = None
    • azure/cosmos/aio/_container.py :: ContainerProxy.query_itemsembedding_generator: Optional["AsyncEmbeddingGenerator"] = None
  3. Stuff into feed_options:
    if embedding_generator is not None:
        feed_options["embeddingGenerator"] = embedding_generator
    Mirror the existing pattern at container.py:1055-1058 (full_text_score_scope).
  4. Type hints + docstrings (Sphinx). Add to _overloads typing files if present.

Non-goals

  • Do NOT yet read the option in any execution path (will be added in 46733).
  • Do NOT yet add the _QueryFeature token (covered by 46731).
  • Do NOT add a built-in implementation. The EGS package owns that.

Files touched

  • sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py (export new protocols)
  • sdk/cosmos/azure-cosmos/azure/cosmos/container.py
  • sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py
  • New: sdk/cosmos/azure-cosmos/azure/cosmos/_embedding_generator.py (or in documents.py)
  • sdk/cosmos/azure-cosmos/CHANGELOG.md

Acceptance

  • from azure.cosmos import EmbeddingGenerator, AsyncEmbeddingGenerator works.
  • query_items(..., embedding_generator=mock) is accepted and stored without error.
  • Mypy / pyright pass on the new protocols.
  • Existing query_items callers unchanged.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Cosmosfeature-requestThis issue requires a new behavior in the product in order be resolved.

    Type

    No type

    Projects

    Status

    No status

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions