scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

Merge pull request #25589 from BerriAI/litellm_oss_staging_04_11_2026

Browse source 
Litellm oss staging 04 11 2026
This commit is contained in:
Sameer Kankute 2026-04-14 23:34:25 +05:30 • committed by GitHub
commit b8f7d61400
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 585 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/constants.py
View file

@ -1060,6 +1060,9 @@ WANDB_MODELS: set = set(
"Qwen/Qwen3-235B-A22B-Thinking-2507",
# moonshotai
"moonshotai/Kimi-K2-Instruct",
"moonshotai/Kimi-K2.5",
# MiniMaxAI
"MiniMaxAI/MiniMax-M2.5",
# meta models
"meta-llama/Llama-3.1-8B-Instruct",
"meta-llama/Llama-3.3-70B-Instruct",

28
litellm/model_prices_and_context_window_backup.json
View file

@ -32414,6 +32414,34 @@
"litellm_provider": "wandb",
"mode": "chat"
},
"wandb/moonshotai/Kimi-K2.5": {
"max_tokens": 262144,
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 6e-07,
"output_cost_per_token": 3e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/inference/coreweave/cw_moonshotai_Kimi-K2.5",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
},
"wandb/MiniMaxAI/MiniMax-M2.5": {
"max_tokens": 197000,
"max_input_tokens": 197000,
"max_output_tokens": 197000,
"input_cost_per_token": 3e-07,
"output_cost_per_token": 1.2e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/inference/coreweave/cw_MiniMaxAI_MiniMax-M2.5",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,

23
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -1,7 +1,7 @@
import asyncio
import json
import time
from datetime import datetime, timedelta, timezone
from datetime import datetime, timezone
from typing import List, Literal, Optional, Union
from litellm._logging import verbose_proxy_logger
@ -652,24 +652,13 @@ class ResetBudgetJob:
) -> LiteLLM_BudgetTableFull:
try:
if budget.budget_duration is not None:
from litellm.litellm_core_utils.duration_parser import (
duration_in_seconds,
from litellm.proxy.common_utils.timezone_utils import (
get_budget_reset_time,
)
duration_s = duration_in_seconds(duration=budget.budget_duration)
# Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account
if (
budget.budget_reset_at is None
and budget.created_at + timedelta(seconds=duration_s) > current_time
):
budget.budget_reset_at = budget.created_at + timedelta(
seconds=duration_s
)
else:
budget.budget_reset_at = current_time + timedelta(
seconds=duration_s
)
budget.budget_reset_at = get_budget_reset_time(
budget_duration=budget.budget_duration
)
except Exception as e:
verbose_proxy_logger.exception(
"Error resetting budget_reset_at for budget: %s. Item: %s", e, budget

11
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -255,7 +255,16 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
return
response_cost: float = standard_logging_payload.get("response_cost", 0)
model = standard_logging_payload.get("model")
# Use model_group (the user-facing model alias, e.g. "gpt-4o") when
# available. The enforcement path (is_key_within_model_budget) receives
# the model name from request_data["model"] which is the model group
# alias, so the spend tracking cache key must use the same name.
# Falling back to the deployment-level "model" field preserves
# behaviour for non-proxy or non-router deployments where model_group
# is None.
model = standard_logging_payload.get(
"model_group"
) or standard_logging_payload.get("model")
virtual_key = standard_logging_payload.get("metadata", {}).get(
"user_api_key_hash"
)

8
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -12,11 +12,9 @@ All /budget management endpoints
"""
#### BUDGET TABLE MANAGEMENT ####
from datetime import timedelta
from fastapi import APIRouter, Depends, HTTPException
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import jsonify_object
@ -86,8 +84,8 @@ async def new_budget(
# if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future
if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None:
budget_obj.budget_reset_at = datetime.utcnow() + timedelta(
seconds=duration_in_seconds(duration=budget_obj.budget_duration)
budget_obj.budget_reset_at = get_budget_reset_time(
budget_duration=budget_obj.budget_duration
)
budget_obj_json = budget_obj.model_dump(exclude_none=True)

4
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1849,7 +1849,7 @@ async def _process_single_key_update(
# Delete cache
await _delete_cache_key_object(
hashed_token=hash_token(key_update_item.key),
hashed_token=_hash_token_if_needed(key_update_item.key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -3726,7 +3726,7 @@ async def _execute_virtual_key_regeneration(
if hashed_api_key or key:
await _delete_cache_key_object(
hashed_token=hash_token(key),
hashed_token=_hash_token_if_needed(key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)

28
model_prices_and_context_window.json
View file

@ -32399,6 +32399,34 @@
"litellm_provider": "wandb",
"mode": "chat"
},
"wandb/moonshotai/Kimi-K2.5": {
"max_tokens": 262144,
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 6e-07,
"output_cost_per_token": 3e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/inference/coreweave/cw_moonshotai_Kimi-K2.5",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
},
"wandb/MiniMaxAI/MiniMax-M2.5": {
"max_tokens": 197000,
"max_input_tokens": 197000,
"max_output_tokens": 197000,
"input_cost_per_token": 3e-07,
"output_cost_per_token": 1.2e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/inference/coreweave/cw_MiniMaxAI_MiniMax-M2.5",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,

149
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
View file

@ -219,6 +219,155 @@ async def test_get_end_user_spend_for_model(budget_limiter):
assert spend == 50.0
@pytest.mark.asyncio
async def test_async_log_success_event_uses_model_group_for_cache_key(budget_limiter):
"""
When model_group is present in StandardLoggingPayload (proxy/router
deployments), spend must be tracked under the model_group name — not the
deployment-level model name — so the cache key matches the one used by
is_key_within_model_budget (which receives request_data["model"], the
model group alias).
Without this, providers that decorate model names (e.g. Vertex AI
"vertex_ai/claude-opus-4-6@default") track spend under a different cache
key than enforcement reads, silently disabling budget limits.
"""
from litellm.proxy.hooks.model_max_budget_limiter import (
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
virtual_key = "test-key-hash"
model_group = "claude-opus-4-6"
deployment_model = "vertex_ai/claude-opus-4-6@default"
budget_duration = "1d"
user_api_key_model_max_budget = {
model_group: {"budget_limit": 50.0, "time_period": budget_duration},
}
kwargs = {
"standard_logging_object": {
"response_cost": 0.10,
"model": deployment_model,
"model_group": model_group,
"metadata": {"user_api_key_hash": virtual_key},
},
"litellm_params": {
"metadata": {
"user_api_key_model_max_budget": user_api_key_model_max_budget,
},
},
}
with patch.object(
budget_limiter,
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
await budget_limiter.async_log_success_event(
kwargs, response_obj=None, start_time=None, end_time=None
)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
# The cache key must use the model_group name, NOT the deployment name
assert spend_key == (
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}"
)
assert call_kwargs["response_cost"] == 0.10
@pytest.mark.asyncio
async def test_async_log_success_event_falls_back_to_model_when_no_model_group(
budget_limiter,
):
"""
When model_group is None (non-proxy / non-router usage), spend tracking
must fall back to using the model field so existing behaviour is preserved.
"""
from litellm.proxy.hooks.model_max_budget_limiter import (
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
virtual_key = "test-key-hash"
model = "gpt-4"
budget_duration = "1d"
user_api_key_model_max_budget = {
model: {"budget_limit": 100.0, "time_period": budget_duration},
}
kwargs = {
"standard_logging_object": {
"response_cost": 0.05,
"model": model,
"model_group": None,
"metadata": {"user_api_key_hash": virtual_key},
},
"litellm_params": {
"metadata": {
"user_api_key_model_max_budget": user_api_key_model_max_budget,
},
},
}
with patch.object(
budget_limiter,
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
await budget_limiter.async_log_success_event(
kwargs, response_obj=None, start_time=None, end_time=None
)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
assert spend_key == (
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}"
)
@pytest.mark.asyncio
async def test_async_log_success_event_end_user_uses_model_group(budget_limiter):
"""
End-user model budget tracking must also use model_group when available,
matching the enforcement path in is_end_user_within_model_budget.
"""
from litellm.proxy.hooks.model_max_budget_limiter import (
END_USER_SPEND_CACHE_KEY_PREFIX,
)
end_user_id = "test-user"
model_group = "claude-sonnet-4-6"
deployment_model = "vertex_ai/claude-sonnet-4-6@default"
budget_duration = "1d"
user_api_key_end_user_model_max_budget = {
model_group: {"budget_limit": 25.0, "time_period": budget_duration},
}
kwargs = {
"standard_logging_object": {
"response_cost": 0.03,
"model": deployment_model,
"model_group": model_group,
"end_user": end_user_id,
"metadata": {"user_api_key_end_user_id": end_user_id},
},
"litellm_params": {
"metadata": {
"user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget,
},
},
}
with patch.object(
budget_limiter,
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
await budget_limiter.async_log_success_event(
kwargs, response_obj=None, start_time=None, end_time=None
)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
assert spend_key == (
f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}"
)
@pytest.mark.asyncio
async def test_async_log_success_event_uses_end_user_model_budget_duration(
budget_limiter,

43
tests/test_budget_management.py
View file

@ -1,12 +1,25 @@
# What is this?
## Unit tests for the /budget/* endpoints
from litellm._uuid import uuid
from datetime import datetime, timedelta
from datetime import datetime, timezone
import aiohttp
import pytest
import pytest_asyncio
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_timezone
def _parse_budget_api_datetime(value: str) -> datetime:
"""Parse ISO timestamps returned by the proxy JSON API."""
if value.endswith("Z"):
value = value[:-1] + "+00:00"
dt = datetime.fromisoformat(value)
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt
async def delete_budget(session, budget_id):
url = "http://0.0.0.0:4000/budget/delete"
@ -61,32 +74,30 @@ async def budget_setup():
@pytest.mark.asyncio
async def test_create_budget_with_duration(budget_setup):
"""
Test creating a budget with a specified duration and verify that the 'budget_reset_at'
timestamp is correctly calculated as 'created_at' plus the budget duration (one day).
This test uses the budget_setup fixture, which handles both the creation and cleanup of the budget.
Test creating a budget with a specified duration and verify that 'budget_reset_at'
matches the next standardized reset (see get_budget_reset_time / new_budget), not
necessarily created_at + wall-clock duration.
"""
# Verify that the response includes a 'budget_reset_at' timestamp.
assert (
budget_setup["budget_reset_at"] is not None
), "The budget_reset_at field should not be None"
# Calculate the expected reset time: created_at + 1 day.
# Replace trailing 'Z' with '+00:00' for Python 3.9 compat (fromisoformat
# only learned to accept 'Z' in Python 3.11).
created_at_str = budget_setup["created_at"].replace("Z", "+00:00")
expected_reset_at_date = datetime.fromisoformat(created_at_str) + timedelta(days=1)
created_at = _parse_budget_api_datetime(budget_setup["created_at"])
expected_reset_at = get_next_standardized_reset_time(
duration=budget_setup["budget_duration"],
current_time=created_at,
timezone_str=get_budget_reset_timezone(),
)
actual_reset_at = _parse_budget_api_datetime(budget_setup["budget_reset_at"])
# Allow for a small tolerance in seconds for the timestamp calculation.
tolerance_seconds = 3
reset_at_str = budget_setup["budget_reset_at"].replace("Z", "+00:00")
actual_reset_at_date = datetime.fromisoformat(reset_at_str)
time_difference = abs(
(actual_reset_at_date - expected_reset_at_date).total_seconds()
(actual_reset_at - expected_reset_at).total_seconds()
)
assert time_difference <= tolerance_seconds, (
f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at_date}, "
f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at}, "
f"but the difference was {time_difference} seconds."
)

134
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -444,6 +444,140 @@ def test_reset_budget_for_keys_linked_to_budgets_empty(
assert len(calls) == 0
@pytest.mark.parametrize(
"budget_duration, expected_day, expected_month",
[
("30d", 1, 7), # 30d → 1st of next month
("1mo", 1, 7), # 1mo → 1st of next month
("1d", 16, 6), # 1d → next midnight (same month)
],
ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"],
)
def test_reset_budget_reset_at_date_calendar_aligned(
budget_duration, expected_day, expected_month
):
"""
Verify that _reset_budget_reset_at_date produces calendar-aligned reset
times (matching get_budget_reset_time), not sliding-window offsets.
"""
from unittest.mock import patch
# Fix "now" to 2023-06-15 10:30:00 UTC for deterministic results
fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"budget_duration": budget_duration,
"budget_reset_at": fixed_now - timedelta(hours=1),
"budget_id": "test-budget",
"created_at": fixed_now - timedelta(days=30),
},
)
with patch(
"litellm.proxy.common_utils.timezone_utils.datetime"
) as mock_dt:
mock_dt.now.return_value = fixed_now
mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
asyncio.run(
ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)
)
assert test_budget.budget_reset_at.day == expected_day
assert test_budget.budget_reset_at.month == expected_month
assert test_budget.budget_reset_at.hour == 0
assert test_budget.budget_reset_at.minute == 0
assert test_budget.budget_reset_at.second == 0
def test_reset_budget_reset_at_date_7d_next_monday():
"""Verify 7d budget duration resets to next Monday at midnight."""
from unittest.mock import patch
# 2023-06-14 is a Wednesday
fixed_now = datetime(2023, 6, 14, 10, 30, 0, tzinfo=timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"budget_duration": "7d",
"budget_reset_at": fixed_now - timedelta(hours=1),
"budget_id": "test-budget",
"created_at": fixed_now - timedelta(days=7),
},
)
with patch(
"litellm.proxy.common_utils.timezone_utils.datetime"
) as mock_dt:
mock_dt.now.return_value = fixed_now
mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
asyncio.run(
ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)
)
# Next Monday after Wednesday June 14 is June 19
assert test_budget.budget_reset_at.day == 19
assert test_budget.budget_reset_at.month == 6
assert test_budget.budget_reset_at.weekday() == 0 # Monday
assert test_budget.budget_reset_at.hour == 0
def test_reset_budget_reset_at_date_none_duration():
"""Verify that budget_reset_at is unchanged when budget_duration is None."""
original_reset_at = datetime(2023, 6, 20, 0, 0, 0, tzinfo=timezone.utc)
now = datetime(2023, 6, 15, 10, 0, 0, tzinfo=timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"budget_duration": None,
"budget_reset_at": original_reset_at,
"budget_id": "test-budget",
"created_at": now - timedelta(days=30),
},
)
asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now))
assert test_budget.budget_reset_at == original_reset_at
def test_reset_budget_reset_at_date_none_reset_at():
"""Verify that budget_reset_at is set correctly even when previously None."""
from unittest.mock import patch
fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"budget_duration": "30d",
"budget_reset_at": None,
"budget_id": "test-budget",
"created_at": fixed_now - timedelta(days=5),
},
)
with patch(
"litellm.proxy.common_utils.timezone_utils.datetime"
) as mock_dt:
mock_dt.now.return_value = fixed_now
mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
asyncio.run(
ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)
)
# Should be set to 1st of next month (July 1)
assert test_budget.budget_reset_at is not None
assert test_budget.budget_reset_at.day == 1
assert test_budget.budget_reset_at.month == 7
def test_budget_table_reset_also_resets_linked_keys(
reset_budget_job, mock_prisma_client
):

180
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -5385,9 +5385,15 @@ async def test_bulk_update_keys_success(monkeypatch):
) as mock_hash:
mock_hash.side_effect = ["hashed-key-1", "hashed-key-2"]
def _hash_for_bulk_success(token: str) -> str:
return {
"test-key-1": "hashed-key-1",
"test-key-2": "hashed-key-2",
}[token]
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
side_effect=["hashed-key-1", "hashed-key-2"],
side_effect=_hash_for_bulk_success,
):
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
@ -5511,9 +5517,15 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
) as mock_hash:
mock_hash.return_value = "hashed-key-1"
def _hash_for_bulk_partial(token: str) -> str:
return {
"test-key-1": "hashed-key-1",
"non-existent-key": "hashed-non-existent-key",
}[token]
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
side_effect=["hashed-key-1", "hashed-non-existent-key"],
side_effect=_hash_for_bulk_partial,
):
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
@ -8746,3 +8758,167 @@ def test_validate_public_image_url_accepts_http_and_noop_empty():
_validate_public_image_url(None, "logo_url")
_validate_public_image_url("", "logo_url")
_validate_public_image_url(" ", "logo_url")
@pytest.mark.asyncio
async def test_process_single_key_update_cache_invalidation_with_token_hash():
"""
_process_single_key_update must pass the token hash as-is (not
double-hashed) to _delete_cache_key_object when the key is already a
pre-hashed token ID rather than an sk- prefixed key.
Without this, cache invalidation silently fails: the wrong cache entry
is deleted while the stale entry (with outdated fields) persists and
gets refreshed indefinitely by update_cache on every successful request.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_process_single_key_update,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequestItem,
)
token_hash = "abc123def456"
existing_key = LiteLLM_VerificationToken(
token=token_hash,
user_id="user-1",
models=["gpt-4"],
team_id=None,
max_budget=None,
tags=None,
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=existing_key
)
mock_updated = MagicMock()
mock_updated.model_dump.return_value = {"max_budget": 100.0}
mock_prisma_client.update_data = AsyncMock(return_value={"data": mock_updated})
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
mock_llm_router = MagicMock()
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data",
return_value={"max_budget": 100.0},
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
return_value=None,
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
) as mock_delete_cache, patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
new_callable=AsyncMock,
):
key_update_item = BulkUpdateKeyRequestItem(
key=token_hash,
max_budget=100.0,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
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,
)
mock_delete_cache.assert_called_once()
call_kwargs = mock_delete_cache.call_args.kwargs
# The token hash should be passed as-is, NOT double-hashed
assert call_kwargs["hashed_token"] == token_hash
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_hash():
"""
_execute_virtual_key_regeneration must pass the token hash as-is (not
double-hashed) to _delete_cache_key_object when the key is a
pre-hashed token ID.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
token_hash = "abc123def456"
existing_key = LiteLLM_VerificationToken(
token=token_hash,
user_id="user-1",
models=["gpt-4"],
team_id=None,
max_budget=None,
tags=None,
)
mock_prisma_client = AsyncMock()
# _execute_virtual_key_regeneration calls dict(updated_token) which
# needs the return value to be iterable as key-value pairs.
class DictLikeResult:
def __init__(self, data):
self._data = data
def __iter__(self):
return iter(self._data.items())
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
return_value=DictLikeResult({"token": "new-hashed-token", "key_name": "sk-...ab12", "user_id": "user-1"})
)
mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock(
return_value=None
)
mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data)
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
) as mock_delete_cache, patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data",
new_callable=AsyncMock,
return_value={},
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=existing_key,
hashed_api_key=token_hash,
key=token_hash,
data=None,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
user_api_key_cache=mock_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
)
mock_delete_cache.assert_called_once()
call_kwargs = mock_delete_cache.call_args.kwargs
# The token hash should be passed as-is, NOT double-hashed
assert call_kwargs["hashed_token"] == token_hash

17
tests/test_litellm/test_cost_calculator.py
View file

@ -93,6 +93,23 @@ def test_baseten_model_api_pricing_entries():
assert model_info["output_cost_per_token"] == output_cost
def test_wandb_model_api_pricing_entries():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
expected_pricing = {
"wandb/moonshotai/Kimi-K2.5": (6e-07, 3e-06),
"wandb/MiniMaxAI/MiniMax-M2.5": (3e-07, 1.2e-06),
}
for model_name, (input_cost, output_cost) in expected_pricing.items():
model_info = litellm.model_cost.get(model_name)
assert model_info is not None, f"Missing model pricing entry: {model_name}"
assert model_info["litellm_provider"] == "wandb"
assert model_info["input_cost_per_token"] == input_cost
assert model_info["output_cost_per_token"] == output_cost
def test_cost_calculator_with_usage(monkeypatch):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")