Azure / Azure/azure-sdk-for-python
[Cosmos] [Embedding V0] Plan-info accessor + _QueryFeature.EmbeddingGeneration + supported-features advertisement
- Dominant language
- Python
- Stars
- 5.6k
- Forks
- 3.4k
- Avg merge
- 1d 21h
- Merged PRs (30d)
- 193
Description
# Plan DTO accessor + `_QueryFeature.EmbeddingGeneration` + supported-features advertisement
Parent: 46729
## Goal
Teach the SDK to (a) read the new `embeddingParameterMap` field out of a query-plan response and (b) advertise the new feature so the gateway emits it.
## Scope
1. `azure/cosmos/documents.py` – add to `_QueryFeature`:
```python
EmbeddingGeneration: Literal["EmbeddingGeneration"] = "EmbeddingGeneration"
```
2. `azure/cosmos/_execution_context/query_execution_info.py` – new accessor on `_PartitionedQueryExecutionInfo`:
```python
EmbeddingParameterMapPath = [QueryInfoPath, "embeddingParameterMap"]
HybridEmbeddingParameterMapPath = [HybridSearchQueryInfoPath, "embeddingParameterMap"]
def get_embedding_parameter_map(self):
return (self._extract(self.HybridEmbeddingParameterMapPath)
or self._extract(self.EmbeddingParameterMapPath)
or {})
```
3. `azure/cosmos/_cosmos_client_connection.py :: _GetQueryPlanThroughGateway` (sync) – build `supported_query_features` per-request, conditionally appending `documents._QueryFeature.EmbeddingGeneration` when `options.get("embeddingGenerator") is not None`.
4. `azure/cosmos/aio/_cosmos_client_connection_async.py :: _GetQueryPlanThroughGateway` (async) – mirror the same change.
5. Move the now-conditional supported-features string out of the inline ladder into a small helper to avoid further duplication.
## Non-goals
- Do NOT yet wire the accessor into the aggregator (covered by 46733).
- Do NOT change the on-wire JSON parser; the existing `_extract` walk handles new keys.
## Files touched
- `sdk/cosmos/azure-cosmos/azure/cosmos/documents.py`
- `sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/query_execution_info.py`
- `sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py`
- `sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py`
- `sdk/cosmos/azure-cosmos/CHANGELOG.md`
## Acceptance
- A canned plan-response dict containing `embeddingParameterMap` returns the right map from `get_embedding_parameter_map()`.
- A plan-response dict with no map returns `{}` (not `None`) so callers can iterate freely.
- When `embedding_generator` is configured, the outgoing `supportedQueryFeatures` header contains `EmbeddingGeneration`; otherwise it does not.
- All existing query-plan tests still pass.
Contributor guide
Assessment
This issue has not been assessed yet.