From 69f30e26c230960373f37c6c97fcb071cf0b0671 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 12 Oct 2026 18:58:26 +0000 Subject: [PATCH] refactor(proxy): type first cleanup run kwargs as a ReadOnly TypedDict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 14 +++++++++----- litellm/proxy/proxy_server.py | 2 +- .../proxy/proxy_server/test_proxy_config.py | 17 +++++++++-------- tests/unit/proxy/test_spend_log_cleanup.py | 4 ++-- 4 files changed, 21 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 4a7c33d61e1..04f67d38a0b 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -5,10 +5,10 @@ from collections.abc import Mapping from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from types import MappingProxyType -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, TypeAlias, TypedDict from pydantic import TypeAdapter +from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache @@ -100,7 +100,11 @@ SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = ( ) -def first_cleanup_run_kwargs(settings: Mapping[str, object], now: datetime) -> Mapping[str, datetime]: +class FirstCleanupRunKwargs(TypedDict, total=False): + next_run_time: ReadOnly[datetime] + + +def first_cleanup_run_kwargs(settings: Mapping[str, object], now: datetime) -> FirstCleanupRunKwargs: """ Extra ``add_job`` kwargs for the interval cleanup job. Empty unless ``maximum_spend_logs_cleanup_run_on_startup`` is true, so the first run otherwise stays one interval out. @@ -109,11 +113,11 @@ def first_cleanup_run_kwargs(settings: Mapping[str, object], now: datetime) -> M settings.get("maximum_spend_logs_cleanup_cron") or settings.get("maximum_spend_logs_cleanup_run_on_startup") is not True ): - return MappingProxyType({}) + return {} delay_seconds: Final = SPEND_LOG_CLEANUP_STARTUP_DELAY_SECONDS + random.randint( 0, SPEND_LOG_CLEANUP_STARTUP_JITTER_SECONDS ) - return MappingProxyType({"next_run_time": now + timedelta(seconds=delay_seconds)}) + return {"next_run_time": now + timedelta(seconds=delay_seconds)} class SpendLogCleanup: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9df43f3a5ef..30fe20edaf7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8014,7 +8014,7 @@ class ProxyConfig: id="spend_log_cleanup_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - **first_cleanup_run_kwargs(general_settings, datetime.now(timezone.utc)), + **first_cleanup_run_kwargs(_current_general_settings(), datetime.now(timezone.utc)), ) verbose_proxy_logger.info("Spend log cleanup rescheduled with trigger: %s", trigger) diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 6fff4f43c3c..95a1879bffa 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -4239,17 +4239,18 @@ def _paused_scheduler(monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("run_on_startup", [None, False, True]) async def test_ProxyConfig__reschedule_spend_log_cleanup_job_first_run_only_moves_up_when_opted_in( - monkeypatch, run_on_startup + monkeypatch: pytest.MonkeyPatch, run_on_startup: bool | None ): - real_scheduler = _paused_scheduler(monkeypatch) - settings = {"maximum_health_check_retention_period": "90d", "maximum_spend_logs_retention_interval": "1d"} - if run_on_startup is not None: - settings["maximum_spend_logs_cleanup_run_on_startup"] = run_on_startup - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", settings) - before = datetime.now(timezone.utc) + real_scheduler: Final = _paused_scheduler(monkeypatch) + opt_in: Final = {} if run_on_startup is None else {"maximum_spend_logs_cleanup_run_on_startup": run_on_startup} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"maximum_health_check_retention_period": "90d", "maximum_spend_logs_retention_interval": "1d", **opt_in}, + ) + before: Final = datetime.now(timezone.utc) try: await ProxyConfig()._reschedule_spend_log_cleanup_job() - delay = (real_scheduler.get_job("spend_log_cleanup_job").next_run_time - before).total_seconds() + delay: Final = (real_scheduler.get_job("spend_log_cleanup_job").next_run_time - before).total_seconds() if run_on_startup: assert delay <= SPEND_LOG_CLEANUP_STARTUP_DELAY_SECONDS + SPEND_LOG_CLEANUP_STARTUP_JITTER_SECONDS + 5 else: diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index f6554ae658b..87249a4af4a 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -6,7 +6,7 @@ import asyncio import logging import math import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from typing import Final @@ -46,7 +46,7 @@ from tests.unit.proxy.db.fake_prisma_engine import engine_call }, ], ) -def test_first_cleanup_run_kwargs_leaves_the_schedule_alone_unless_opted_in(settings): +def test_first_cleanup_run_kwargs_leaves_the_schedule_alone_unless_opted_in(settings: Mapping[str, object]): assert dict(first_cleanup_run_kwargs(settings, datetime(2026, 1, 1, tzinfo=timezone.utc))) == {}