Azure / Azure/azure-sdk-for-python
[Cosmos] [Embedding V0] Public surface: EmbeddingGenerator protocols + embedding_generator on query_items
- Dominant language
- Python
- Stars
- 5.6k
- Forks
- 3.4k
- Avg merge
- 1d 21h
- Merged PRs (30d)
- 193
Description
# 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`):
```python
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_items` – `embedding_generator: Optional["EmbeddingGenerator"] = None`
- `azure/cosmos/aio/_container.py :: ContainerProxy.query_items` – `embedding_generator: Optional["AsyncEmbeddingGenerator"] = None`
3. Stuff into `feed_options`:
```python
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.
Contributor guide
Assessment
This issue has not been assessed yet.