mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
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
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:
parent
72fee17069
commit
3883f590f1
2 changed files with 28 additions and 2 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue