Skip to content

[WIP] Semantic Reranking: Add key-based authentication via environment variable#47917

Open
aayush3011 wants to merge 2 commits into
Azure:mainfrom
aayush3011:users/akataria/reranker_key_based
Open

[WIP] Semantic Reranking: Add key-based authentication via environment variable#47917
aayush3011 wants to merge 2 commits into
Azure:mainfrom
aayush3011:users/akataria/reranker_key_based

Conversation

@aayush3011

@aayush3011 aayush3011 commented Jul 7, 2026

Copy link
Copy Markdown
Member

Description

Adds key-based authentication for the semantic reranker inference service so key-based auth can be validated end-to-end without AAD credentials.

When the new AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY environment variable is set, the inference pipeline authenticates by sending the key in the Ocp-Apim-Subscription-Key header via AzureKeyCredentialPolicy. When the variable is not set, behavior is unchanged - the service continues to use the existing AAD bearer-token policy.

All SDK Contribution checklist:

  • The pull request does not introduce [breaking changes]
  • CHANGELOG is updated for new features, bug fixes or other significant changes.
  • I have read the contribution guidelines.

General Guidelines and Best Practices

  • Title of the pull request is clear and informative.
  • There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

  • Pull request includes test coverage for the included changes.

Copilot AI review requested due to automatic review settings July 7, 2026 19:17
@aayush3011 aayush3011 requested a review from a team as a code owner July 7, 2026 19:17
@github-actions github-actions Bot added the Cosmos label Jul 7, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds an opt-in key-based authentication path for the Cosmos semantic reranker inference calls, enabled via the AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY environment variable, while preserving the existing AAD bearer-token behavior when the key is not set.

Changes:

  • Add key-auth support to sync/async inference pipelines using AzureKeyCredentialPolicy with the Ocp-Apim-Subscription-Key header when the key env var is present.
  • Allow lazy inference service initialization when either AAD credentials or the key env var is present (sync + async client connections).
  • Extend unit tests to validate key-auth policy selection and header injection, and update user-facing error messages to mention the key env var option.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py Adds unit tests covering key-based auth selection and header injection; includes env var setup/teardown for the new key.
sdk/cosmos/azure-cosmos/azure/cosmos/container.py Updates the sync semantic_rerank error message to include the key env var authentication option.
sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py Adds key-based auth policy creation and uses it when building the async inference pipeline.
sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py Enables async inference service lazy-init when the key env var is set (even without AAD credentials).
sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py Updates the async semantic_rerank error message to include the key env var authentication option.
sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py Adds key-based auth policy creation and uses it when building the sync inference pipeline.
sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py Enables sync inference service lazy-init when the key env var is set (even without AAD credentials).
sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py Introduces constants for the key env var name and subscription key header name.

Comment on lines 17 to 31
import azure.cosmos.exceptions as exceptions
from azure.cosmos._cosmos_client_connection import CosmosClientConnection as _SyncCosmosClientConnection
from azure.cosmos._inference_auth_policy import InferenceServiceBearerTokenPolicy
from azure.cosmos._inference_service import _InferenceService as _SyncInferenceService
from azure.cosmos.aio._cosmos_client_connection_async import (
CosmosClientConnection as _AsyncCosmosClientConnection,
)
from azure.cosmos.aio._inference_auth_policy_async import AsyncInferenceServiceBearerTokenPolicy
from azure.cosmos.aio._inference_service_async import _InferenceService as _AsyncInferenceService
from azure.cosmos.documents import ConnectionPolicy

_INFERENCE_ENDPOINT_ENV_VAR = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT"
_INFERENCE_KEY_ENV_VAR = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY"
_INFERENCE_KEY_HEADER = "Ocp-Apim-Subscription-Key"
_MALFORMED_INPUT_ENV_VAR = "AZURE_COSMOS_CHARSET_DECODER_ERROR_ACTION_ON_MALFORMED_INPUT"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants