From 194270d5a22ba69e9bfc13704425bbe8f3643515 Mon Sep 17 00:00:00 2001 From: Lovro Seder Date: Sun, 1 Feb 2026 00:26:53 +0100 Subject: [PATCH] Fix Azure AI Anthropic CountTokens 401 auth error (#20069) Add x-api-key header to CountTokens handler to match chat completion authentication. Azure AI Anthropic requires this header per Microsoft's native API format. --- .../anthropic/count_tokens/transformation.py | 14 ++- ...e_anthropic_count_tokens_transformation.py | 111 ++++++++++++++++++ 2 files changed, 119 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index e284595cc8a..09b83b7c971 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -30,30 +30,32 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): """ Get the required headers for the Azure AI Anthropic CountTokens API. - Uses Azure authentication (api-key header) instead of Anthropic's x-api-key. + Azure AI Anthropic uses Anthropic's native API format, which requires the + x-api-key header for authentication (in addition to Azure's api-key header). Args: api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config Returns: - Dictionary of required headers with Azure authentication + Dictionary of required headers with both x-api-key and Azure authentication """ - # Start with base headers + # Start with base headers including x-api-key for Anthropic API compatibility headers = { "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + "x-api-key": api_key, # Azure AI Anthropic requires this header } - # Use Azure authentication + # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} if "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj = GenericLiteLLMParams(**litellm_params) - # Get Azure auth headers + # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM._base_validate_azure_environment( headers={}, litellm_params=litellm_params_obj ) @@ -68,7 +70,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Get the Azure AI Anthropic CountTokens API endpoint. Args: - api_base: The Azure AI API base URL + api_base: The Azure AI API base URL (e.g., https://my-resource.services.ai.azure.com or https://my-resource.services.ai.azure.com/anthropic) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py new file mode 100644 index 00000000000..78806831685 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -0,0 +1,111 @@ +""" +Tests for Azure AI Anthropic CountTokens transformation. + +Verifies that the CountTokens API uses the correct authentication headers. +""" +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + + +from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( + AzureAIAnthropicCountTokensConfig, +) + + +class TestAzureAIAnthropicCountTokensConfig: + """Test Azure AI Anthropic CountTokens configuration and headers.""" + + def test_get_required_headers_includes_x_api_key(self): + """ + Test that get_required_headers includes x-api-key header. + + Azure AI Anthropic uses Anthropic's native API format which requires + the x-api-key header for authentication (not just Azure's api-key). + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-api-key-12345" + + headers = config.get_required_headers(api_key=api_key) + + # Verify x-api-key header is set + assert "x-api-key" in headers + assert headers["x-api-key"] == api_key + + # Verify base headers are present + assert headers["Content-Type"] == "application/json" + assert headers["anthropic-version"] == "2023-06-01" + assert "anthropic-beta" in headers + + def test_get_required_headers_includes_azure_api_key(self): + """ + Test that get_required_headers includes Azure api-key header. + + Both x-api-key and api-key headers should be present. + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-azure-key-67890" + + headers = config.get_required_headers(api_key=api_key) + + # Verify both authentication headers are set + assert "x-api-key" in headers + assert "api-key" in headers + assert headers["x-api-key"] == api_key + assert headers["api-key"] == api_key + + def test_get_required_headers_with_litellm_params(self): + """ + Test that get_required_headers works with litellm_params. + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-key" + litellm_params = {"api_key": "param-key", "custom_field": "value"} + + headers = config.get_required_headers( + api_key=api_key, litellm_params=litellm_params + ) + + # x-api-key should use the direct api_key parameter + assert headers["x-api-key"] == api_key + # Azure api-key should come from litellm_params + assert headers["api-key"] == "param-key" + + def test_get_count_tokens_endpoint_with_base_url(self): + """Test endpoint generation from base URL.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + ) + + def test_get_count_tokens_endpoint_with_anthropic_path(self): + """Test endpoint generation when base URL already includes /anthropic.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com/anthropic" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + ) + + def test_get_count_tokens_endpoint_with_trailing_slash(self): + """Test endpoint generation with trailing slash in base URL.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com/" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + )