From 1411a227aaa9c4eeb48c0ad3bcd591a4ca9b0c96 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Jan 2026 13:58:32 -0800 Subject: [PATCH 01/45] bulk update keys endpoint --- litellm/proxy/_types.py | 2 + .../key_management_endpoints.py | 367 ++++++++++++++ .../key_management_endpoints.py | 42 ++ .../test_key_management_endpoints.py | 467 ++++++++++++++++++ 4 files changed, 878 insertions(+) create mode 100644 litellm/types/proxy/management_endpoints/key_management_endpoints.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c854d81ec7..f1f2c259f1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -227,6 +227,7 @@ class KeyManagementRoutes(str, enum.Enum): KEY_REGENERATE_WITH_PATH_PARAM = "/key/{key_id}/regenerate" KEY_BLOCK = "/key/block" KEY_UNBLOCK = "/key/unblock" + KEY_BULK_UPDATE = "/key/bulk_update" # info and health routes KEY_INFO = "/key/info" @@ -494,6 +495,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_LIST.value, KeyManagementRoutes.KEY_BLOCK.value, KeyManagementRoutes.KEY_UNBLOCK.value, + KeyManagementRoutes.KEY_BULK_UPDATE.value, ] management_routes = [ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ab87e862ea..f9fb0e5f49 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -37,6 +37,13 @@ from litellm.proxy._experimental.mcp_server.db import ( ) from litellm.proxy._types import * from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + BulkUpdateKeyResponse, + FailedKeyUpdate, + SuccessfulKeyUpdate, +) from litellm.proxy.auth.auth_checks import ( _cache_key_object, _delete_cache_key_object, @@ -1438,6 +1445,205 @@ def is_different_team( return data.team_id != existing_key_row.team_id +def _validate_max_budget(max_budget: Optional[float]) -> None: + """ + Validate that max_budget is not negative. + + Args: + max_budget: The max_budget value to validate + + Raises: + HTTPException: If max_budget is negative + """ + if max_budget is not None and max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {max_budget}" + }, + ) + + +async def _get_and_validate_existing_key( + token: str, prisma_client: Optional[PrismaClient] +) -> LiteLLM_VerificationToken: + """ + Get existing key from database and validate it exists. + + Args: + token: The key token to look up + prisma_client: Prisma client instance + + Returns: + LiteLLM_VerificationToken: The existing key row + + Raises: + HTTPException: If key is not found + """ + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + + existing_key_row = await prisma_client.get_data( + token=token, + table_name="key", + query_type="find_unique", + ) + + if existing_key_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Key not found: {token}"}, + ) + + return existing_key_row + + +async def _process_single_key_update( + key_update_item: BulkUpdateKeyRequestItem, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: Any, + llm_router: Optional[Router], +) -> Dict[str, Any]: + """ + Process a single key update with all validations and checks. + + This function encapsulates all the logic for updating a single key, + including validation, permission checks, team checks, and database updates. + + Args: + key_update_item: The key update request item + user_api_key_dict: The authenticated user's API key info + litellm_changed_by: Optional header for tracking who made the change + prisma_client: Prisma client instance + user_api_key_cache: User API key cache + proxy_logging_obj: Proxy logging object + llm_router: LLM router instance + + Returns: + Dict containing the updated key information + + Raises: + HTTPException: For various validation and permission errors + """ + # Validate max_budget + _validate_max_budget(key_update_item.max_budget) + + # Get and validate existing key + existing_key_row = await _get_and_validate_existing_key( + token=key_update_item.key, + prisma_client=prisma_client, + ) + + # Check team member permissions + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, + ) + + # Create UpdateKeyRequest from BulkUpdateKeyRequestItem + update_key_request = UpdateKeyRequest( + key=key_update_item.key, + budget_id=key_update_item.budget_id, + max_budget=key_update_item.max_budget, + team_id=key_update_item.team_id, + tags=key_update_item.tags, + ) + + # Get team object and check team limits if team_id is provided + team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + if update_key_request.team_id is not None: + team_obj = await get_team_object( + team_id=update_key_request.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None: + await _check_team_key_limits( + team_table=team_obj, + data=update_key_request, + prisma_client=prisma_client, + ) + + # Validate team change if team is being changed + if is_different_team( + data=update_key_request, existing_key_row=existing_key_row + ): + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + if team_obj is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Team object not found for team change validation" + }, + ) + validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) + + # Prepare update data + non_default_values = await prepare_key_update_data( + data=update_key_request, existing_key_row=existing_key_row + ) + + # Update key in database + _data = {**non_default_values, "token": key_update_item.key} + response = await prisma_client.update_data( + token=key_update_item.key, data=_data + ) + + # Delete cache + await _delete_cache_key_object( + hashed_token=hash_token(key_update_item.key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # Trigger async hook + asyncio.create_task( + KeyManagementEventHooks.async_key_updated_hook( + data=update_key_request, + existing_key_row=existing_key_row, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + + if response is None: + raise ValueError("Failed to update key got response = None") + + # Extract and format updated key info + updated_key_info = response.get("data", {}) + if hasattr(updated_key_info, "model_dump"): + updated_key_info = updated_key_info.model_dump() + elif hasattr(updated_key_info, "dict"): + updated_key_info = updated_key_info.dict() + + updated_key_info.pop("token", None) + + return updated_key_info + + @router.post( "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1684,6 +1890,167 @@ async def update_key_fn( ) +@router.post( + "/key/bulk_update", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + response_model=BulkUpdateKeyResponse, +) +@management_endpoint_wrapper +async def bulk_update_keys( + data: BulkUpdateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Bulk update multiple keys at once. + + This endpoint allows updating multiple keys in a single request. Each key update + is processed independently - if some updates fail, others will still succeed. + + Parameters: + - keys: List[BulkUpdateKeyRequestItem] - List of key update requests, each containing: + - key: str - The key identifier (token) to update + - budget_id: Optional[str] - Budget ID associated with the key + - max_budget: Optional[float] - Max budget for key + - team_id: Optional[str] - Team ID associated with key + - tags: Optional[List[str]] - Tags for organizing keys + + Returns: + - total_requested: int - Total number of keys requested for update + - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info + - failed_updates: List[FailedKeyUpdate] - List of failed updates with key_info and failed_reason + + Example request: + ```bash + curl --location 'http://0.0.0.0:4000/key/bulk_update' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "keys": [ + { + "key": "sk-1234", + "max_budget": 100.0, + "team_id": "team-123", + "tags": ["production", "api"] + }, + { + "key": "sk-5678", + "budget_id": "budget-456", + "tags": ["staging"] + } + ] + }' + ``` + """ + from litellm.proxy.proxy_server import ( + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins can perform bulk key updates" + }, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + + if not data.keys: + raise HTTPException( + status_code=400, + detail={"error": "No keys provided for update"}, + ) + + MAX_BATCH_SIZE = 500 + if len(data.keys) > MAX_BATCH_SIZE: + raise HTTPException( + status_code=400, + detail={ + "error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys." + }, + ) + + successful_updates: List[SuccessfulKeyUpdate] = [] + failed_updates: List[FailedKeyUpdate] = [] + + for key_update_item in data.keys: + try: + # Process single key update using reusable function + updated_key_info = await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + ) + + successful_updates.append( + SuccessfulKeyUpdate( + key=key_update_item.key, + key_info=updated_key_info, + ) + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Failed to update key {key_update_item.key}: {e}" + ) + + if isinstance(e, HTTPException): + error_detail = e.detail + if isinstance(error_detail, dict): + error_message = error_detail.get("error", str(e)) + else: + error_message = str(error_detail) + else: + error_message = str(e) + + key_info = None + try: + existing_key_row = await prisma_client.get_data( + token=key_update_item.key, + table_name="key", + query_type="find_unique", + ) + if existing_key_row is not None: + if hasattr(existing_key_row, "model_dump"): + key_info = existing_key_row.model_dump() + elif hasattr(existing_key_row, "dict"): + key_info = existing_key_row.dict() + if key_info: + key_info.pop("token", None) + except Exception: + pass + + failed_updates.append( + FailedKeyUpdate( + key=key_update_item.key, + key_info=key_info, + failed_reason=error_message, + ) + ) + + return BulkUpdateKeyResponse( + total_requested=len(data.keys), + successful_updates=successful_updates, + failed_updates=failed_updates, + ) + + def validate_key_team_change( key: LiteLLM_VerificationToken, team: LiteLLM_TeamTable, diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py new file mode 100644 index 0000000000..b1d25455d1 --- /dev/null +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -0,0 +1,42 @@ +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel + + +class BulkUpdateKeyRequestItem(BaseModel): + """Individual key update request item""" + + key: str # Key identifier (token) + budget_id: Optional[str] = None # Budget ID associated with the key + max_budget: Optional[float] = None # Max budget for key + team_id: Optional[str] = None # Team ID associated with key + tags: Optional[List[str]] = None # Tags for organizing keys + + +class BulkUpdateKeyRequest(BaseModel): + """Request for bulk key updates""" + + keys: List[BulkUpdateKeyRequestItem] + + +class SuccessfulKeyUpdate(BaseModel): + """Successfully updated key with its updated information""" + + key: str + key_info: Dict[str, Any] + + +class FailedKeyUpdate(BaseModel): + """Failed key update with reason""" + + key: str + key_info: Optional[Dict[str, Any]] = None + failed_reason: str + + +class BulkUpdateKeyResponse(BaseModel): + """Response for bulk key update operations""" + + total_requested: int + successful_updates: List[SuccessfulKeyUpdate] + failed_updates: List[FailedKeyUpdate] diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 7d31f76209..a57378e579 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -30,10 +30,13 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, _check_team_key_limits, _common_key_generation_helper, + _get_and_validate_existing_key, _list_key_helper, _persist_deleted_verification_tokens, + _process_single_key_update, _save_deleted_verification_token_records, _transform_verification_tokens_to_deleted_records, + _validate_max_budget, can_modify_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, @@ -4223,3 +4226,467 @@ async def test_update_key_with_router_settings(monkeypatch): # Verify router_settings can be deserialized and matches input deserialized_settings = json.loads(result["router_settings"]) assert deserialized_settings == router_settings_data + + +@pytest.mark.asyncio +async def test_validate_max_budget(): + """ + Test _validate_max_budget helper function. + + Tests: + 1. Positive max_budget should pass + 2. Zero max_budget should pass + 3. Negative max_budget should raise HTTPException + 4. None max_budget should pass + """ + from fastapi import HTTPException + + # Test Case 1: Positive max_budget should pass + try: + _validate_max_budget(100.0) + _validate_max_budget(0.0) + except HTTPException: + pytest.fail("_validate_max_budget raised HTTPException for valid values") + + # Test Case 2: None max_budget should pass + try: + _validate_max_budget(None) + except HTTPException: + pytest.fail("_validate_max_budget raised HTTPException for None") + + # Test Case 3: Negative max_budget should raise HTTPException + with pytest.raises(HTTPException) as exc_info: + _validate_max_budget(-10.0) + + assert exc_info.value.status_code == 400 + assert "max_budget cannot be negative" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_get_and_validate_existing_key(): + """ + Test _get_and_validate_existing_key helper function. + + Tests: + 1. Successfully retrieve existing key + 2. Key not found raises HTTPException + 3. Database not connected raises HTTPException + """ + from fastapi import HTTPException + + # Test Case 1: Successfully retrieve existing key + mock_prisma_client = AsyncMock() + mock_key = LiteLLM_VerificationToken( + token="test-key-123", + user_id="user-123", + models=["gpt-4"], + team_id=None, + ) + mock_prisma_client.get_data = AsyncMock(return_value=mock_key) + + result = await _get_and_validate_existing_key( + token="test-key-123", + prisma_client=mock_prisma_client, + ) + + assert result == mock_key + mock_prisma_client.get_data.assert_called_once_with( + token="test-key-123", + table_name="key", + query_type="find_unique", + ) + + # Test Case 2: Key not found raises HTTPException + mock_prisma_client.get_data = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc_info: + await _get_and_validate_existing_key( + token="non-existent-key", + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 404 + assert "Key not found" in str(exc_info.value.detail) + + # Test Case 3: Database not connected raises HTTPException + with pytest.raises(HTTPException) as exc_info: + await _get_and_validate_existing_key( + token="test-key-123", + prisma_client=None, + ) + + assert exc_info.value.status_code == 500 + assert "Database not connected" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_process_single_key_update(): + """ + Test _process_single_key_update helper function. + + Tests successful key update with all validations passing. + """ + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequestItem, + ) + + # Setup mocks + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_llm_router = MagicMock() + + # Mock existing key + existing_key = LiteLLM_VerificationToken( + token="test-key-123", + user_id="user-123", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + # Mock updated key response + updated_key_data = { + "user_id": "user-123", + "models": ["gpt-4"], + "team_id": None, + "max_budget": 100.0, + "tags": ["production"], + } + + mock_prisma_client.get_data = AsyncMock(return_value=existing_key) + mock_updated_key_obj = MagicMock() + mock_updated_key_obj.model_dump.return_value = updated_key_data + mock_prisma_client.update_data = AsyncMock( + return_value={"data": mock_updated_key_obj} + ) + + # Mock prepare_key_update_data + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data" + ) as mock_prepare: + mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]} + + # Mock TeamMemberPermissionChecks + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint" + ) as mock_permission_check: + mock_permission_check.return_value = None + + # Mock _delete_cache_key_object + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache: + mock_delete_cache.return_value = None + + # Mock hash_token (imported from litellm.proxy._types) + with patch( + "litellm.proxy._types.hash_token" + ) as mock_hash: + mock_hash.return_value = "hashed-test-key-123" + + # Mock KeyManagementEventHooks + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" + ): + # Create update request + key_update_item = BulkUpdateKeyRequestItem( + key="test-key-123", + max_budget=100.0, + tags=["production"], + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + # Call the function + result = await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + llm_router=mock_llm_router, + ) + + # Verify results + assert result is not None + assert "token" not in result # Token should be removed + assert result.get("max_budget") == 100.0 + assert result.get("tags") == ["production"] + + # Verify mocks were called + mock_prisma_client.get_data.assert_called_once() + mock_prisma_client.update_data.assert_called_once() + mock_delete_cache.assert_called_once() + + +@pytest.mark.asyncio +async def test_bulk_update_keys_success(monkeypatch): + """ + Test /key/bulk_update endpoint with successful updates. + + Tests: + 1. Multiple keys updated successfully + 2. Response contains correct counts and data + """ + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + ) + from litellm.proxy.management_endpoints.key_management_endpoints import ( + bulk_update_keys, + ) + from litellm.proxy.proxy_server import ( + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + # Setup mocks + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_llm_router = MagicMock() + + # Mock existing keys + existing_key_1 = LiteLLM_VerificationToken( + token="test-key-1", + user_id="user-123", + models=["gpt-4"], + team_id=None, + max_budget=None, + ) + existing_key_2 = LiteLLM_VerificationToken( + token="test-key-2", + user_id="user-123", + models=["gpt-3.5-turbo"], + team_id=None, + max_budget=50.0, + ) + + # Mock updated key responses + updated_key_1_data = { + "user_id": "user-123", + "models": ["gpt-4"], + "max_budget": 100.0, + "tags": ["production"], + } + updated_key_2_data = { + "user_id": "user-123", + "models": ["gpt-3.5-turbo"], + "max_budget": 200.0, + "tags": ["staging"], + } + + mock_prisma_client.get_data = AsyncMock( + side_effect=[existing_key_1, existing_key_2] + ) + mock_updated_key_1_obj = MagicMock() + mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data + mock_updated_key_2_obj = MagicMock() + mock_updated_key_2_obj.model_dump.return_value = updated_key_2_data + mock_prisma_client.update_data = AsyncMock( + side_effect=[ + {"data": mock_updated_key_1_obj}, + {"data": mock_updated_key_2_obj}, + ] + ) + + # Patch dependencies + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router) + + # Mock helper functions + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data" + ) as mock_prepare: + mock_prepare.side_effect = [ + {"max_budget": 100.0, "tags": ["production"]}, + {"max_budget": 200.0, "tags": ["staging"]}, + ] + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint" + ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ): + with patch( + "litellm.proxy._types.hash_token" + ) as mock_hash: + mock_hash.side_effect = ["hashed-key-1", "hashed-key-2"] + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" + ): + # Create request + request_data = BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem( + key="test-key-1", + max_budget=100.0, + tags=["production"], + ), + BulkUpdateKeyRequestItem( + key="test-key-2", + max_budget=200.0, + tags=["staging"], + ), + ] + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + # Call endpoint + response = await bulk_update_keys( + data=request_data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # Verify response + assert response.total_requested == 2 + assert len(response.successful_updates) == 2 + assert len(response.failed_updates) == 0 + assert response.successful_updates[0].key == "test-key-1" + assert response.successful_updates[1].key == "test-key-2" + + +@pytest.mark.asyncio +async def test_bulk_update_keys_partial_failures(monkeypatch): + """ + Test /key/bulk_update endpoint with partial failures. + + Tests: + 1. Some keys update successfully, others fail + 2. Response contains both successful and failed updates + 3. Failed updates include error messages + """ + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + ) + from litellm.proxy.management_endpoints.key_management_endpoints import ( + bulk_update_keys, + ) + + # Setup mocks + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_llm_router = MagicMock() + + # Mock existing keys + existing_key_1 = LiteLLM_VerificationToken( + token="test-key-1", + user_id="user-123", + models=["gpt-4"], + team_id=None, + max_budget=None, + ) + + # Mock updated key response for successful update + updated_key_1_data = { + "user_id": "user-123", + "models": ["gpt-4"], + "max_budget": 100.0, + "tags": ["production"], + } + + # First key exists, second key doesn't exist + mock_prisma_client.get_data = AsyncMock( + side_effect=[existing_key_1, None] # Second key not found + ) + mock_updated_key_1_obj = MagicMock() + mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data + mock_prisma_client.update_data = AsyncMock( + return_value={"data": mock_updated_key_1_obj} + ) + + # Patch dependencies + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router) + + # Mock helper functions + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data" + ) as mock_prepare: + mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]} + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint" + ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ): + with patch( + "litellm.proxy._types.hash_token" + ) as mock_hash: + mock_hash.return_value = "hashed-key-1" + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" + ): + # Create request with one valid and one invalid key + request_data = BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem( + key="test-key-1", + max_budget=100.0, + tags=["production"], + ), + BulkUpdateKeyRequestItem( + key="non-existent-key", + max_budget=200.0, + tags=["staging"], + ), + ] + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + # Call endpoint + response = await bulk_update_keys( + data=request_data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # Verify response + assert response.total_requested == 2 + assert len(response.successful_updates) == 1 + assert len(response.failed_updates) == 1 + assert response.successful_updates[0].key == "test-key-1" + assert response.failed_updates[0].key == "non-existent-key" + assert "Key not found" in response.failed_updates[0].failed_reason From 93d6aae4a36049010f668f4b91f0bcaf9e13d3f5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Jan 2026 14:06:54 -0800 Subject: [PATCH 02/45] mypy linting --- .../management_endpoints/key_management_endpoints.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f9fb0e5f49..380e8bddc9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1568,7 +1568,7 @@ async def _process_single_key_update( check_db_only=True, ) - if team_obj is not None: + if team_obj is not None and prisma_client is not None: await _check_team_key_limits( team_table=team_obj, data=update_key_request, @@ -1606,6 +1606,12 @@ async def _process_single_key_update( ) # Update key in database + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + _data = {**non_default_values, "token": key_update_item.key} response = await prisma_client.update_data( token=key_update_item.key, data=_data From 8c4ccdc313c9af5405dc24f99b562bcb2d38dbff Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 28 Jan 2026 02:34:40 +0100 Subject: [PATCH 03/45] test(proxy): add regression tests for vertex passthrough model names with slashes (#19855) Added test cases for custom model names containing slashes in Vertex AI passthrough URLs (e.g., gcp/google/gemini-2.5-flash). Test cases: - gcp/google/gemini-2.5-flash - gcp/google/gemini-3-flash-preview - custom/model --- tests/local_testing/test_auth_utils.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index 72f799a6cf..d36f96b1a3 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -356,6 +356,25 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none(): "/openai/deployments/my-deployment/chat/completions", "my-deployment" ), + # Custom model_name with slashes (e.g., gcp/google/gemini-2.5-flash) + # This is the NVIDIA P0 bug fix - regex should capture full model name including slashes + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gcp/google/gemini-2.5-flash:generateContent", + "gcp/google/gemini-2.5-flash" + ), + # Another custom model_name with slashes + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/global/publishers/google/models/gcp/google/gemini-3-flash-preview:generateContent", + "gcp/google/gemini-3-flash-preview" + ), + # Model name with single slash + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/custom/model:generateContent", + "custom/model" + ), ], ) def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model): From d0939075bc84cc27fd6e37c7df72f4eb9af439e8 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 28 Jan 2026 07:06:18 +0530 Subject: [PATCH 04/45] fix: guardrails issues streaming-response regex (#19901) --- litellm/proxy/common_request_processing.py | 12 +- .../litellm_content_filter/content_filter.py | 170 +++++++++--------- .../litellm_content_filter/patterns.json | 12 +- litellm/proxy/utils.py | 13 +- litellm/types/guardrails.py | 25 ++- .../content_filter/test_content_filter.py | 84 +++++---- 6 files changed, 180 insertions(+), 136 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0d3e61b75c..51f3e6482a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -650,11 +650,15 @@ class ProxyBaseLLMRequestProcessing: ) tasks = [] + # Start the moderation check (during_call_hook) as early as possible + # This gives it a head start to mask/validate input while the proxy handles routing tasks.append( - proxy_logging_obj.during_call_hook( - data=self.data, - user_api_key_dict=user_api_key_dict, - call_type=route_type, # type: ignore + asyncio.create_task( + proxy_logging_obj.during_call_hook( + data=self.data, + user_api_key_dict=user_api_key_dict, + call_type=route_type, # type: ignore + ) ) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index c9bd0135a0..083a407e9c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -198,6 +198,15 @@ class ContentFilterGuardrail(CustomGuardrail): for pattern_config in normalized_patterns: self._add_pattern(pattern_config) + # Warn if using during_call with MASK action (unstable) + if self.event_hook == GuardrailEventHooks.during_call and any( + p["action"] == ContentFilterAction.MASK for p in self.compiled_patterns + ): + verbose_proxy_logger.warning( + f"ContentFilterGuardrail '{self.guardrail_name}': 'during_call' mode with 'MASK' action is unstable due to race conditions. " + "Use 'pre_call' mode for reliable request masking." + ) + # Load blocked words - always initialize as dict self.blocked_words: Dict[str, Tuple[ContentFilterAction, Optional[str]]] = {} for word in normalized_blocked_words: @@ -905,11 +914,15 @@ class ContentFilterGuardrail(CustomGuardrail): elif isinstance(e.detail, str): e.detail = e.detail + " (Image description): " + description else: - e.detail = "Content blocked: Image description detected" + description + e.detail = ( + "Content blocked: Image description detected" + description + ) raise e def _count_masked_entities( - self, detections: List[ContentFilterDetection], masked_entity_count: Dict[str, int] + self, + detections: List[ContentFilterDetection], + masked_entity_count: Dict[str, int], ) -> None: """ Count masked entities by type from detections. @@ -964,9 +977,11 @@ class ContentFilterGuardrail(CustomGuardrail): dict(detection) for detection in detections ] if status != "success": - guardrail_json_response = exception_str if exception_str else [ - dict(detection) for detection in detections - ] + guardrail_json_response = ( + exception_str + if exception_str + else [dict(detection) for detection in detections] + ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, @@ -1066,99 +1081,84 @@ class ContentFilterGuardrail(CustomGuardrail): Process streaming response chunks and check for blocked content. For BLOCK action: Raises HTTPException immediately when blocked content is detected. - For MASK action: Content passes through (masking streaming responses is not supported). + For MASK action: Content is buffered to handle patterns split across chunks. """ + accumulated_full_text = "" + yielded_masked_text_len = 0 + buffer_size = 50 # Increased buffer to catch patterns split across many chunks - # Accumulate content as we iterate through chunks - accumulated_content = "" + verbose_proxy_logger.info( + f"ContentFilterGuardrail: Starting robust streaming masking for model {request_data.get('model')}" + ) async for item in response: - # Accumulate content from this chunk before checking if isinstance(item, ModelResponseStream) and item.choices: + delta_content = "" + is_final = False for choice in item.choices: if hasattr(choice, "delta") and choice.delta: content = getattr(choice.delta, "content", None) if content and isinstance(content, str): - accumulated_content += content + delta_content += content + if getattr(choice, "finish_reason", None): + is_final = True - # Check accumulated content for blocked patterns/keywords after processing all choices - # Only check for BLOCK actions, not MASK (masking streaming is not supported) - if accumulated_content: - try: - # Check patterns - pattern_match = self._check_patterns(accumulated_content) - if pattern_match: - matched_text, pattern_name, action = pattern_match - if action == ContentFilterAction.BLOCK: - error_msg = ( - f"Content blocked: {pattern_name} pattern detected" - ) - verbose_proxy_logger.warning(error_msg) - raise HTTPException( - status_code=403, - detail={ - "error": error_msg, - "pattern": pattern_name, - }, - ) + accumulated_full_text += delta_content - # Check blocked words - blocked_word_match = self._check_blocked_words( - accumulated_content - ) - if blocked_word_match: - keyword, action, description = blocked_word_match - if action == ContentFilterAction.BLOCK: - error_msg = ( - f"Content blocked: keyword '{keyword}' detected" - ) - if description: - error_msg += f" ({description})" - verbose_proxy_logger.warning(error_msg) - raise HTTPException( - status_code=403, - detail={ - "error": error_msg, - "keyword": keyword, - "description": description, - }, - ) + # Check for blocking or apply masking + # Add a space at the end if it's the final chunk to trigger word boundaries (\b) + text_to_check = accumulated_full_text + if is_final: + text_to_check += " " - # Check category keywords - all_exceptions = [] - for category in self.loaded_categories.values(): - all_exceptions.extend(category.exceptions) - category_match = self._check_category_keywords( - accumulated_content, all_exceptions - ) - if category_match: - keyword, category_name, severity, action = category_match - if action == ContentFilterAction.BLOCK: - error_msg = ( - f"Content blocked: {category_name} category keyword '{keyword}' detected " - f"(severity: {severity})" - ) - verbose_proxy_logger.warning(error_msg) - raise HTTPException( - status_code=403, - detail={ - "error": error_msg, - "category": category_name, - "keyword": keyword, - "severity": severity, - }, - ) - except HTTPException: - # Re-raise HTTPException (blocked content detected) - raise - except Exception as e: - # Log other exceptions but don't block the stream - verbose_proxy_logger.warning( - f"Error checking content filter in streaming: {e}" - ) + try: + masked_text = self._filter_single_text(text_to_check) + if is_final and masked_text.endswith(" "): + masked_text = masked_text[:-1] + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.error( + f"ContentFilterGuardrail: Error in masking: {e}" + ) + masked_text = text_to_check # Fallback to current text - # Yield the chunk (only if no exception was raised above) - yield item + # Determine how much can be safely yielded + if is_final: + safe_to_yield_len = len(masked_text) + else: + safe_to_yield_len = max(0, len(masked_text) - buffer_size) + + if safe_to_yield_len > yielded_masked_text_len: + new_masked_content = masked_text[ + yielded_masked_text_len:safe_to_yield_len + ] + # Modify the chunk to contain only the new masked content + if ( + item.choices + and hasattr(item.choices[0], "delta") + and item.choices[0].delta + ): + item.choices[0].delta.content = new_masked_content + yielded_masked_text_len = safe_to_yield_len + yield item + else: + # Hold content by yielding empty content chunk (keeps metadata/structure) + if ( + item.choices + and hasattr(item.choices[0], "delta") + and item.choices[0].delta + ): + item.choices[0].delta.content = "" + yield item + else: + # Not a ModelResponseStream or no choices - yield as is + yield item + + # Any remaining content (should have been handled by is_final, but just in case) + if yielded_masked_text_len < len(accumulated_full_text): + # We already reached the end of the generator + pass @staticmethod def get_config_model(): diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json index f2427b5b92..1eff7804b4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json @@ -108,7 +108,7 @@ { "name": "ipv6", "display_name": "IP Address (IPv6)", - "pattern": "\\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}\\b", + "pattern": "(? 0 - assert guardrail_info["start_time"] < guardrail_info["end_time"] + assert guardrail_info["duration"] >= 0 + assert guardrail_info["start_time"] <= guardrail_info["end_time"] # Verify detections are logged assert "guardrail_response" in guardrail_info @@ -839,15 +839,21 @@ class TestContentFilterGuardrail: assert "action" in detection assert detection["action"] == "MASK" # Verify sensitive content (matched_text) is NOT included - assert "matched_text" not in detection, "Sensitive content should not be logged" + assert ( + "matched_text" not in detection + ), "Sensitive content should not be logged" # Verify blocked word detection structure - blocked_word_detections = [d for d in detections if d.get("type") == "blocked_word"] + blocked_word_detections = [ + d for d in detections if d.get("type") == "blocked_word" + ] assert len(blocked_word_detections) > 0 for detection in blocked_word_detections: assert detection["type"] == "blocked_word" assert "keyword" in detection - assert detection["keyword"] == "confidential" # Config keyword, not user content + assert ( + detection["keyword"] == "confidential" + ) # Config keyword, not user content assert "action" in detection assert detection["action"] == "MASK" assert "description" in detection @@ -896,7 +902,9 @@ class TestContentFilterGuardrail: assert "metadata" in request_data assert "standard_logging_guardrail_information" in request_data["metadata"] - guardrail_info_list = request_data["metadata"]["standard_logging_guardrail_information"] + guardrail_info_list = request_data["metadata"][ + "standard_logging_guardrail_information" + ] assert len(guardrail_info_list) == 1 guardrail_info = guardrail_info_list[0] @@ -909,4 +917,6 @@ class TestContentFilterGuardrail: # If detections are logged, verify they don't contain sensitive content for detection in detections: if detection.get("type") == "pattern": - assert "matched_text" not in detection, "Sensitive content should not be logged" + assert ( + "matched_text" not in detection + ), "Sensitive content should not be logged" From 54a83e75cc0aaabdfa3a2448bc71233cd048b326 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 28 Jan 2026 07:08:02 +0530 Subject: [PATCH 05/45] fix: add fix for migration issue and and stable linux debain (#19843) --- .../migration.sql | 10 +++++----- litellm/proxy/proxy_server.py | 15 ++++++++++++++- litellm/proxy/schema.prisma | 1 + 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql index a9d9528bd2..43eb240142 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql @@ -1,12 +1,12 @@ -- DropIndex -DROP INDEX "LiteLLM_PromptTable_prompt_id_key"; +DROP INDEX IF EXISTS "LiteLLM_PromptTable_prompt_id_key"; -- AlterTable -ALTER TABLE "LiteLLM_PromptTable" ADD COLUMN "version" INTEGER NOT NULL DEFAULT 1; +ALTER TABLE "LiteLLM_PromptTable" +ADD COLUMN "version" INTEGER NOT NULL DEFAULT 1; -- CreateIndex -CREATE INDEX "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable"("prompt_id"); +CREATE INDEX "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable" ("prompt_id"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable"("prompt_id", "version"); - +CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable" ("prompt_id", "version"); \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 183c25ed46..4c8f31e930 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5098,7 +5098,20 @@ class ProxyStartupEvent: except Exception as e: raise e - await prisma_client.connect() + try: + await prisma_client.connect() + except Exception as e: + if "P3018" in str(e) or "P3009" in str(e): + verbose_proxy_logger.debug( + "CRITICAL: DATABASE MIGRATION FAILED" + ) + verbose_proxy_logger.debug( + "Your database is in a 'dirty' state." + ) + verbose_proxy_logger.debug( + "FIX: Run 'prisma migrate resolve --applied '" + ) + raise e ## Start RDS IAM token refresh background task if enabled ## # This proactively refreshes IAM tokens before they expire, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d7aa6e9f0d..d46c2db763 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -5,6 +5,7 @@ datasource client { generator client { provider = "prisma-client-py" + binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"] } // Budget / Rate Limits for an org From 4717f742eb8216c4b56edb60e816c96f9e2dc028 Mon Sep 17 00:00:00 2001 From: Jay Prajapati <79649559+jayy-77@users.noreply.github.com> Date: Wed, 28 Jan 2026 07:17:27 +0530 Subject: [PATCH 06/45] fix: filter unsupported beta headers for Bedrock Invoke API (#19877) - Add whitelist-based filtering for anthropic_beta headers - Only allow Bedrock-supported beta flags (computer-use, tool-search, etc.) - Filter out unsupported flags like mcp-servers, structured-outputs - Remove output_format parameter from Bedrock Invoke requests - Force tool-based structured outputs when response_format is used Fixes #16726 --- .../anthropic_claude3_transformation.py | 37 +++++- ...ations_anthropic_claude3_transformation.py | 105 ++++++++++++++++++ 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 53e0822979..c936b2cd23 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -53,13 +53,26 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model: str, drop_params: bool, ) -> dict: - return AnthropicConfig.map_openai_params( + # Force tool-based structured outputs for Bedrock Invoke + # (similar to VertexAI fix in #19201) + # Bedrock Invoke doesn't support output_format parameter + original_model = model + if "response_format" in non_default_params: + # Use a model name that forces tool-based approach + model = "claude-3-sonnet-20240229" + + optional_params = AnthropicConfig.map_openai_params( self, non_default_params, optional_params, model, drop_params, ) + + # Restore original model name + model = original_model + + return optional_params def transform_request( @@ -90,6 +103,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): _anthropic_request.pop("model", None) _anthropic_request.pop("stream", None) + # Bedrock Invoke doesn't support output_format parameter + _anthropic_request.pop("output_format", None) if "anthropic_version" not in _anthropic_request: _anthropic_request["anthropic_version"] = self.anthropic_version @@ -117,6 +132,26 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if "opus-4" in model.lower() or "opus_4" in model.lower(): beta_set.add("tool-search-tool-2025-10-19") + # Filter out beta headers that Bedrock Invoke doesn't support + # AWS Bedrock only supports a specific whitelist of beta flags + # Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + BEDROCK_SUPPORTED_BETAS = { + "computer-use-2024-10-22", # Legacy computer use + "computer-use-2025-01-24", # Current computer use (Claude 3.7 Sonnet) + "token-efficient-tools-2025-02-19", # Tool use (Claude 3.7+ and Claude 4+) + "interleaved-thinking-2025-05-14", # Interleaved thinking (Claude 4+) + "output-128k-2025-02-19", # 128K output tokens (Claude 3.7 Sonnet) + "dev-full-thinking-2025-05-14", # Developer mode for raw thinking (Claude 4+) + "context-1m-2025-08-07", # 1 million tokens (Claude Sonnet 4) + "context-management-2025-06-27", # Context management (Claude Sonnet/Haiku 4.5) + "effort-2025-11-24", # Effort parameter (Claude Opus 4.5) + "tool-search-tool-2025-10-19", # Tool search (Claude Opus 4.5) + "tool-examples-2025-10-29", # Tool use examples (Claude Opus 4.5) + } + + # Only keep beta headers that Bedrock supports + beta_set = {beta for beta in beta_set if beta in BEDROCK_SUPPORTED_BETAS} + if beta_set: _anthropic_request["anthropic_beta"] = list(beta_set) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 98b392a353..5c1b4cbd38 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -464,3 +464,108 @@ def test_opus_4_5_model_detection(): for model in non_opus_4_5_models: assert not config._is_claude_opus_4_5(model), \ f"Should not detect {model} as Opus 4.5" + + +def test_structured_outputs_beta_header_filtered_for_bedrock_invoke(): + """ + Test that unsupported beta headers are filtered out for Bedrock Invoke API. + + Bedrock Invoke API only supports a specific whitelist of beta flags and returns + "invalid beta flag" error for others (e.g., structured-outputs, mcp-servers). + This test ensures unsupported headers are filtered while keeping supported ones. + + Fixes: https://github.com/BerriAI/litellm/issues/16726 + """ + config = AmazonAnthropicClaudeConfig() + + messages = [{"role": "user", "content": "test"}] + + # Test 1: structured-outputs beta header (unsupported) + headers = {"anthropic-beta": "structured-outputs-2025-11-13"} + + result = config.transform_request( + model="anthropic.claude-4-0-sonnet-20250514-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers=headers, + ) + + # Verify structured-outputs beta is filtered out + anthropic_beta = result.get("anthropic_beta", []) + assert not any("structured-outputs" in beta for beta in anthropic_beta), \ + f"structured-outputs beta should be filtered, got: {anthropic_beta}" + + # Test 2: mcp-servers beta header (unsupported - the main issue from #16726) + headers = {"anthropic-beta": "mcp-servers-2025-12-04"} + + result = config.transform_request( + model="anthropic.claude-4-0-sonnet-20250514-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers=headers, + ) + + # Verify mcp-servers beta is filtered out + anthropic_beta = result.get("anthropic_beta", []) + assert not any("mcp-servers" in beta for beta in anthropic_beta), \ + f"mcp-servers beta should be filtered, got: {anthropic_beta}" + + # Test 3: Mix of supported and unsupported beta headers + headers = {"anthropic-beta": "computer-use-2024-10-22,mcp-servers-2025-12-04,structured-outputs-2025-11-13"} + + result = config.transform_request( + model="anthropic.claude-4-0-sonnet-20250514-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers=headers, + ) + + # Verify only supported betas are kept + anthropic_beta = result.get("anthropic_beta", []) + assert not any("structured-outputs" in beta for beta in anthropic_beta), \ + f"structured-outputs beta should be filtered, got: {anthropic_beta}" + assert not any("mcp-servers" in beta for beta in anthropic_beta), \ + f"mcp-servers beta should be filtered, got: {anthropic_beta}" + assert any("computer-use" in beta for beta in anthropic_beta), \ + f"computer-use beta should be kept, got: {anthropic_beta}" + + +def test_output_format_removed_from_bedrock_invoke_request(): + """ + Test that output_format parameter is removed from Bedrock Invoke requests. + + Bedrock Invoke API doesn't support the output_format parameter (only supported + in Anthropic Messages API). This test ensures it's removed to prevent errors. + """ + config = AmazonAnthropicClaudeConfig() + + messages = [{"role": "user", "content": "test"}] + + # Create a request with output_format via map_openai_params + non_default_params = { + "response_format": {"type": "json_object"} + } + optional_params = {} + + # This should trigger tool-based structured outputs + optional_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="anthropic.claude-4-0-sonnet-20250514-v1:0", + drop_params=False, + ) + + result = config.transform_request( + model="anthropic.claude-4-0-sonnet-20250514-v1:0", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + # Verify output_format is not in the request + assert "output_format" not in result, \ + f"output_format should be removed for Bedrock Invoke, got keys: {result.keys()}" From 6a9d41234f05fd1617aed605fdc5d33be4cfcf8e Mon Sep 17 00:00:00 2001 From: Jay Prajapati <79649559+jayy-77@users.noreply.github.com> Date: Wed, 28 Jan 2026 07:21:13 +0530 Subject: [PATCH 07/45] fix: allow tool_choice for Azure GPT-5 chat models (#19813) * fix: don't treat gpt-5-chat as GPT-5 reasoning * fix: mark azure gpt-5-chat as supporting tool_choice * test: cover gpt-5-chat params on azure/openai --- litellm/llms/azure/chat/gpt_5_transformation.py | 8 +++++++- .../llms/openai/chat/gpt_5_transformation.py | 4 +++- .../model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- .../chat/test_azure_gpt5_transformation.py | 11 +++++++++++ .../llms/openai/test_gpt5_transformation.py | 17 +++++++++++++++++ 6 files changed, 42 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 506b7fdfe5..eeb55911ec 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -22,7 +22,8 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix used for manual routing. """ - return "gpt-5" in model or "gpt5_series" in model + # gpt-5-chat* is a chat model and shouldn't go through GPT-5 reasoning restrictions. + return ("gpt-5" in model and "gpt-5-chat" not in model) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> List[str]: """Get supported parameters for Azure OpenAI GPT-5 models. @@ -37,6 +38,11 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): """ params = OpenAIGPT5Config.get_supported_openai_params(self, model=model) + # Azure supports tool_choice for GPT-5 deployments, but the base GPT-5 config + # can drop it when the deployment name isn't in the OpenAI model registry. + if "tool_choice" not in params: + params.append("tool_choice") + # Only gpt-5.2 has been verified to support logprobs on Azure if self.is_model_gpt_5_2_model(model): azure_supported_params = ["logprobs", "top_logprobs"] diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 3fffa335fd..05c003c8b7 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -19,7 +19,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - return "gpt-5" in model + # gpt-5-chat* behaves like a regular chat model (supports temperature, etc.) + # Don't route it through GPT-5 reasoning-specific parameter restrictions. + return "gpt-5" in model and "gpt-5-chat" not in model @classmethod def is_model_gpt_5_codex_model(cls, model: str) -> bool: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4ac4159558..f5c680990c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3130,7 +3130,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true }, "azure/gpt-5-chat-latest": { @@ -3162,7 +3162,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true }, "azure/gpt-5-codex": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4ac4159558..f5c680990c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3130,7 +3130,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true }, "azure/gpt-5-chat-latest": { @@ -3162,7 +3162,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true }, "azure/gpt-5-codex": { diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 199a16d859..25f3d1364f 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -16,6 +16,17 @@ def test_azure_gpt5_supports_reasoning_effort(config: AzureOpenAIGPT5Config): ) +def test_azure_gpt5_allows_tool_choice_for_deployment_names(): + supported_params = litellm.get_supported_openai_params( + model="gpt-5-chat-2025-08-07", custom_llm_provider="azure" + ) + assert supported_params is not None + assert "tool_choice" in supported_params + # gpt-5-chat* should not be treated as a GPT-5 reasoning model + assert "reasoning_effort" not in supported_params + assert "temperature" in supported_params + + def test_azure_gpt5_maps_max_tokens(config: AzureOpenAIGPT5Config): params = config.map_openai_params( non_default_params={"max_tokens": 5}, diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index fd25d302d0..386f264a4d 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -20,6 +20,23 @@ def test_gpt5_supports_reasoning_effort(config: OpenAIConfig): assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5-mini") +def test_gpt5_chat_does_not_support_reasoning_effort(config: OpenAIConfig): + assert ( + "reasoning_effort" + not in config.get_supported_openai_params(model="gpt-5-chat-latest") + ) + + +def test_gpt5_chat_supports_temperature(config: OpenAIConfig): + params = config.map_openai_params( + non_default_params={"temperature": 0.3}, + optional_params={}, + model="gpt-5-chat-latest", + drop_params=False, + ) + assert params["temperature"] == 0.3 + + def test_gpt5_maps_max_tokens(config: OpenAIConfig): params = config.map_openai_params( non_default_params={"max_tokens": 10}, From d6cf4df3cbc53da3ea4b906681c56088ee74e0a5 Mon Sep 17 00:00:00 2001 From: Teo Stocco Date: Tue, 27 Jan 2026 18:02:37 -0800 Subject: [PATCH 08/45] fix: tool with antropic #19800 (#19805) --- litellm/llms/anthropic/chat/transformation.py | 17 ++++-- .../test_anthropic_chat_transformation.py | 53 +++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 82eccee596..f0eaf12fb0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -290,10 +290,19 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): elif tool_choice == "none": _tool_choice = AnthropicMessagesToolChoice(type="none") elif isinstance(tool_choice, dict): - _tool_name = tool_choice.get("function", {}).get("name") - _tool_choice = AnthropicMessagesToolChoice(type="tool") - if _tool_name is not None: - _tool_choice["name"] = _tool_name + if "type" in tool_choice and "function" not in tool_choice: + tool_type = tool_choice.get("type") + if tool_type == "auto": + _tool_choice = AnthropicMessagesToolChoice(type="auto") + elif tool_type == "required" or tool_type == "any": + _tool_choice = AnthropicMessagesToolChoice(type="any") + elif tool_type == "none": + _tool_choice = AnthropicMessagesToolChoice(type="none") + else: + _tool_name = tool_choice.get("function", {}).get("name") + if _tool_name is not None: + _tool_choice = AnthropicMessagesToolChoice(type="tool") + _tool_choice["name"] = _tool_name if parallel_tool_use is not None: # Anthropic uses 'disable_parallel_tool_use' flag to determine if parallel tool use is allowed diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 7e96c4634f..bd3fa93e6a 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -548,6 +548,59 @@ def test_map_tool_choice_dict_type_function_with_name(): assert result["name"] == "my_tool" +def test_map_tool_choice_dict_type_auto(): + """ + Test that dict {"type": "auto"} maps to Anthropic type='auto'. + This handles Cursor's format for tool_choice. + """ + config = AnthropicConfig() + result = config._map_tool_choice( + tool_choice={"type": "auto"}, + parallel_tool_use=None, + ) + assert result is not None + assert result["type"] == "auto" + + +def test_map_tool_choice_dict_type_required(): + """ + Test that dict {"type": "required"} maps to Anthropic type='any'. + """ + config = AnthropicConfig() + result = config._map_tool_choice( + tool_choice={"type": "required"}, + parallel_tool_use=None, + ) + assert result is not None + assert result["type"] == "any" + + +def test_map_tool_choice_dict_type_none(): + """ + Test that dict {"type": "none"} maps to Anthropic type='none'. + """ + config = AnthropicConfig() + result = config._map_tool_choice( + tool_choice={"type": "none"}, + parallel_tool_use=None, + ) + assert result is not None + assert result["type"] == "none" + + +def test_map_tool_choice_dict_type_function_without_name(): + """ + Test that dict {"type": "function"} without name is handled gracefully. + Should return None since there's no valid tool name. + """ + config = AnthropicConfig() + result = config._map_tool_choice( + tool_choice={"type": "function"}, + parallel_tool_use=None, + ) + assert result is None + + def test_transform_response_with_prefix_prompt(): import httpx From 920ef665a3f1e1cd306e9a01aeda8e060df24d2b Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Tue, 27 Jan 2026 21:15:04 -0500 Subject: [PATCH 09/45] inspect BadRequestError after all other policy types (#19878) As indicated by https://docs.litellm.ai/docs/exception_mapping, BadRequestError is used as the base type for multiple exceptions. As such, it should be tested last in handling retry policies. This updates the integration test that validates retry policies work as expected. Fixes #19876 --- litellm/router.py | 10 +++++----- litellm/router_utils/get_retry_from_policy.py | 10 +++++----- tests/local_testing/test_completion_with_retries.py | 1 + 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 09d71b6b49..a3c3afa932 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8729,11 +8729,6 @@ class Router: if allowed_fails_policy is None: return None - if ( - isinstance(exception, litellm.BadRequestError) - and allowed_fails_policy.BadRequestErrorAllowedFails is not None - ): - return allowed_fails_policy.BadRequestErrorAllowedFails if ( isinstance(exception, litellm.AuthenticationError) and allowed_fails_policy.AuthenticationErrorAllowedFails is not None @@ -8754,6 +8749,11 @@ class Router: and allowed_fails_policy.ContentPolicyViolationErrorAllowedFails is not None ): return allowed_fails_policy.ContentPolicyViolationErrorAllowedFails + if ( + isinstance(exception, litellm.BadRequestError) + and allowed_fails_policy.BadRequestErrorAllowedFails is not None + ): + return allowed_fails_policy.BadRequestErrorAllowedFails def _initialize_alerting(self): from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index 48df43ef81..ec326ebb50 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -43,11 +43,6 @@ def get_num_retries_from_retry_policy( if isinstance(retry_policy, dict): retry_policy = RetryPolicy(**retry_policy) - if ( - isinstance(exception, BadRequestError) - and retry_policy.BadRequestErrorRetries is not None - ): - return retry_policy.BadRequestErrorRetries if ( isinstance(exception, AuthenticationError) and retry_policy.AuthenticationErrorRetries is not None @@ -65,6 +60,11 @@ def get_num_retries_from_retry_policy( and retry_policy.ContentPolicyViolationErrorRetries is not None ): return retry_policy.ContentPolicyViolationErrorRetries + if ( + isinstance(exception, BadRequestError) + and retry_policy.BadRequestErrorRetries is not None + ): + return retry_policy.BadRequestErrorRetries def reset_retry_policy() -> RetryPolicy: diff --git a/tests/local_testing/test_completion_with_retries.py b/tests/local_testing/test_completion_with_retries.py index 6eb3ad460e..585e1ee261 100644 --- a/tests/local_testing/test_completion_with_retries.py +++ b/tests/local_testing/test_completion_with_retries.py @@ -60,6 +60,7 @@ async def test_completion_with_retry_policy(sync_mode): retry_number = 1 retry_policy = RetryPolicy( + BadRequestErrorRetries=10, ContentPolicyViolationErrorRetries=retry_number, # run 3 retries for ContentPolicyViolationErrors AuthenticationErrorRetries=0, # run 0 retries for AuthenticationErrorRetries ) From 807ba011ebd12b1981012f419345161f65ecd135 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 27 Jan 2026 23:16:58 -0300 Subject: [PATCH 10/45] fix(main): use local tiktoken cache in lazy loading (#19774) The lazy loading implementation for encoding in __getattr__ was calling tiktoken.get_encoding() directly without first setting TIKTOKEN_CACHE_DIR. This caused tiktoken to attempt downloading the encoding file from the internet instead of using the local copy bundled with litellm. This fix uses _get_default_encoding() from _lazy_imports which properly sets TIKTOKEN_CACHE_DIR before loading tiktoken, ensuring the local cache is used. --- litellm/main.py | 7 +++-- .../test_litellm/test_eager_tiktoken_load.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 5b8c569a39..319a59771f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7300,8 +7300,11 @@ def _get_encoding(): def __getattr__(name: str) -> Any: """Lazy import handler for main module""" if name == "encoding": - # Lazy load encoding to avoid heavy tiktoken import at module load time - _encoding = tiktoken.get_encoding("cl100k_base") + # Use _get_default_encoding which properly sets TIKTOKEN_CACHE_DIR + # before loading tiktoken, ensuring the local cache is used + # instead of downloading from the internet + from litellm._lazy_imports import _get_default_encoding + _encoding = _get_default_encoding() # Cache it in the module's __dict__ for subsequent accesses import sys diff --git a/tests/test_litellm/test_eager_tiktoken_load.py b/tests/test_litellm/test_eager_tiktoken_load.py index 1264c68b99..33dd57fad8 100644 --- a/tests/test_litellm/test_eager_tiktoken_load.py +++ b/tests/test_litellm/test_eager_tiktoken_load.py @@ -78,6 +78,35 @@ def test_lazy_loading_default(): assert len(tokens) > 0, "Encoding should work" +def test_tiktoken_cache_dir_set_on_lazy_load(): + """Test that TIKTOKEN_CACHE_DIR is set when encoding is lazy loaded. + + This ensures the local tiktoken cache is used instead of downloading + from the internet. Regression test for issue #19768. + """ + # Remove environment variables to ensure clean state + if "LITELLM_DISABLE_LAZY_LOADING" in os.environ: + del os.environ["LITELLM_DISABLE_LAZY_LOADING"] + if "TIKTOKEN_CACHE_DIR" in os.environ: + del os.environ["TIKTOKEN_CACHE_DIR"] + + # Clear any cached modules + modules_to_clear = [k for k in sys.modules.keys() if k.startswith("litellm")] + for module in modules_to_clear: + del sys.modules[module] + + # Import litellm fresh + import litellm + + # Access encoding (triggers lazy load) + _ = litellm.encoding + + # Verify TIKTOKEN_CACHE_DIR is now set and points to local tokenizers + assert "TIKTOKEN_CACHE_DIR" in os.environ, "TIKTOKEN_CACHE_DIR should be set after lazy loading encoding" + cache_dir = os.environ["TIKTOKEN_CACHE_DIR"] + assert "tokenizers" in cache_dir, f"TIKTOKEN_CACHE_DIR should point to tokenizers directory, got: {cache_dir}" + + @pytest.fixture(autouse=True) def cleanup_env(): """Clean up environment variable after each test""" From 64c102e3c2f1f52dd77767f6bbd8954fb3e56e2a Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 27 Jan 2026 23:18:47 -0300 Subject: [PATCH 11/45] fix(gemini): subtract implicit cached tokens from text_tokens for correct cost calculation (#19775) When Gemini uses implicit caching, it returns cachedContentTokenCount but NOT cacheTokensDetails. Previously, text_tokens was not adjusted in this case, causing costs to be calculated as if all tokens were non-cached. This fix subtracts cachedContentTokenCount from text_tokens when no cacheTokensDetails is present (implicit caching), ensuring correct cost calculation with the reduced cache_read pricing. --- .../vertex_and_google_ai_studio_gemini.py | 10 +++ tests/test_litellm/test_cost_calculator.py | 83 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b78ac8f9e9..a9ac21bb56 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1657,7 +1657,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## This is necessary because promptTokensDetails includes both cached and non-cached tokens ## See: https://github.com/BerriAI/litellm/issues/18750 if cached_text_tokens is not None and prompt_text_tokens is not None: + # Explicit caching: subtract cached tokens per modality from cacheTokensDetails prompt_text_tokens = prompt_text_tokens - cached_text_tokens + elif ( + cached_tokens is not None + and prompt_text_tokens is not None + and cached_text_tokens is None + ): + # Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails) + # Subtract from text tokens since implicit caching is primarily for text content + # See: https://github.com/BerriAI/litellm/issues/16341 + prompt_text_tokens = prompt_text_tokens - cached_tokens if cached_audio_tokens is not None and prompt_audio_tokens is not None: prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens if cached_image_tokens is not None and prompt_image_tokens is not None: diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 9d968d482c..f0e5cafda5 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1607,3 +1607,86 @@ def test_gemini_without_cache_tokens_details(): assert usage.prompt_tokens_details.text_tokens >= 0 print("āœ… Gemini without cacheTokensDetails works correctly") + + +def test_gemini_implicit_caching_cost_calculation(): + """ + Test for Issue #16341: Gemini implicit cached tokens not counted in spend log + + When Gemini uses implicit caching, it returns cachedContentTokenCount but NOT + cacheTokensDetails. In this case, we should subtract cachedContentTokenCount + from text_tokens to correctly calculate costs. + + See: https://github.com/BerriAI/litellm/issues/16341 + """ + from litellm import completion_cost + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.utils import Choices, Message, ModelResponse + + # Simulate Gemini response with implicit caching (cachedContentTokenCount only) + completion_response = { + "usageMetadata": { + "promptTokenCount": 10000, + "candidatesTokenCount": 5, + "totalTokenCount": 10005, + "cachedContentTokenCount": 8000, # Implicit caching - no cacheTokensDetails + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 10000}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], + } + } + + usage = VertexGeminiConfig._calculate_usage(completion_response) + + # Verify parsing + assert ( + usage.cache_read_input_tokens == 8000 + ), f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" + assert ( + usage.prompt_tokens_details.cached_tokens == 8000 + ), f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" + + # CRITICAL: text_tokens should be (10000 - 8000) = 2000, NOT 10000 + # This is the fix for issue #16341 + assert ( + usage.prompt_tokens_details.text_tokens == 2000 + ), f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" + + # Verify cost calculation uses cached token pricing + response = ModelResponse( + id="mock-id", + model="gemini-2.0-flash", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="Hello!"), + finish_reason="stop", + ) + ], + usage=usage, + ) + + cost = completion_cost( + completion_response=response, + model="gemini-2.0-flash", + custom_llm_provider="gemini", + ) + + # Get model pricing for verification + import litellm + + model_info = litellm.get_model_info("gemini/gemini-2.0-flash") + input_cost = model_info.get("input_cost_per_token", 0) + cache_read_cost = model_info.get("cache_read_input_token_cost", input_cost) + output_cost = model_info.get("output_cost_per_token", 0) + + # Expected cost: (2000 * input) + (8000 * cache_read) + (5 * output) + expected_cost = (2000 * input_cost) + (8000 * cache_read_cost) + (5 * output_cost) + + assert abs(cost - expected_cost) < 1e-9, ( + f"Cost calculation is wrong. Got ${cost:.6f}, expected ${expected_cost:.6f}. " + f"Cached tokens may not be using reduced pricing." + ) + + print("āœ… Issue #16341 fix verified: Gemini implicit caching cost calculated correctly") From 905e9cd6c9634f6b0bf1fcda685ace5124da1b56 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 28 Jan 2026 13:30:27 -0800 Subject: [PATCH 12/45] breakdown by team and keys --- .../EntityUsageExport/ExportTypeSelector.tsx | 10 +- .../src/components/EntityUsageExport/types.ts | 2 +- .../src/components/EntityUsageExport/utils.ts | 91 +++++++++++++++++++ 3 files changed, 101 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx index 17a833deac..6bc8594860 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx @@ -17,11 +17,19 @@ const ExportTypeSelector: React.FC = ({ value, onChange + +