fix(proxy): parse MAX_API_KEYS_IN_USAGE_BREAKDOWN with the ranged env helper
Some checks failed
ai-gateway image / ai-gateway release image (push) Has been cancelled

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 23:51:11 +00:00
parent 72fee17069
commit 3883f590f1
2 changed files with 28 additions and 2 deletions

View file

@ -1,5 +1,4 @@
import asyncio
import os
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
@ -11,6 +10,7 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.litellm_core_utils.env_utils import get_env_int_in_range
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
@ -53,7 +53,12 @@ _PRISMA_TO_PG_TABLE: Final[Mapping[str, str]] = {
"litellm_dailytagspend": "LiteLLM_DailyTagSpend",
}
MAX_API_KEYS_IN_USAGE_BREAKDOWN: Final = int(os.getenv("MAX_API_KEYS_IN_USAGE_BREAKDOWN", "100"))
MAX_API_KEYS_IN_USAGE_BREAKDOWN: Final = get_env_int_in_range(
"MAX_API_KEYS_IN_USAGE_BREAKDOWN",
default=100,
minimum=1,
maximum=10_000,
)
class DailySpendRecord(Protocol):

View file

@ -1,3 +1,4 @@
import importlib
import re
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
@ -10,6 +11,7 @@ from psycopg.rows import dict_row
from pytest_postgresql import factories
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.proxy.management_endpoints import common_daily_activity
from litellm.proxy.management_endpoints.common_daily_activity import (
MAX_API_KEYS_IN_USAGE_BREAKDOWN,
_adjust_dates_for_timezone,
@ -1508,6 +1510,25 @@ async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
assert set(day.breakdown.endpoints["/v1/chat/completions"].api_key_breakdown) == expected_top
@pytest.mark.parametrize("configured", ["0", "-5", "not-an-int", "", "10001"])
def test_invalid_api_key_cap_config_falls_back_to_default(monkeypatch, configured):
monkeypatch.setenv("MAX_API_KEYS_IN_USAGE_BREAKDOWN", configured)
try:
assert importlib.reload(common_daily_activity).MAX_API_KEYS_IN_USAGE_BREAKDOWN == 100
finally:
monkeypatch.delenv("MAX_API_KEYS_IN_USAGE_BREAKDOWN")
importlib.reload(common_daily_activity)
def test_valid_api_key_cap_config_is_honoured(monkeypatch):
monkeypatch.setenv("MAX_API_KEYS_IN_USAGE_BREAKDOWN", "250")
try:
assert importlib.reload(common_daily_activity).MAX_API_KEYS_IN_USAGE_BREAKDOWN == 250
finally:
monkeypatch.delenv("MAX_API_KEYS_IN_USAGE_BREAKDOWN")
importlib.reload(common_daily_activity)
def _no_spend_record():
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
return SimpleNamespace(