From 07f8a17fffdd4700ff4c7d7433b9c62056522329 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 7 Sep 2026 16:23:53 -0700 Subject: [PATCH] fix(prometheus): distinguish missing usage and preserve config persistence Keep quoted database-storage config values intact for legacy persistence readers while using a local boolean for early callback discovery. Distinguish absent provider usage from an explicitly reported zero when labeling latency metrics. --- basedpyright-code-budget.json | 2 +- litellm/integrations/prometheus.py | 11 ++++- litellm/proxy/proxy_server.py | 18 +++---- ruff-strict-budget.json | 6 +-- ..._prometheus_input_sequence_length_label.py | 37 +++++++++++++- .../proxy/proxy_server/test_lifecycle.py | 3 +- .../proxy/proxy_server/test_proxy_config.py | 20 ++++---- tests/test_litellm/proxy/test_proxy_server.py | 48 +++++++++++++++++++ type-discipline-budget.json | 2 +- 9 files changed, 122 insertions(+), 25 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..931ebc46f51 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -117,7 +117,7 @@ "limit": 110 }, "reportUnnecessaryComparison": { - "limit": 687 + "limit": 686 }, "reportUnnecessaryContains": { "limit": 4 diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f2b19815563..d8a7027f62f 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1417,6 +1417,9 @@ class PrometheusLogger(CustomLogger): f"inside track_prometheus_metrics, model {model}, response_cost {response_cost}, tokens_used {tokens_used}, end_user_id {end_user_id}, user_api_key {user_api_key}" ) + reported_usage: Final = ( + response_obj.get("usage") if isinstance(response_obj, dict) else getattr(response_obj, "usage", None) + ) enum_values: Final = UserAPIKeyLabelValues( end_user=end_user_id, hashed_api_key=user_api_key, @@ -1445,7 +1448,13 @@ class PrometheusLogger(CustomLogger): stream=(str(standard_logging_payload.get("stream")) if litellm.prometheus_emit_stream_label else None), service_tier=get_service_tier_from_standard_logging_payload(standard_logging_payload), input_sequence_length=( - get_input_sequence_length_bucket(standard_logging_payload.get("prompt_tokens")) + get_input_sequence_length_bucket( + standard_logging_payload.get("prompt_tokens") + if standard_logging_payload.get("prompt_tokens") + or reported_usage is not None + or kwargs.get("combined_usage_object") is not None + else None + ) if self._emit_input_sequence_length_label else None ), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8e7920fc507..61f72e222b3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1189,15 +1189,15 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: general_settings=general_settings ) - store_model_in_db = ( # rebind-ok: startup publishes the combined env, YAML and DB setting before callback construction - await ProxyStartupEvent.resolve_store_model_in_db(prisma_client=prisma_client, configured=store_model_in_db) + should_load_db_litellm_settings: Final = await ProxyStartupEvent.resolve_store_model_in_db( + prisma_client=prisma_client, configured=store_model_in_db ) await ProxyStartupEvent._initialize_startup_logging( llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, redis_usage_cache=transaction_buffer_redis_cache, prisma_client=prisma_client, - should_load_db_litellm_settings=store_model_in_db, + should_load_db_litellm_settings=should_load_db_litellm_settings, proxy_config_obj=proxy_config, ) @@ -9024,8 +9024,9 @@ class ProxyStartupEvent: ) @staticmethod - async def resolve_store_model_in_db(prisma_client: PrismaClient | None, configured: bool) -> bool: - if (get_secret_bool("STORE_MODEL_IN_DB", configured) or configured) is True: + async def resolve_store_model_in_db(prisma_client: PrismaClient | None, configured: bool | str) -> bool: + default: Final = str_to_bool(configured) if isinstance(configured, str) else configured is True + if (get_secret_bool("STORE_MODEL_IN_DB", default) or configured) is True: return True if prisma_client is None: return False @@ -9591,8 +9592,9 @@ class ProxyStartupEvent: prisma_client.spend_logs_queue_monitor_task = monitor_task # rebind-ok: the client owns its monitor handle ### ADD NEW MODELS ### - store_model_in_db = await cls.resolve_store_model_in_db( - prisma_client=prisma_client, configured=store_model_in_db + store_model_in_db = ( # rebind-ok: preserve legacy YAML values unless env or DB explicitly enables storage + await cls.resolve_store_model_in_db(prisma_client=prisma_client, configured=store_model_in_db) + or store_model_in_db ) config_reload_interval_seconds = proxy_config_reload_interval_seconds @@ -17152,7 +17154,7 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie "forgiving it. Applies to key, user, team, team member, org, tag and end-user budgets." ), }, - "prometheus_emit_input_sequence_length_label": { # mutable-ok: nested registry literal; LIT002 exempts only TypedDict-annotated top-level literals + "prometheus_emit_input_sequence_length_label": { # mutable-ok: frozen with the registry below "type": "Boolean", "description": ( "Break latency and time-to-first-token metrics into input token length buckets. " diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..3ec8eea9807 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -21,7 +21,7 @@ "limit": 112 }, "ANN206": { - "limit": 133 + "limit": 132 }, "ANN401": { "limit": 119 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2915 }, "C401": { "limit": 8 @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 190 + "limit": 189 }, "PLW1508": { "limit": 190 diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py index e6f386abc03..a3527cf1016 100644 --- a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py +++ b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py @@ -25,7 +25,7 @@ FLAG: Final = "prometheus_emit_input_sequence_length_label" def _clear_prometheus_registry() -> None: - for collector in list(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] + for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # test registry reset REGISTRY.unregister(collector) @@ -164,6 +164,41 @@ async def test_logger_emits_bucket_from_its_startup_label_set( assert all(sample.labels["input_sequence_length"] == "4k-16k" for sample in samples) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "combined_usage", "expected"), + ( + ({"id": "moderation", "results": []}, None, "unknown"), + ({"usage": None}, None, "unknown"), + ({"usage": {"prompt_tokens": 0}}, None, "0-1k"), + (litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=0)), None, "0-1k"), + (None, litellm.Usage(prompt_tokens=0), "0-1k"), + ), +) +async def test_logger_distinguishes_missing_usage_from_reported_zero( + monkeypatch: pytest.MonkeyPatch, response: object, combined_usage: object, expected: str +): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=response if isinstance(response, dict) else None + ) + + await logger.async_log_success_event( + {**_success_kwargs(now, prompt_tokens=usage["prompt_tokens"]), "combined_usage_object": combined_usage}, + response, + now, + now, + ) + + samples: Final = _latency_bucket_samples() + assert samples + assert all(sample.labels["input_sequence_length"] == expected for sample in samples) + + @pytest.mark.asyncio async def test_logger_built_with_flag_off_emits_no_bucket_label(monkeypatch: pytest.MonkeyPatch): now: Final = datetime.datetime.now() diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 97d49bfb908..e82047d40d7 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -794,7 +794,6 @@ async def test_resolve_store_model_in_db_preserves_legacy_config_and_env_precede db_value: bool, expected: bool, ): - monkeypatch.setattr(ps, "store_model_in_db", configured) if environment is None: monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) else: @@ -804,7 +803,7 @@ async def test_resolve_store_model_in_db_preserves_legacy_config_and_env_precede return_value=MagicMock(param_value={"store_model_in_db": db_value}) ) - assert await ProxyStartupEvent.resolve_store_model_in_db(prisma_client, ps.store_model_in_db) is expected + assert await ProxyStartupEvent.resolve_store_model_in_db(prisma_client, configured) is expected assert prisma_client.db.litellm_config.find_first.await_count == ( 0 if configured is True or environment == "true" else 1 ) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 9bbe8fcc12d..8e8406d564f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -12,7 +12,7 @@ import json import os import re from types import SimpleNamespace -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -2914,17 +2914,21 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch @pytest.mark.asyncio @pytest.mark.parametrize( - ("db_value", "expected"), - ((True, True), (False, False), ("true", True), ("false", False), ("invalid", False)), + ("param_value", "expected"), + ( + ({"prometheus_emit_input_sequence_length_label": "true"}, True), + ('{"prometheus_emit_input_sequence_length_label": "true"}', True), + ('{"prometheus_emit_input_sequence_length_label": false}', False), + (None, False), + ('["prometheus_emit_input_sequence_length_label"]', False), + ), ) -async def test_apply_safe_litellm_settings_overrides_normalizes_input_sequence_length_flag( - monkeypatch: pytest.MonkeyPatch, db_value: object, expected: bool +async def test_apply_safe_litellm_settings_overrides_parses_input_sequence_length_record( + monkeypatch: pytest.MonkeyPatch, param_value: object, expected: bool ): from litellm.proxy import proxy_server - config_record = SimpleNamespace( - param_value={"prometheus_emit_input_sequence_length_label": db_value} - ) + config_record: Final = None if param_value is None else SimpleNamespace(param_value=param_value) monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=config_record)) monkeypatch.setattr(litellm, "prometheus_emit_input_sequence_length_label", False) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7773c857d6b..cdec7acce05 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7452,6 +7452,54 @@ async def test_store_model_in_db_db_override_when_config_false(): assert mock_proxy_config.get_credentials.call_count == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("configured", ("true", "false")) +async def test_scheduled_startup_preserves_quoted_config_persistence( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, configured: str +): + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.proxy.proxy_server import ProxyConfig, ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + prisma: Final = MagicMock(insert_data=AsyncMock()) + prisma.db.litellm_config.find_first = AsyncMock(return_value=None) + proxy_logging: Final = MagicMock(spec=ProxyLogging) + proxy_logging.slack_alerting_instance = MagicMock() + proxy_logging.db_spend_update_writer = MagicMock() + config: Final = _mock_scheduled_proxy_config() + scheduler: Final = AsyncIOScheduler() + config_file: Final = tmp_path / "config.yaml" + config_file.write_text("general_settings: {}\n") + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + monkeypatch.setattr(proxy_server_module, "prisma_client", prisma) + monkeypatch.setattr(proxy_server_module, "general_settings", {"store_model_in_db": configured}) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", configured) + monkeypatch.setattr(proxy_server_module, "proxy_config", config) + monkeypatch.setattr(proxy_server_module, "user_config_file_path", str(config_file)) + monkeypatch.setattr(proxy_server_module, "AsyncIOScheduler", MagicMock(return_value=scheduler)) + + try: + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=prisma, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=proxy_logging, + ) + await ProxyConfig().save_config(new_config={"litellm_settings": {"enable_anthropic_prompt_caching": True}}) + + assert proxy_server_module.store_model_in_db == configured + assert scheduler.get_job("add_deployment_job") is None + assert config_file.read_text() == "general_settings: {}\n" + prisma.insert_data.assert_awaited_once_with( + data={"litellm_settings": {"enable_anthropic_prompt_caching": True}}, table_name="config" + ) + finally: + scheduler.shutdown(wait=False) + + @pytest.mark.asyncio async def test_store_model_in_db_db_check_skipped_when_already_true(monkeypatch): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..970dc38072a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16424 }, "LIT011": { "limit": 5506