mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(anthropic): count tokens through the federated credential
A federated deployment holds no static key, so the token counter found none, returned None, and the caller silently fell back to the local tokenizer. A workload identity deployment could therefore never get Anthropic's authoritative count. It now mints one when no static key is configured. The minted token is an sk-ant-oat, which the existing header builder already sends as a Bearer rather than an x-api-key, so nothing else on that path changes.
This commit is contained in:
parent
0a551902ca
commit
33ed7b14bb
2 changed files with 63 additions and 5 deletions
|
|
@ -54,13 +54,19 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
|
||||
# Get Anthropic API key from deployment config or environment
|
||||
api_key = litellm_params.get("api_key")
|
||||
if not api_key:
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
from litellm.llms.anthropic.wif import aget_anthropic_wif_token
|
||||
|
||||
static_key: Final = litellm_params.get("api_key") or os.getenv("ANTHROPIC_API_KEY")
|
||||
# A federated deployment holds no static key by design. Without a minted one this returns
|
||||
# None and the caller silently falls back to the local tokenizer, so a workload identity
|
||||
# deployment would never reach Anthropic's authoritative count. The minted token is an
|
||||
# sk-ant-oat, which get_required_headers already sends as a Bearer rather than x-api-key.
|
||||
api_key: Final = static_key or await aget_anthropic_wif_token(
|
||||
litellm_params, litellm_params.get("api_base"), model_to_use
|
||||
)
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No Anthropic API key found for token counting")
|
||||
verbose_logger.warning("No Anthropic credential found for token counting")
|
||||
return None
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ Regression test for https://github.com/BerriAI/litellm/issues/22040
|
|||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
|
||||
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
|
|
@ -78,3 +80,53 @@ class TestCountTokensOAuthHeaders:
|
|||
beta_value = headers.get("anthropic-beta", "")
|
||||
assert "token-counting" in beta_value, f"token-counting beta missing from OAuth headers: {beta_value}"
|
||||
assert "oauth-2025-04-20" in beta_value, f"oauth beta missing from OAuth headers: {beta_value}"
|
||||
|
||||
|
||||
class TestCountTokensUsesWorkloadIdentity:
|
||||
"""A federated deployment holds no static key. Without minting one, count_tokens returns None
|
||||
and the caller silently falls back to the local tokenizer, so the number a federated
|
||||
deployment reports would never come from Anthropic."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_federated_deployment_mints_and_counts(self, monkeypatch):
|
||||
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
|
||||
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
minted = "sk-ant-oat01-minted-for-count"
|
||||
|
||||
async def fake_mint(_params, _api_base, _model):
|
||||
return minted
|
||||
|
||||
monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
async def fake_request(**kwargs):
|
||||
seen.update(kwargs)
|
||||
return {"input_tokens": 42}
|
||||
|
||||
monkeypatch.setattr(
|
||||
token_counter_module.anthropic_count_tokens_handler,
|
||||
"handle_count_tokens_request",
|
||||
fake_request,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "fdrl_x",
|
||||
"anthropic_organization_id": "org-x",
|
||||
}
|
||||
},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.total_tokens == 42
|
||||
assert seen["api_key"] == minted
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue