Skip to content

[Cosmos] [Embedding V0] Plan-info accessor + _QueryFeature.EmbeddingGeneration + supported-features advertisement #46731

@ananth7592

Description

@ananth7592

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:
    EmbeddingGeneration: Literal["EmbeddingGeneration"] = "EmbeddingGeneration"
  2. azure/cosmos/_execution_context/query_execution_info.py – new accessor on _PartitionedQueryExecutionInfo:
    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.

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