Merge e7b92ce1c2 into c2c2a623c0

litellm/proxy/proxy_server.py

@@ -12984,10 +12984,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             verbose_proxy_logger.exception(
                 "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment"
             )
-    if deployment is not None:
-        litellm_model_name = deployment.get("litellm_params", {}).get("model")
-        model_info = deployment.get("model_info", {})
-        load_credentials_from_list(deployment.get("litellm_params", {}))
+    request_deployment: Final = copy.deepcopy(deployment)
+    if request_deployment is not None:
+        litellm_model_name = request_deployment.get("litellm_params", {}).get("model")
+        model_info = request_deployment.get("model_info", {})
+        load_credentials_from_list(request_deployment.get("litellm_params", {}))
         # remove the custom_llm_provider_prefix in the litellm_model_name
         if "/" in litellm_model_name:
             litellm_model_name = litellm_model_name.split("/", 1)[1]
@@ -12999,9 +13000,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     # Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
     provider_counter: BaseTokenCounter | None = None
     custom_llm_provider: str | None = None
-    if call_endpoint is True and deployment is not None:
+    if call_endpoint is True and request_deployment is not None:
         # Auto-route to the correct provider based on model
-        provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use)
+        provider_counter, _model, custom_llm_provider = _get_provider_token_counter(request_deployment, model_to_use)
         if _model is not None:
             model_to_use = _model
 
@@ -13012,7 +13013,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             model_to_use=model_to_use,
             messages=messages,
             contents=contents,
-            deployment=deployment,
+            deployment=request_deployment,
             request_model=request.model,
             tools=tools,
             system=system,

tests/test_litellm/proxy/proxy_server/test_routes_utils.py

@@ -10,12 +10,14 @@ from __future__ import annotations
 
 import asyncio
 import json
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 import litellm
 from litellm.proxy import proxy_server
 from litellm.router_utils import pattern_match_deployments
+from litellm.types.utils import CredentialItem, TokenCountResponse
 
 from .conftest import normalize  # type: ignore[import-not-found]
 
@@ -93,6 +95,61 @@ def test_token_counter_missing_input_returns_400(
     assert "prompt or messages or contents" in response.text
 
 
+def test_token_counter_named_credentials_do_not_mutate_router_deployment(client, auth_as, monkeypatch):
+    """Provider token counting must not pin rotating credentials to the router."""
+    named_credential = CredentialItem(
+        credential_name="rotating-bedrock",
+        credential_values={"aws_session_token": "temporary-session-token"},
+        credential_info={"custom_llm_provider": "bedrock"},
+    )
+    router_deployment = {
+        "litellm_params": {
+            "model": "bedrock/anthropic.claude-3-sonnet",
+            "litellm_credential_name": "rotating-bedrock",
+        },
+        "model_info": {},
+    }
+
+    mock_router = MagicMock()
+    mock_router.async_get_available_deployment = AsyncMock(return_value=router_deployment)
+    mock_counter = MagicMock()
+    mock_counter.should_use_token_counting_api.return_value = True
+    mock_counter.count_tokens = AsyncMock(
+        return_value=TokenCountResponse(
+            total_tokens=1,
+            request_model="claude-bedrock",
+            model_used="anthropic.claude-3-sonnet",
+            tokenizer_type="bedrock_api",
+        )
+    )
+
+    monkeypatch.setattr(litellm, "credential_list", [named_credential])
+    monkeypatch.setattr(proxy_server, "llm_router", mock_router)
+    monkeypatch.setattr(
+        proxy_server,
+        "_get_provider_token_counter",
+        lambda deployment, model: (
+            mock_counter,
+            "anthropic.claude-3-sonnet",
+            "bedrock",
+        ),
+    )
+
+    with auth_as():
+        response = client.post(
+            "/utils/token_counter?call_endpoint=true",
+            json={
+                "model": "claude-bedrock",
+                "messages": [{"role": "user", "content": "hello"}],
+            },
+        )
+
+    assert response.status_code == 200
+    hydrated_deployment = mock_counter.count_tokens.await_args.kwargs["deployment"]
+    assert hydrated_deployment["litellm_params"]["aws_session_token"] == "temporary-session-token"
+    assert "aws_session_token" not in router_deployment["litellm_params"]
+
+
 # ---------------------------------------------------------------------------
 # GET /utils/supported_openai_params
 # ---------------------------------------------------------------------------
@@ -117,9 +174,7 @@ def patched_supported_params(monkeypatch):
 def test_supported_openai_params_happy_path(client, auth_as, patched_supported_params):
     """Pins ``GET /utils/supported_openai_params``."""
     with auth_as():
-        response = client.get(
-            "/utils/supported_openai_params", params={"model": "gpt-4"}
-        )
+        response = client.get("/utils/supported_openai_params", params={"model": "gpt-4"})
     assert response.status_code == 200
     assert normalize(response.json()) == {
         "supported_openai_params": ["max_tokens", "temperature", "top_p"],