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

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"

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:

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

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
             )
 

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:

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
 

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
             )
 

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()