From 3e2cf75cfa662a6640d4e57340057ef745429369 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 7 Jul 2026 12:15:04 -0700 Subject: [PATCH] Adding changes for jey based reranker api --- .../azure-cosmos/azure/cosmos/_constants.py | 2 + .../azure/cosmos/_cosmos_client_connection.py | 4 +- .../azure/cosmos/_inference_service.py | 28 +++++-- .../azure/cosmos/aio/_container.py | 4 +- .../aio/_cosmos_client_connection_async.py | 4 +- .../cosmos/aio/_inference_service_async.py | 24 +++++- .../azure-cosmos/azure/cosmos/container.py | 4 +- .../tests/test_semantic_reranker_unit.py | 73 +++++++++++++++++++ 8 files changed, 131 insertions(+), 12 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py index 73ba0649a859..badf1fe4c734 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py @@ -70,6 +70,8 @@ class _Constants: AAD_DEFAULT_SCOPE: str = "https://cosmos.azure.com/.default" INFERENCE_SERVICE_DEFAULT_SCOPE = "https://dbinference.azure.com/.default" SEMANTIC_RERANKER_INFERENCE_ENDPOINT: str = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT" + SEMANTIC_RERANKER_INFERENCE_KEY: str = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY" + SEMANTIC_RERANKER_INFERENCE_KEY_HEADER: str = "Ocp-Apim-Subscription-Key" # Health Check Retry Policy constants AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES: str = "AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index a5d01a7a12c9..ab6ea4782435 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -350,7 +350,9 @@ def _set_client_consistency_level( def _get_inference_service(self) -> Optional[_InferenceService]: """Get inference service instance, lazily initializing on first access.""" - if self._inference_service is None and self.aad_credentials: + if self._inference_service is None and ( + self.aad_credentials or os.environ.get(Constants.SEMANTIC_RERANKER_INFERENCE_KEY) + ): with self._inference_service_lock: if self._inference_service is None: try: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py index 563493e1b96e..d397a00c703a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py @@ -26,10 +26,11 @@ from urllib3.util.retry import Retry from azure.core import PipelineClient +from azure.core.credentials import AzureKeyCredential from azure.core.exceptions import DecodeError, ServiceRequestError, ServiceResponseError -from azure.core.pipeline.policies import (ContentDecodePolicy, CustomHookPolicy, DistributedTracingPolicy, - HeadersPolicy, HTTPPolicy, NetworkTraceLoggingPolicy, ProxyPolicy, - UserAgentPolicy) +from azure.core.pipeline.policies import (AzureKeyCredentialPolicy, ContentDecodePolicy, CustomHookPolicy, + DistributedTracingPolicy, HeadersPolicy, HTTPPolicy, + NetworkTraceLoggingPolicy, ProxyPolicy, UserAgentPolicy) from azure.core.pipeline.transport import HttpRequest from azure.core.utils import CaseInsensitiveDict @@ -77,14 +78,31 @@ def __init__(self, cosmos_client_connection): self._inference_request_timeout = self._client_connection.connection_policy.InferenceRequestTimeout self._inference_pipeline_client = self._create_inference_pipeline_client() + def _create_auth_policy(self): + """Create the authentication policy for inference requests. + + When the ``AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY`` environment variable is set, key-based + authentication is used and the key is sent in the ``Ocp-Apim-Subscription-Key`` header. Otherwise + the request is authenticated with the client's AAD credentials using a standard bearer token. + + :returns: The authentication policy to use for inference requests. + :rtype: ~azure.core.pipeline.policies.AzureKeyCredentialPolicy or + ~azure.cosmos._inference_auth_policy.InferenceServiceBearerTokenPolicy + """ + inference_key = os.environ.get(Constants.SEMANTIC_RERANKER_INFERENCE_KEY) + if inference_key: + return AzureKeyCredentialPolicy( + AzureKeyCredential(inference_key), Constants.SEMANTIC_RERANKER_INFERENCE_KEY_HEADER + ) + return InferenceServiceBearerTokenPolicy(self._aad_credentials, self._token_scope) + def _create_inference_pipeline_client(self) -> PipelineClient: """Create a pipeline for inference requests. :returns: A PipelineClient configured for inference calls. :rtype: ~azure.core.PipelineClient """ - access_token = self._aad_credentials - auth_policy = InferenceServiceBearerTokenPolicy(access_token, self._token_scope) + auth_policy = self._create_auth_policy() connection_policy = self._client_connection.connection_policy retry_policy = None diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index b4d760b5c5f4..5f1113a5aa66 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1370,7 +1370,9 @@ async def semantic_rerank( inference_service = self.client_connection._get_inference_service() if inference_service is None: raise CosmosHttpResponseError( - message="Semantic reranking requires AAD credentials (inference service not initialized).", + message="Semantic reranking requires AAD credentials or the " + "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY environment variable " + "(inference service not initialized).", response=None ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 1f9f2ca87369..047027e5f2ee 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -308,7 +308,9 @@ def _set_container_properties_cache(self, container_link: str, properties: Optio def _get_inference_service(self) -> Optional[_InferenceService]: """Get async inference service instance, lazily initializing on first access.""" - if self._inference_service is None and self.aad_credentials: + if self._inference_service is None and ( + self.aad_credentials or os.environ.get(Constants.SEMANTIC_RERANKER_INFERENCE_KEY) + ): try: self._inference_service = _InferenceService(self) except ValueError as e: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py index c190bf43526b..20bfe10bd54f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py @@ -27,8 +27,9 @@ from urllib3.util.retry import Retry from azure.core import AsyncPipelineClient +from azure.core.credentials import AzureKeyCredential from azure.core.exceptions import DecodeError, ServiceRequestError, ServiceResponseError -from azure.core.pipeline.policies import (AsyncHTTPPolicy, ContentDecodePolicy, +from azure.core.pipeline.policies import (AsyncHTTPPolicy, AzureKeyCredentialPolicy, ContentDecodePolicy, DistributedTracingPolicy, HeadersPolicy, NetworkTraceLoggingPolicy, ProxyPolicy, UserAgentPolicy) from azure.core.pipeline.transport import HttpRequest @@ -78,14 +79,31 @@ def __init__(self, cosmos_client_connection): self._inference_request_timeout = self._client_connection.connection_policy.InferenceRequestTimeout self._inference_pipeline_client = self._create_inference_pipeline_client() + def _create_auth_policy(self): + """Create the authentication policy for inference requests. + + When the ``AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY`` environment variable is set, key-based + authentication is used and the key is sent in the ``Ocp-Apim-Subscription-Key`` header. Otherwise + the request is authenticated with the client's AAD credentials using a standard bearer token. + + :returns: The authentication policy to use for inference requests. + :rtype: ~azure.core.pipeline.policies.AzureKeyCredentialPolicy or + ~azure.cosmos.aio._inference_auth_policy_async.AsyncInferenceServiceBearerTokenPolicy + """ + inference_key = os.environ.get(Constants.SEMANTIC_RERANKER_INFERENCE_KEY) + if inference_key: + return AzureKeyCredentialPolicy( + AzureKeyCredential(inference_key), Constants.SEMANTIC_RERANKER_INFERENCE_KEY_HEADER + ) + return AsyncInferenceServiceBearerTokenPolicy(self._aad_credentials, self._token_scope) + def _create_inference_pipeline_client(self) -> AsyncPipelineClient: """Create a pipeline for inference requests. :returns: An AsyncPipelineClient configured for inference calls. :rtype: ~azure.core.AsyncPipelineClient """ - access_token = self._aad_credentials - auth_policy = AsyncInferenceServiceBearerTokenPolicy(access_token, self._token_scope) + auth_policy = self._create_auth_policy() connection_policy = self._client_connection.connection_policy diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 665ae2e4f869..1408f11281f9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1140,7 +1140,9 @@ def semantic_rerank( inference_service = self.client_connection._get_inference_service() if inference_service is None: raise CosmosHttpResponseError( - message="Semantic reranking requires AAD credentials (inference service not initialized).", + message="Semantic reranking requires AAD credentials or the " + "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_KEY environment variable " + "(inference service not initialized).", response=None ) diff --git a/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py b/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py index 91f7d570924b..57aa5aaa774b 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py @@ -10,17 +10,24 @@ from unittest.mock import MagicMock, patch from azure.core.exceptions import DecodeError, ServiceRequestError, ServiceResponseError +from azure.core.pipeline import PipelineContext, PipelineRequest +from azure.core.pipeline.policies import AzureKeyCredentialPolicy +from azure.core.pipeline.transport import HttpRequest 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" @@ -33,6 +40,8 @@ def setUp(self): os.environ[_INFERENCE_ENDPOINT_ENV_VAR] = "https://example.com" self._saved_malformed = os.environ.get(_MALFORMED_INPUT_ENV_VAR) os.environ.pop(_MALFORMED_INPUT_ENV_VAR, None) + self._saved_key = os.environ.get(_INFERENCE_KEY_ENV_VAR) + os.environ.pop(_INFERENCE_KEY_ENV_VAR, None) def tearDown(self): """Restore the inference endpoint env var after each test.""" @@ -44,6 +53,10 @@ def tearDown(self): os.environ[_MALFORMED_INPUT_ENV_VAR] = self._saved_malformed else: os.environ.pop(_MALFORMED_INPUT_ENV_VAR, None) + if self._saved_key is not None: + os.environ[_INFERENCE_KEY_ENV_VAR] = self._saved_key + else: + os.environ.pop(_INFERENCE_KEY_ENV_VAR, None) def _create_mock_connection(self, inference_request_timeout=5): """Create a mock cosmos client connection with configurable inference timeout.""" @@ -589,6 +602,66 @@ def test_async_get_inference_service_caches_instance(self): self.assertIsNotNone(first) self.assertIs(first, second) + # ── key-based authentication tests ── + + def test_sync_inference_uses_key_auth_policy_when_key_env_var_set(self): + """Test that the sync inference service uses key-based auth when the key env var is set.""" + os.environ[_INFERENCE_KEY_ENV_VAR] = "test-key" + service = _SyncInferenceService(self._create_mock_connection()) + self.assertIsInstance(service._create_auth_policy(), AzureKeyCredentialPolicy) + + def test_async_inference_uses_key_auth_policy_when_key_env_var_set(self): + """Test that the async inference service uses key-based auth when the key env var is set.""" + os.environ[_INFERENCE_KEY_ENV_VAR] = "test-key" + mock_connection = self._create_mock_connection() + mock_connection.connection_policy.DisableSSLVerification = False + service = _AsyncInferenceService(mock_connection) + self.assertIsInstance(service._create_auth_policy(), AzureKeyCredentialPolicy) + + def test_sync_inference_uses_bearer_auth_policy_without_key_env_var(self): + """Test that the sync inference service falls back to AAD bearer auth without the key env var.""" + service = _SyncInferenceService(self._create_mock_connection()) + self.assertIsInstance(service._create_auth_policy(), InferenceServiceBearerTokenPolicy) + + def test_async_inference_uses_bearer_auth_policy_without_key_env_var(self): + """Test that the async inference service falls back to AAD bearer auth without the key env var.""" + mock_connection = self._create_mock_connection() + mock_connection.connection_policy.DisableSSLVerification = False + service = _AsyncInferenceService(mock_connection) + self.assertIsInstance(service._create_auth_policy(), AsyncInferenceServiceBearerTokenPolicy) + + def test_key_auth_policy_sets_subscription_key_header(self): + """Test that key-based auth injects the key into the Ocp-Apim-Subscription-Key header.""" + os.environ[_INFERENCE_KEY_ENV_VAR] = "super-secret-key" + service = _SyncInferenceService(self._create_mock_connection()) + request = PipelineRequest(HttpRequest("POST", "https://example.com"), PipelineContext(None)) + service._create_auth_policy().on_request(request) + self.assertEqual(request.http_request.headers.get(_INFERENCE_KEY_HEADER), "super-secret-key") + + def test_sync_get_inference_service_returns_service_with_key_and_no_aad(self): + """Test that _get_inference_service() creates a service via key auth when no AAD + credentials are present but the key env var is set.""" + os.environ[_INFERENCE_KEY_ENV_VAR] = "test-key" + mock_conn = self._create_mock_connection() + mock_conn.aad_credentials = None + mock_conn._inference_service = None + mock_conn._inference_service_lock = threading.Lock() + + result = _SyncCosmosClientConnection._get_inference_service(mock_conn) + self.assertIsNotNone(result) + + def test_async_get_inference_service_returns_service_with_key_and_no_aad(self): + """Test that async _get_inference_service() creates a service via key auth when no AAD + credentials are present but the key env var is set.""" + os.environ[_INFERENCE_KEY_ENV_VAR] = "test-key" + mock_conn = self._create_mock_connection() + mock_conn.aad_credentials = None + mock_conn._inference_service = None + mock_conn.connection_policy.DisableSSLVerification = False + + result = _AsyncCosmosClientConnection._get_inference_service(mock_conn) + self.assertIsNotNone(result) + if __name__ == "__main__": unittest.main()