[WIP] Semantic Reranking: Add key-based authentication via environment variable#47917
Open
aayush3011 wants to merge 2 commits into
Open
[WIP] Semantic Reranking: Add key-based authentication via environment variable#47917aayush3011 wants to merge 2 commits into
aayush3011 wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
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
AzureKeyCredentialPolicywith theOcp-Apim-Subscription-Keyheader 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" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_KEYenvironment variable is set, the inference pipeline authenticates by sending the key in theOcp-Apim-Subscription-Keyheader viaAzureKeyCredentialPolicy. When the variable is not set, behavior is unchanged - the service continues to use the existing AAD bearer-token policy.All SDK Contribution checklist:
General Guidelines and Best Practices
Testing Guidelines