Azure / Azure/azure-sdk-for-python
[Cosmos] [Tracking] Embedding Generation in the Python SDK (V0)
- Dominant language
- Python
- Stars
- 5.6k
- Forks
- 3.4k
- Avg merge
- 1d 21h
- Merged PRs (30d)
- 193
Description
# [Cosmos] [Tracking] Embedding Generation in the Python SDK (V0)
## Summary
V0 plan (parity with .NET tracking issue [Azure/azure-cosmos-dotnet-v3#5830](https://github.com/Azure/azure-cosmos-dotnet-v3/issues/5830)) for adding **automatic embedding generation** for hybrid / vector queries in the Python Cosmos SDK (`azure-cosmos`).
When a customer writes:
```sql
SELECT TOP 10 *
FROM c
ORDER BY VectorDistance(c.text, 'big brown cat')
```
the gateway rewrites it to:
```sql
SELECT TOP 10 *
FROM c
ORDER BY VectorDistance(c.embedding, @documentdb-hybridsearchquery-embedding-0)
```
…and returns an `embeddingParameterMap` mapping `@documentdb-hybridsearchquery-embedding-0` → `'big brown cat'`. The SDK calls the registered `EmbeddingGenerator`, injects the returned vectors as query parameters, and executes per-partition.
## Architecture
### Protocols (core package — `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]]: ...
```
Passed via `embedding_generator=` keyword on `query_items` (sync + async).
### Provider package (`azure-cosmos-embedding`)
Separate optional PyPI package — analogous to the .NET `Microsoft.Azure.Cosmos.Embedding`. Ships `AzureOpenAIEmbeddingGenerator` (implements both protocols). Customers who use a different provider implement the protocol directly. `azure-cosmos` has **no dependency** on this package.
### VectorEmbeddingPolicy extension
The `embeddingSource` block (new in V0) is added to each entry in `vectorEmbeddings`. It carries the endpoint, deployment name, auth type, and source paths:
```json
"embeddingSource": {
"sourcePaths": ["/title", "/abstract"],
"deploymentName": "text-embedding-3-small",
"modelName": "text-embedding-3-small",
"endpoint": "https://embedding-south-central.cognitiveservices.azure.com/",
"authType": "ApiKey"
}
```
A new `EmbeddingSource` TypedDict (and `VectorEmbedding` / `VectorEmbeddingPolicy` typed equivalents) are added to `azure-cosmos` to support typed construction of `AzureOpenAIEmbeddingGenerator.from_embedding_source(...)`.
### Flow
1. SDK → Gateway: query plan request for raw user query.
2. Gateway → SDK: rewritten plan + `embeddingParameterMap` (key → original text).
3. SDK → `EmbeddingGenerator`: one batched call with all texts.
4. SDK → backend: per-partition execution with embeddings injected as parameters.
> **Hard GW dependency:** Python has no `serviceinterop` DLL. All query plans come from the gateway. Python SDK changes can land immediately, but end-to-end testing requires the gateway change to be deployed.
## Sub-issues
### Core package (`azure-cosmos`)
- [x] #46730 — Public surface: `EmbeddingGenerator` / `AsyncEmbeddingGenerator` protocols + `embedding_generator` on `query_items`
- [ ] #46731 — Plan-info accessor + `_QueryFeature.EmbeddingGeneration` + supported-features advertisement
- [ ] #46764 — VectorEmbeddingPolicy: `EmbeddingSource` TypedDict + typed `VectorEmbedding` / `VectorEmbeddingPolicy`
- [ ] #46732 — `_resolve_embeddings` helper on hybrid-search aggregator (sync + async)
- [ ] #46733 — Wire resolver into `_run_hybrid_search`, plumb generator through dispatcher, **and** add diagnostics span (merged from #46734)
- [ ] #46735 — Unit tests
- [ ] #46736 — Emulator tests + sample
### Provider package (`azure-cosmos-embedding`)
- [ ] #46765 — Package scaffold (setup.py, namespace, ci.yml registration)
- [ ] #46766 — `AzureOpenAIEmbeddingGenerator` implementation
- [ ] #46767 — CI/CD and release pipeline
## Open items
1. Confirm wire shape of `embeddingParameterMap` in GW JSON response — object vs array of `{key,value}`.
2. Confirm the exact string for the `EmbeddingGeneration` supported-feature token.
3. `azure-cosmos-embedding` version: ship as `1.0.0b1` (beta) at Build, or stage post-Build?
4. Auth: confirm which auth types (ApiKey, Entra/RBAC) the Foundry / AOAI endpoint supports in V0.
5. Python SDK owner for this feature: Aayush.
## Notes
- **Hard GW dependency:** Python end-to-end tests (#46736) cannot run until the gateway change ships on the release train (R2D process underway).
- **Sync + async dual paths:** every sub-issue covers both sync (`azure.cosmos.container`) and async (`azure.cosmos.aio._container`) surfaces unless stated.
## References
- .NET tracking issue: [Azure/azure-cosmos-dotnet-v3#5830](https://github.com/Azure/azure-cosmos-dotnet-v3/issues/5830)
- Python files touched (core):
- `sdk/cosmos/azure-cosmos/azure/cosmos/documents.py` — `_QueryFeature` token
- `sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py` — sync aggregator
- `sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py` — async aggregator
- `sdk/cosmos/azure-cosmos/azure/cosmos/container.py` — `query_items` keyword
- `sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py` — async `query_items` keyword
Contributor guide
Assessment
This issue has not been assessed yet.