diff --git a/litellm/__init__.py b/litellm/__init__.py index 0fe3d091c2a..f6604a3b4e5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -551,6 +551,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, ma model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount +apply_search_free_tier: bool = False cost_margin_config: Dict[ str, Union[float, Dict[str, float]] ] = {} # Provider-specific or global cost margins. Examples: diff --git a/litellm/constants.py b/litellm/constants.py index bc94760b187..a0dcf44d121 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1855,6 +1855,9 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = get_env_float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", 300.0) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", 30.0) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = get_env_int("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", 100000) +SPEND_LOG_SIZE_CHECK_JOB_NAME: Final = "spend_log_size_check" +SPEND_LOG_SIZE_CHECK_INTERVAL: Final = os.getenv("SPEND_LOG_SIZE_CHECK_INTERVAL", "1d") +SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES: Final = get_env_int("SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES", 10) SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = get_env_float( "SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", 5.0 ) @@ -1976,6 +1979,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "public_model_groups", "public_model_groups_links", "cost_discount_config", + "apply_search_free_tier", "cost_margin_config", "block_requests_for_models_without_pricing", "budget_exceeded_throttle_percentage", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 903210c57e5..503877a9d4d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28948,6 +28948,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_token_batches": 1.5e-07, @@ -29062,6 +29063,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3-pro-image-preview": { @@ -29114,6 +29116,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { @@ -29164,6 +29167,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image-preview": { @@ -29207,6 +29211,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite-image": { @@ -29302,6 +29307,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29363,6 +29369,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -29422,6 +29429,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29508,6 +29516,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-08, "input_cost_per_token_batches": 5e-08, @@ -29562,6 +29571,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29608,6 +29618,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29795,6 +29806,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, @@ -29832,6 +29844,7 @@ "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29912,6 +29925,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29968,6 +29982,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3-flash-preview": { @@ -29992,6 +30007,7 @@ "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30074,6 +30090,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -30138,6 +30155,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30198,6 +30216,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30224,6 +30243,7 @@ "prompt_cache_min_tokens": 4096.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30369,6 +30389,7 @@ "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30455,6 +30476,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30782,6 +30804,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "supports_reasoning": false }, @@ -30826,6 +30849,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-lite-preview": { @@ -30887,6 +30911,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -31036,6 +31061,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_audio_token_priority": 1.8e-06, @@ -31148,6 +31174,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "supports_reasoning": true }, @@ -31212,6 +31239,7 @@ "output_cost_per_token_batches": 1.5e-06, "rpm": 1000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31362,6 +31390,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 1.5e-07, "input_cost_per_audio_token_priority": 5.4e-07, @@ -31567,6 +31596,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -31664,6 +31694,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07, @@ -31725,6 +31756,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -31781,6 +31813,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 5e-07, @@ -31844,6 +31877,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -31909,6 +31943,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -31970,6 +32005,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32031,6 +32067,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32126,6 +32163,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32189,6 +32227,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32217,6 +32256,7 @@ "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -32331,6 +32371,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -32395,6 +32436,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32455,6 +32497,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32480,6 +32523,7 @@ "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -52449,6 +52493,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -52511,6 +52556,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -52571,6 +52617,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 523f478d771..315533d6d0f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2884,6 +2884,30 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "deleted. Only historical tag usage analytics are affected; tag budgets read the lifetime counter." ), ) + maximum_spend_logs_size: str | None = Field( + None, + description=( + "Size rule for spend logs (e.g. '180GB'): when LiteLLM_SpendLogs plus LiteLLM_SpendLogToolIndex " + "(heap, indexes and TOAST, measured with pg_total_relation_size) hold more live data than this, the " + "cleanup job deletes the oldest spend logs first until they are under " + "maximum_spend_logs_cleanup_target_size. Independent of maximum_spend_logs_retention_period: a row is " + "deleted when either rule says so. Unset means no size rule." + ), + ) + maximum_spend_logs_cleanup_target_size: str | None = Field( + None, + description=( + "Level a size-triggered cleanup run deletes down to (e.g. '20GB'). Must be below maximum_spend_logs_size. " + "Unset means the run stops as soon as the tables are back under maximum_spend_logs_size." + ), + ) + maximum_spend_logs_size_check_interval: str | None = Field( + None, + description=( + "How often the size rule is checked (e.g. '10m'). Defaults to '1d'. Ignored when " + "maximum_spend_logs_cleanup_cron is set, which then drives both the age and the size cleanup." + ), + ) use_spend_logs_partitioning: bool | None = Field( None, description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.", 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 b57952b4992..268f0230d7a 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,6 +1,8 @@ import asyncio +import math import re import time +from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -19,6 +21,7 @@ from litellm.constants import ( SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP, SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS, SPEND_LOG_RUN_LOOPS, + SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.db.db_span import db_span @@ -35,7 +38,7 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( from litellm.proxy.utils import PrismaClient from litellm.types.llms.base import LiteLLMBaseModel -StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"] +StopReason: TypeAlias = Literal["exhausted", "target_reached", "budget_exhausted", "batch_cap_reached", "aborted"] _SIZE_RE: Final = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([KMGT]?)I?B?\s*$", re.IGNORECASE) _SIZE_UNITS: Final = {"": 1, "K": 1024, "M": 1024**2, "G": 1024**3, "T": 1024**4} @@ -61,6 +64,14 @@ def size_in_bytes(value: str | int) -> int: return size +@dataclass(frozen=True, slots=True) +class SizeRule: + """The size cleanup rule as configured: start cleaning above `limit_bytes`, stop under `target_bytes`.""" + + limit_bytes: int + target_bytes: int + + @dataclass(frozen=True, slots=True) class StorageProbe: """ @@ -218,6 +229,42 @@ class TableCleanupResult: stop_reason: StopReason +@dataclass(frozen=True, slots=True) +class _SizeStepState: + """Progress of the size pass across its re-check steps.""" + + logs_deleted: int + index_deleted: int + batches_left: int + stop_reason: StopReason | None + + def stopped(self, reason: StopReason) -> "_SizeStepState": + return _SizeStepState( + logs_deleted=self.logs_deleted, + index_deleted=self.index_deleted, + batches_left=self.batches_left, + stop_reason=reason, + ) + + +@dataclass(frozen=True, slots=True) +class SizePassReport: + """What one size pass measured and achieved, for the run summary and the storage gauges.""" + + rule: SizeRule + before: SpendLogStorage + after: SpendLogStorage + results: tuple[TableCleanupResult, ...] + + @property + def target_reached(self) -> bool: + return self.after.estimated_live_bytes <= self.rule.target_bytes + + @property + def above_target_bytes(self) -> int: + return max(0, self.after.estimated_live_bytes - self.rule.target_bytes) + + class _RunProgress: """How far one cleanup run has got, reported if that run is cancelled""" @@ -277,6 +324,13 @@ _SAVE_SIZE_STATE_SQL: Final = """ """ +class _OffsetRow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + startTime: datetime + + +_OFFSET_ROWS: Final = TypeAdapter(list[_OffsetRow]) + _SPEND_LOG_TABLES: Final = ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex") # Partition children are counted with their parent: a partitioned parent has no storage of its own. @@ -456,6 +510,53 @@ class SpendLogCleanup: self.retention_seconds = self._retention_seconds_for("maximum_spend_logs_retention_period") return self.retention_seconds is not None + def _setting(self, setting_name: str) -> object: + return self.general_settings.get(setting_name) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped settings dict + + def _size_bytes_setting(self, setting_name: str) -> int | None: + """Parse one size setting into bytes, or None when unset or unusable.""" + raw: Final = self._setting(setting_name) + if raw is None: + return None + if not isinstance(raw, (str, int)): + verbose_proxy_logger.warning("Invalid %s value: %r, ignoring it", setting_name, raw) + return None + try: + return size_in_bytes(raw) + except ValueError as e: + verbose_proxy_logger.warning("Invalid %s value: %r (%s), ignoring it", setting_name, raw, e) + return None + + def _size_rule(self) -> SizeRule | None: + """ + The size rule from settings, or None when maximum_spend_logs_size is unset. + + A target at or above the limit cannot be what the operator meant, so it is reported and the run + cleans only until it is back under the limit, which is also what an unset target means. + """ + limit_bytes: Final = self._size_bytes_setting("maximum_spend_logs_size") + if limit_bytes is None: + return None + target_bytes: Final = self._size_bytes_setting("maximum_spend_logs_cleanup_target_size") + if target_bytes is None: + return SizeRule(limit_bytes=limit_bytes, target_bytes=limit_bytes) + if target_bytes >= limit_bytes: + verbose_proxy_logger.error( + "maximum_spend_logs_cleanup_target_size (%s bytes) must be below maximum_spend_logs_size (%s bytes); " + "cleaning down to maximum_spend_logs_size instead", + target_bytes, + limit_bytes, + ) + return SizeRule(limit_bytes=limit_bytes, target_bytes=limit_bytes) + return SizeRule(limit_bytes=limit_bytes, target_bytes=target_bytes) + + def _size_pass_runs_with_age_pass(self) -> bool: + """ + With maximum_spend_logs_cleanup_cron set there is one schedule for both rules, so the age run carries + the size pass; otherwise the size pass has its own job on maximum_spend_logs_size_check_interval. + """ + return bool(self._setting("maximum_spend_logs_cleanup_cron")) + async def _probe_storage(self, prisma_client: PrismaClient, deadline: float) -> StorageProbe | None: """Read what the spend-log tables weigh from the catalog, or None when the probe fails or has no budget.""" if time.monotonic() >= deadline: @@ -508,6 +609,32 @@ class SpendLogCleanup: except Exception as e: # noqa: BLE001 - a lost save costs one re-seed from the catalog, it must not fail the run verbose_proxy_logger.warning("Could not save the spend-log size cleanup state: %s: %s", type(e).__name__, e) + async def _size_step_cutoff(self, prisma_client: PrismaClient, deadline: float, rows: int) -> datetime | None: + """ + The startTime that puts the `rows` oldest rows behind it. + + Expressing a size step as a cutoff lets the size pass reuse the retention DELETE unchanged, including + the tool-index rows that expire with their spend logs. When fewer rows than that remain, the step is + everything up to now. + """ + if time.monotonic() >= deadline: + return None + timeout_ms: Final = self._timeout_ms(deadline) + async with db_span("size_step_cutoff", "LiteLLM_SpendLogs"), bounded_tx(prisma_client, timeout_ms) as tx: + await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}")) + found: Final = _OFFSET_ROWS.validate_python( + await tx.query_raw( + trusted_sql( + 'SELECT "startTime" FROM "LiteLLM_SpendLogs" ORDER BY "startTime" ASC OFFSET $1 LIMIT 1' + ), + rows, + ) + ) + if not found: + return datetime.now(timezone.utc) + start_time: Final = found[0].startTime + return start_time if start_time.tzinfo is not None else start_time.replace(tzinfo=timezone.utc) + def _timeout_ms(self, deadline: float) -> int: """ The per-statement bound in milliseconds: the batch timeout, or whatever @@ -621,14 +748,17 @@ class SpendLogCleanup: key_columns: tuple[str, ...], time_column: str, deadline: float, + batch_cap: int | None = None, ) -> TableCleanupResult: """ Delete a table's rows older than the cutoff in batches. Stops at whichever bound is reached first: the backlog running out, the - shared wall-clock deadline, the per-table batch cap, or too many - consecutive batch failures. + shared wall-clock deadline, the per-table batch cap (`batch_cap` when a + caller spreads one table's cap over several calls, else the configured + one), or too many consecutive batch failures. """ + max_batches: Final = self.max_batches if batch_cap is None else batch_cap key_list: Final = ", ".join(f'"{col}"' for col in key_columns) delete_sql: Final = f""" DELETE FROM "{table_name}" @@ -651,7 +781,7 @@ class SpendLogCleanup: return await self._finish_table( prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline ) - if run_count >= self.max_batches: + if run_count >= max_batches: verbose_proxy_logger.info( "Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name ) @@ -891,6 +1021,163 @@ class SpendLogCleanup: verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted) return (logs_result, index_result) + async def _clean_spend_logs_by_size( + self, prisma_client: PrismaClient, rule: SizeRule, deadline: float + ) -> SizePassReport | None: + """ + Bring the spend-log tables under the size target, oldest rows first. + + Live bytes are the files on disk less the reclaimable count carried in LiteLLM_Config (see + `SizeCleanupState`), rolled forward from the catalog's cumulative insert and delete counters. Rows go in steps: the rows still needed to reach the target follow + from bytes per row, and each step turns that many of the oldest rows (at most + SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches) into a cutoff and runs the same bounded DELETE the age + rule uses, for the spend logs and the tool index rows derived from them. Progress is counted from the + deletes themselves, not re-read from pg_stat_user_tables, whose counters trail a DELETE by up to a + second. The run's batch cap and budget apply to the whole pass, so a pass that stops short reports how + far above the target it still is and the next check continues from there. + """ + probe: Final = await self._probe_storage(prisma_client, deadline) + if probe is None: + return None + if probe.stats_are_stale: + verbose_proxy_logger.warning( + "Spend-log row statistics are stale (allocated=%s bytes, live_rows=0, dead_rows=0); skipping the " + 'size pass until pg_stat_user_tables is repopulated, ANALYZE "LiteLLM_SpendLogs" refreshes it', + probe.allocated_bytes, + ) + return None + state: Final = advance_size_state(await self._load_size_state(prisma_client, deadline), probe) + before: Final = SpendLogStorage( + allocated_bytes=probe.allocated_bytes, + reclaimable_bytes=state.reclaimable_bytes, + bytes_per_row=state.bytes_per_row, + live_rows=probe.live_rows, + ) + SpendLogCleanupMetrics.set_storage_bytes("allocated", before.allocated_bytes) + SpendLogCleanupMetrics.set_storage_bytes("estimated_live", before.estimated_live_bytes) + if before.estimated_live_bytes <= rule.limit_bytes: + verbose_proxy_logger.info( + "Spend-log storage is under maximum_spend_logs_size (estimated_live=%s allocated=%s reclaimable=%s " + "limit=%s bytes)", + before.estimated_live_bytes, + before.allocated_bytes, + before.reclaimable_bytes, + rule.limit_bytes, + ) + SpendLogCleanupMetrics.set_storage_bytes("above_target", 0) + await self._save_size_state(prisma_client, deadline, state) + return SizePassReport(rule=rule, before=before, after=before, results=()) + verbose_proxy_logger.warning( + "Spend-log storage is over maximum_spend_logs_size (estimated_live=%s allocated=%s reclaimable=%s " + "limit=%s bytes); deleting the oldest spend logs down to %s bytes", + before.estimated_live_bytes, + before.allocated_bytes, + before.reclaimable_bytes, + rule.limit_bytes, + rule.target_bytes, + ) + rows_to_delete: Final = max( + 0, math.ceil((before.estimated_live_bytes - rule.target_bytes) / before.bytes_per_row) + ) + + def rows_still_needed(rows_deleted: int) -> int: + return rows_to_delete - rows_deleted + + results: Final = await self._delete_oldest_rows_for_size(prisma_client, deadline, rows_still_needed) + rows_deleted: Final = sum(result.rows_deleted for result in results) + remeasured: Final = await self._probe_storage(prisma_client, deadline) + after: Final = before.after_deleting( + rows_deleted, remeasured.allocated_bytes if remeasured is not None else before.allocated_bytes + ) + await self._save_size_state( + prisma_client, + deadline, + after.to_state(inserted_total=probe.inserted_total, deleted_total=probe.deleted_total + rows_deleted), + ) + report: Final = SizePassReport(rule=rule, before=before, after=after, results=results) + SpendLogCleanupMetrics.set_storage_bytes("allocated", after.allocated_bytes) + SpendLogCleanupMetrics.set_storage_bytes("estimated_live", after.estimated_live_bytes) + SpendLogCleanupMetrics.set_storage_bytes("above_target", report.above_target_bytes) + return report + + async def _delete_oldest_rows_for_size( + self, + prisma_client: PrismaClient, + deadline: float, + rows_still_needed: Callable[[int], int], + ) -> tuple[TableCleanupResult, ...]: + state = _SizeStepState( # rebind-ok: one value threaded through the bounded step loop + logs_deleted=0, index_deleted=0, batches_left=self.max_batches, stop_reason=None + ) + for _ in range(self.max_batches // SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES + 2): + if state.stop_reason is not None: + break + state = await self._run_size_step(prisma_client, deadline, rows_still_needed, state) + final_reason: Final[StopReason] = state.stop_reason or "batch_cap_reached" + return ( + TableCleanupResult( + table_name="LiteLLM_SpendLogs (size rule)", rows_deleted=state.logs_deleted, stop_reason=final_reason + ), + TableCleanupResult( + table_name="LiteLLM_SpendLogToolIndex (size rule)", + rows_deleted=state.index_deleted, + stop_reason=final_reason, + ), + ) + + async def _run_size_step( + self, + prisma_client: PrismaClient, + deadline: float, + rows_still_needed: Callable[[int], int], + state: "_SizeStepState", + ) -> "_SizeStepState": + """One step: delete the rows still needed, at most SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches of the oldest.""" + needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted) + if needed <= 0: + return state.stopped("target_reached") + if time.monotonic() >= deadline: + return state.stopped("budget_exhausted") + if state.batches_left <= 0: + return state.stopped("batch_cap_reached") + step_cap: Final = min( + SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES, state.batches_left, math.ceil(needed / self.batch_size) + ) + cutoff: Final = await self._size_step_cutoff(prisma_client, deadline, step_cap * self.batch_size) + if cutoff is None: + return state.stopped("budget_exhausted") + logs_result: Final = await self._delete_old_rows_batched( + prisma_client, + cutoff, + table_name="LiteLLM_SpendLogs", + key_columns=("request_id", "startTime"), + time_column="startTime", + deadline=deadline, + batch_cap=step_cap, + ) + index_result: Final = await self._delete_old_rows_batched( + prisma_client, + cutoff, + table_name="LiteLLM_SpendLogToolIndex", + key_columns=("request_id", "tool_name"), + time_column="start_time", + deadline=deadline, + batch_cap=step_cap, + ) + advanced: Final = _SizeStepState( + logs_deleted=state.logs_deleted + logs_result.rows_deleted, + index_deleted=state.index_deleted + index_result.rows_deleted, + batches_left=state.batches_left - step_cap, + stop_reason=None, + ) + if rows_still_needed(advanced.logs_deleted + advanced.index_deleted) <= 0: + return advanced.stopped("target_reached") + if logs_result.stop_reason in ("aborted", "budget_exhausted") or index_result.stop_reason == "aborted": + return advanced.stopped(logs_result.stop_reason if logs_result.stop_reason != "exhausted" else "aborted") + if logs_result.rows_deleted == 0: + return advanced.stopped("exhausted") + return advanced + async def _clean_session_rollup( self, prisma_client: PrismaClient, retention_seconds: int, deadline: float ) -> tuple[TableCleanupResult, ...]: @@ -968,22 +1255,60 @@ class SpendLogCleanup: return "completed" @staticmethod - def _log_run_summary(outcome: RunOutcome, results: tuple[TableCleanupResult, ...], elapsed_seconds: float) -> None: + def _storage_summary(report: SizePassReport | None) -> str: + """ + The size pass in one clause: what was measured, where it ended against the target, and whether the + operator still has to act to get the disk back. + """ + if report is None: + return "" + rows_deleted: Final = sum(result.rows_deleted for result in report.results) + reclaim_note: Final = ( + "; allocated bytes stay until autovacuum reuses them, only VACUUM FULL or pg_repack on the spend-log " + "tables returns them to the filesystem" + if rows_deleted > 0 and report.after.allocated_bytes >= report.before.allocated_bytes + else "" + ) + return ( + f" storage: allocated_before={report.before.allocated_bytes} " + f"estimated_live_before={report.before.estimated_live_bytes} " + f"allocated_after={report.after.allocated_bytes} estimated_live_after={report.after.estimated_live_bytes} " + f"reclaimable_after={report.after.reclaimable_bytes} " + f"limit={report.rule.limit_bytes} target={report.rule.target_bytes} " + f"target_reached={str(report.target_reached).lower()} above_target={report.above_target_bytes}{reclaim_note}" + ) + + @classmethod + def _log_run_summary( + cls, + outcome: RunOutcome, + results: tuple[TableCleanupResult, ...], + elapsed_seconds: float, + size_report: SizePassReport | None = None, + ) -> None: per_table: Final = ", ".join( f"{result.table_name}: deleted={result.rows_deleted} stop_reason={result.stop_reason}" for result in results ) - message: Final = "Spend log cleanup run finished: outcome=%s elapsed=%.1fs [%s]" - if outcome == "completed": - verbose_proxy_logger.info(message, outcome, elapsed_seconds, per_table) + message: Final = "Spend log cleanup run finished: outcome=%s elapsed=%.1fs [%s]%s" + storage: Final = cls._storage_summary(size_report) + if outcome == "completed" and (size_report is None or size_report.target_reached): + verbose_proxy_logger.info(message, outcome, elapsed_seconds, per_table, storage) return - verbose_proxy_logger.warning(message, outcome, elapsed_seconds, per_table) + verbose_proxy_logger.warning(message, outcome, elapsed_seconds, per_table, storage) async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ - Main cleanup function. Deletes old spend logs in batches. + The scheduled retention run: every age rule, plus the size rule when one cron drives both. If pod_lock_manager is available, ensures only one pod runs cleanup. If no pod_lock_manager, runs cleanup without distributed locking. """ + await self._run(prisma_client, age_rules=True, size_rule=self._size_pass_runs_with_age_pass()) + + async def cleanup_spend_logs_by_size(self, prisma_client: PrismaClient) -> None: + """The size-check run scheduled on maximum_spend_logs_size_check_interval: the size rule only.""" + await self._run(prisma_client, age_rules=False, size_rule=True) + + async def _run(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None: lock_acquired = False run_started_at: Final = time.monotonic() progress: Final = _RunProgress() @@ -992,19 +1317,23 @@ class SpendLogCleanup: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() - delete_spend_logs: Final = self._should_delete_spend_logs() - autorouter_retention_seconds: Final = self._retention_seconds_for( - "maximum_autorouter_session_retention_period" + delete_spend_logs: Final = age_rules and self._should_delete_spend_logs() + active_size_rule: Final = self._size_rule() if size_rule else None + autorouter_retention_seconds: Final = ( + self._retention_seconds_for("maximum_autorouter_session_retention_period") if age_rules else None ) - health_check_retention_seconds: Final = self._retention_seconds_for("maximum_health_check_retention_period") - daily_tag_spend_retention_seconds: Final = self._retention_seconds_for( - "maximum_daily_tag_spend_retention_period" + health_check_retention_seconds: Final = ( + self._retention_seconds_for("maximum_health_check_retention_period") if age_rules else None + ) + daily_tag_spend_retention_seconds: Final = ( + self._retention_seconds_for("maximum_daily_tag_spend_retention_period") if age_rules else None ) if ( not delete_spend_logs and autorouter_retention_seconds is None and health_check_retention_seconds is None and daily_tag_spend_retention_seconds is None + and active_size_rule is None ): SpendLogCleanupMetrics.record_run("skipped_disabled") return @@ -1032,11 +1361,13 @@ class SpendLogCleanup: return deadline: Final = time.monotonic() + self.run_budget_seconds + size_group: Final = int(active_size_rule is not None) configured_group_count: Final = ( int(delete_spend_logs and self.retention_seconds is not None) + int(autorouter_retention_seconds is not None) + int(health_check_retention_seconds is not None) + int(daily_tag_spend_retention_seconds is not None) + + size_group ) spend_log_results: Final = ( @@ -1051,9 +1382,12 @@ class SpendLogCleanup: int(autorouter_retention_seconds is not None) + int(health_check_retention_seconds is not None) + int(daily_tag_spend_retention_seconds is not None) + + size_group ) - remaining_groups_after_sessions: Final = int(health_check_retention_seconds is not None) + int( - daily_tag_spend_retention_seconds is not None + remaining_groups_after_sessions: Final = ( + int(health_check_retention_seconds is not None) + + int(daily_tag_spend_retention_seconds is not None) + + size_group ) session_results: Final = ( await self._clean_session_rollup( @@ -1074,15 +1408,27 @@ class SpendLogCleanup: else () ) daily_tag_spend_results: Final = ( - await self._clean_daily_tag_spend(prisma_client, daily_tag_spend_retention_seconds, deadline) + await self._clean_daily_tag_spend( + prisma_client, + daily_tag_spend_retention_seconds, + self._group_deadline(deadline, 1 + size_group), + ) if daily_tag_spend_retention_seconds is not None else () ) + size_report: Final = ( + await self._clean_spend_logs_by_size(prisma_client, active_size_rule, deadline) + if active_size_rule is not None + else None + ) + size_results: Final = size_report.results if size_report is not None else () - results: Final = spend_log_results + session_results + health_check_results + daily_tag_spend_results + results: Final = ( + spend_log_results + session_results + health_check_results + daily_tag_spend_results + size_results + ) outcome: Final = self._run_outcome(results) SpendLogCleanupMetrics.record_run(outcome) - self._log_run_summary(outcome, results, time.monotonic() - run_started_at) + self._log_run_summary(outcome, results, time.monotonic() - run_started_at, size_report) except asyncio.CancelledError: verbose_proxy_logger.error( diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py index 340aeab938c..048fa8127a4 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py @@ -32,6 +32,9 @@ RunOutcome: TypeAlias = Literal[ _BATCH_DURATION_BUCKETS: Final = (0.005, 0.025, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0) _TABLE_LABEL: Final = ("table",) _OUTCOME_LABEL: Final = ("outcome",) +_MEASURE_LABEL: Final = ("measure",) + +StorageMeasure: TypeAlias = Literal["allocated", "estimated_live", "above_target"] class SpendLogCleanupMetrics: @@ -49,6 +52,7 @@ class SpendLogCleanupMetrics: rows_remaining: "PrometheusGauge | None" = None batch_failures: "PrometheusCounter | None" = None runs: "PrometheusCounter | None" = None + storage_bytes: "PrometheusGauge | None" = None @classmethod def _ensure_initialized(cls) -> None: @@ -90,6 +94,15 @@ class SpendLogCleanupMetrics: "Retention cleanup runs, labelled by why the run ended", labelnames=_OUTCOME_LABEL, ) + cls.storage_bytes = Gauge( + "litellm_spend_log_storage_bytes", + "Disk held by LiteLLM_SpendLogs plus LiteLLM_SpendLogToolIndex as last measured by the cleanup " + "job: allocated is pg_total_relation_size (heap, indexes, TOAST), estimated_live scales it by the " + "live share of rows from pg_stat_user_tables, above_target is estimated_live minus the size " + "cleanup target (0 once the target is reached)", + labelnames=_MEASURE_LABEL, + multiprocess_mode="livemax", + ) except Exception as e: # noqa: BLE001 - a metrics problem must never fail the cleanup run # Covers the extra being absent, a duplicate registration (repeated # imports under a test runner), and registry misconfiguration alike. @@ -120,3 +133,9 @@ class SpendLogCleanupMetrics: cls._ensure_initialized() if cls.runs is not None: cls.runs.labels(outcome=outcome).inc() + + @classmethod + def set_storage_bytes(cls, measure: StorageMeasure, value: int) -> None: + cls._ensure_initialized() + if cls.storage_bytes is not None: + cls.storage_bytes.labels(measure=measure).set(value) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 8c6d5ab69ab..b18b4c75ec0 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,7 +1,7 @@ import asyncio import traceback from collections.abc import Callable, Mapping, Sequence -from datetime import datetime +from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, cast @@ -35,6 +35,7 @@ from litellm.proxy.db.db_spend_update_writer import ( debitable_model_access_groups, get_llm_router, ) +from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( @@ -491,8 +492,21 @@ class _ProxyDBLogger(CustomLogger): key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) + raw_cost: Final = _str_keyed_fields( + sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload + ).get("response_cost") + priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None response_cost: Final = ( - sl_object.get("response_cost", None) if sl_object is not None else kwargs.get("response_cost", None) + await apply_search_free_tier( + kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload + sl_object=sl_object, # pyright: ignore[reportUnknownArgumentType] # payload carries untyped lists + response_cost=priced_cost, + completion_response=completion_response, + counter=proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + now=datetime.now(timezone.utc), + ) + if litellm.apply_search_free_tier and sl_object is not None and priced_cost is not None + else priced_cost ) tags: Final = _get_request_tags_for_cost_tracking( sl_object=sl_object, diff --git a/litellm/proxy/hooks/search_free_tier.py b/litellm/proxy/hooks/search_free_tier.py new file mode 100644 index 00000000000..4ef300dfbac --- /dev/null +++ b/litellm/proxy/hooks/search_free_tier.py @@ -0,0 +1,123 @@ +from collections.abc import MutableMapping +from datetime import datetime +from types import MappingProxyType +from typing import Final, Protocol + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time +from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage +from litellm.llms.gemini.cost_calculator import cost_per_web_search_request +from litellm.types.utils import CostBreakdown, ModelInfo, SearchFreeRequests, StandardLoggingPayload, Usage + +_PERIOD_DURATION: Final = MappingProxyType({"day": "1d", "month": "1mo"}) +_PERIOD_STAMP: Final = MappingProxyType({"day": "%Y-%m-%d", "month": "%Y-%m"}) +_TTL_GRACE_SECONDS: Final = 3600 +_COUNTED_KWARG: Final = "search_free_tier_counted" +_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown) + + +class SharedCounter(Protocol): + async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: ... + + +class _UsageCarrier(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + usage: Usage | None = None + + +def freed_units(allowance: int, units: int, total_after: int) -> int: + """Units of this request that landed under the allowance, given the shared total after they were added.""" + return max(0, min(units, allowance - (total_after - units))) + + +def search_free_tier_key(allowance: SearchFreeRequests, now: datetime) -> str: + return f"litellm:search_free_tier:{allowance['shared_group']}:{now.strftime(_PERIOD_STAMP[allowance['period']])}" + + +def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int: + reset: Final = get_next_standardized_reset_time(_PERIOD_DURATION[allowance["period"]], now, "UTC") + return int((reset - now).total_seconds()) + _TTL_GRACE_SECONDS + + +def _usage_of(completion_response: object) -> Usage | None: + try: + return _UsageCarrier.model_validate(completion_response, from_attributes=True).usage + except ValidationError: + return None + + +def _grounding_queries(usage: Usage) -> int: + details: Final = usage.prompt_tokens_details + from_details: Final = getattr(details, "web_search_requests", None) if details is not None else None + return from_details or get_web_search_requests_from_usage(usage) or 0 + + +def _billable_units(usage: Usage, model_info: ModelInfo) -> int: + queries: Final = _grounding_queries(usage) + if queries == 0: + return 0 + return queries if model_info.get("web_search_billing_unit") == "per_query" else 1 + + +def _record_freed_cost( + kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it + sl_object: StandardLoggingPayload, + original_cost: float, + freed_cost: float, +) -> float: + reduced_cost: Final = original_cost - freed_cost + breakdown: Final = sl_object["cost_breakdown"] or CostBreakdown(total_cost=original_cost) + sl_object["cost_breakdown"] = _COST_BREAKDOWN.validate_python( # rebind-ok: the spend row reads this payload + { + **breakdown, + "total_cost": reduced_cost, + "original_cost": breakdown.get("original_cost", original_cost), + "discount_amount": breakdown.get("discount_amount", 0.0) + freed_cost, + } + ) + sl_object["response_cost"] = reduced_cost # rebind-ok: the spend row reads this payload + kwargs["response_cost"] = reduced_cost # rebind-ok: the spend row reads kwargs["response_cost"] + return reduced_cost + + +@with_service_target("search_free_tier") +async def apply_search_free_tier( + kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it + sl_object: StandardLoggingPayload, + response_cost: float, + completion_response: object, + counter: SharedCounter | None, + now: datetime, +) -> float: + """Zero the grounding queries that landed under the model's shared free allowance, everything else is unchanged.""" + model_info: Final = sl_object["model_map_information"]["model_map_value"] + allowance: Final = model_info.get("search_free_requests") if model_info is not None else None + usage: Final = _usage_of(completion_response) + if model_info is None or allowance is None or usage is None or sl_object["cache_hit"]: + return response_cost + units: Final = _billable_units(usage, model_info) + if units == 0 or kwargs.get(_COUNTED_KWARG): + return response_cost + if counter is None: + verbose_proxy_logger.warning( + "apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price" + ) + return response_cost + key: Final = search_free_tier_key(allowance, now) + try: + total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now)) + except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row + verbose_proxy_logger.warning( + "search free tier counter %s is unavailable, grounding is billed at full price", key + ) + return response_cost + kwargs[_COUNTED_KWARG] = True # rebind-ok: the proxy runs this callback twice per request, count it once + freed: Final = freed_units(allowance["count"], units, int(total_after)) + if freed == 0: + return response_cost + freed_cost: Final = min(response_cost, cost_per_web_search_request(usage, model_info) * freed / units) + return _record_freed_cost(kwargs, sl_object, response_cost, freed_cost) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0af6036efa7..c6d3e12a378 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -340,6 +340,7 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, SPEND_CAPTURE_RATE_CHECK_JOB_ID, + SPEND_LOG_SIZE_CHECK_INTERVAL, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) @@ -5522,20 +5523,44 @@ def _current_general_settings() -> Mapping[str, object]: return general_settings -_CLEANUP_SCHEDULE_KEYS: Final = ( +_RETENTION_RULE_KEYS: Final = ( "maximum_spend_logs_retention_period", "maximum_autorouter_session_retention_period", "maximum_health_check_retention_period", "maximum_daily_tag_spend_retention_period", +) +_CLEANUP_SCHEDULE_KEYS: Final = ( + *_RETENTION_RULE_KEYS, "maximum_spend_logs_cleanup_cron", "maximum_spend_logs_retention_interval", + "maximum_spend_logs_size", + "maximum_spend_logs_cleanup_target_size", + "maximum_spend_logs_size_check_interval", ) +_SPEND_LOG_CLEANUP_JOB_ID: Final = "spend_log_cleanup_job" +_SPEND_LOG_SIZE_CHECK_JOB_ID: Final = "spend_log_size_check_job" def _cleanup_schedule_of(settings: Mapping[str, object]) -> tuple[object, ...]: return tuple(settings.get(key) for key in _CLEANUP_SCHEDULE_KEYS) +def _wants_retention_job(settings: Mapping[str, object]) -> bool: + """The retention job runs the age rules, and the size rule too when one cron drives both.""" + return any(settings.get(key) is not None for key in _RETENTION_RULE_KEYS) or ( + settings.get("maximum_spend_logs_size") is not None and bool(settings.get("maximum_spend_logs_cleanup_cron")) + ) + + +def _wants_size_check_job(settings: Mapping[str, object]) -> bool: + """The size rule gets its own job on maximum_spend_logs_size_check_interval unless a cleanup cron is set.""" + return settings.get("maximum_spend_logs_size") is not None and not settings.get("maximum_spend_logs_cleanup_cron") + + +def _wants_cleanup_jobs(settings: Mapping[str, object]) -> bool: + return _wants_retention_job(settings) or _wants_size_check_job(settings) + + @lru_cache(maxsize=4096) def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None: verbose_proxy_logger.warning( @@ -8052,39 +8077,73 @@ class ProxyConfig: if scheduler is None: return - wants_job: Final = any( - general_settings.get(key) is not None - for key in ( - "maximum_spend_logs_retention_period", - "maximum_autorouter_session_retention_period", - "maximum_health_check_retention_period", - "maximum_daily_tag_spend_retention_period", - ) - ) - if not wants_job: - if scheduler.get_job("spend_log_cleanup_job") is not None: - scheduler.remove_job("spend_log_cleanup_job") - verbose_proxy_logger.info("Removed existing spend log cleanup job") - return - - trigger: Final = self._spend_log_cleanup_trigger() - if trigger is None: - return from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SpendLogCleanup, ) - scheduler.add_job( - SpendLogCleanup().cleanup_old_spend_logs, - trigger, + if not _wants_retention_job(general_settings): + if scheduler.get_job(_SPEND_LOG_CLEANUP_JOB_ID) is not None: + scheduler.remove_job(_SPEND_LOG_CLEANUP_JOB_ID) + verbose_proxy_logger.info("Removed existing spend log cleanup job") + else: + trigger: Final = self._spend_log_cleanup_trigger() + if trigger is not None: + scheduler.add_job( + SpendLogCleanup().cleanup_old_spend_logs, + trigger, + args=[prisma_client], + id=_SPEND_LOG_CLEANUP_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info("Spend log cleanup rescheduled with trigger: %s", trigger) + + if not _wants_size_check_job(general_settings): # pyright: ignore[reportUnknownArgumentType] # general_settings is an untyped module dict + if scheduler.get_job(_SPEND_LOG_SIZE_CHECK_JOB_ID) is not None: # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + scheduler.remove_job(_SPEND_LOG_SIZE_CHECK_JOB_ID) # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + verbose_proxy_logger.info("Removed existing spend log size check job") + return + size_trigger: Final = self._spend_log_size_check_trigger() + if size_trigger is None: + return + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + SpendLogCleanup().cleanup_spend_logs_by_size, + size_trigger, args=[prisma_client], - id="spend_log_cleanup_job", + id=_SPEND_LOG_SIZE_CHECK_JOB_ID, replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info("Spend log cleanup rescheduled with trigger: %s", trigger) + verbose_proxy_logger.info("Spend log size check rescheduled with trigger: %s", size_trigger) - def _spend_log_cleanup_trigger(self) -> BaseTrigger | None: + def _spend_log_size_check_trigger(self) -> object | None: + return self._staggered_interval_trigger( + _SPEND_LOG_SIZE_CHECK_JOB_ID, + "maximum_spend_logs_size_check_interval", + general_settings.get("maximum_spend_logs_size_check_interval", SPEND_LOG_SIZE_CHECK_INTERVAL), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # general_settings is an untyped module dict + ) + + @staticmethod + def _staggered_interval_trigger(job_id: str, setting_name: str, interval: object) -> object | None: + if not isinstance(interval, str): + verbose_proxy_logger.error("Invalid %s value: %r", setting_name, interval) + return None + # this runs against a started scheduler, which the startup stagger sweep + # cannot reach, so the offset is applied here or the job reconverges across + # replicas the first time an admin edits the settings + try: + interval_seconds: Final = duration_in_seconds(interval) + return stagger_trigger( + job_id=job_id, + trigger=IntervalTrigger(seconds=interval_seconds), + period_seconds=interval_seconds, + settings=parse_stagger_settings(general_settings), # pyright: ignore[reportUnknownArgumentType] # general_settings is an untyped module dict + ) + except (ValueError, OverflowError): + verbose_proxy_logger.error("Invalid %s value: %r", setting_name, interval) + return None + + def _spend_log_cleanup_trigger(self) -> object | None: cleanup_cron: Final[object] = general_settings.get("maximum_spend_logs_cleanup_cron") if cleanup_cron: from apscheduler.triggers.cron import CronTrigger @@ -8095,24 +8154,11 @@ class ProxyConfig: verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron) return None return cron_trigger - retention_interval: Final[object] = general_settings.get("maximum_spend_logs_retention_interval", "1d") - if not isinstance(retention_interval, str): - verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value: %r", retention_interval) - return None - # this runs against a started scheduler, which the startup stagger sweep - # cannot reach, so the offset is applied here or the job reconverges across - # replicas the first time an admin edits the retention settings - try: - interval_seconds: Final = duration_in_seconds(retention_interval) - return stagger_trigger( - job_id="spend_log_cleanup_job", - trigger=IntervalTrigger(seconds=interval_seconds), - period_seconds=interval_seconds, - settings=parse_stagger_settings(general_settings), - ) - except (ValueError, OverflowError): - verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value: %r", retention_interval) - return None + return self._staggered_interval_trigger( + _SPEND_LOG_CLEANUP_JOB_ID, + "maximum_spend_logs_retention_interval", + general_settings.get("maximum_spend_logs_retention_interval", "1d"), + ) async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | None) -> None: global general_settings @@ -8269,8 +8315,11 @@ class ProxyConfig: if scheduler is not None and scheduler.state == STATE_STOPPED: return schedule: Final = self._resolved_cleanup_schedule() - wants_job: Final = any(value is not None for value in schedule[:4]) - has_job: Final = scheduler is not None and scheduler.get_job("spend_log_cleanup_job") is not None + wants_job: Final = _wants_cleanup_jobs(self.settings) + has_job: Final = scheduler is not None and ( + scheduler.get_job(_SPEND_LOG_CLEANUP_JOB_ID) is not None + or scheduler.get_job(_SPEND_LOG_SIZE_CHECK_JOB_ID) is not None # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + ) baseline: Final = ( self._last_cleanup_schedule_attempt if has_job and self._last_cleanup_schedule_attempt is not None @@ -11417,12 +11466,29 @@ class ProxyStartupEvent: ### SPEND LOG CLEANUP ### cleanup_settings: Final = _current_general_settings() - if ( - cleanup_settings.get("maximum_spend_logs_retention_period") is not None - or cleanup_settings.get("maximum_autorouter_session_retention_period") is not None - or cleanup_settings.get("maximum_health_check_retention_period") is not None - or cleanup_settings.get("maximum_daily_tag_spend_retention_period") is not None - ): + if _wants_size_check_job(cleanup_settings): + size_check_interval: Final = cleanup_settings.get( + "maximum_spend_logs_size_check_interval", SPEND_LOG_SIZE_CHECK_INTERVAL + ) + try: + if not isinstance(size_check_interval, str): + raise TypeError(size_check_interval) + size_check_seconds: Final = duration_in_seconds(size_check_interval) + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + SpendLogCleanup().cleanup_spend_logs_by_size, + "interval", + seconds=size_check_seconds + random.randint(0, 60), + args=[prisma_client], + id=_SPEND_LOG_SIZE_CHECK_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info("Spend log size check scheduled every %s", size_check_interval) + except (TypeError, ValueError, OverflowError): + verbose_proxy_logger.error( + "Invalid maximum_spend_logs_size_check_interval value: %r", size_check_interval + ) + if _wants_retention_job(cleanup_settings): spend_log_cleanup: Final = SpendLogCleanup() cleanup_cron: Final = cleanup_settings.get("maximum_spend_logs_cleanup_cron") @@ -11437,7 +11503,7 @@ class ProxyStartupEvent: spend_log_cleanup.cleanup_old_spend_logs, cron_trigger, args=[prisma_client], - id="spend_log_cleanup_job", + id=_SPEND_LOG_CLEANUP_JOB_ID, replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) @@ -11456,7 +11522,7 @@ class ProxyStartupEvent: "interval", seconds=interval_seconds + random.randint(0, 60), args=[prisma_client], - id="spend_log_cleanup_job", + id=_SPEND_LOG_CLEANUP_JOB_ID, replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) @@ -11464,6 +11530,7 @@ class ProxyStartupEvent: verbose_proxy_logger.error( "Invalid maximum_spend_logs_retention_interval value: %r", retention_interval ) + if _wants_cleanup_jobs(cleanup_settings): proxy_config.record_cleanup_schedule_attempt(cleanup_settings) ### CHECK BATCH COST ### if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2721a0c1cc1..e268cc90ffd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -245,6 +245,12 @@ class SearchContextCostPerQuery(TypedDict, total=False): search_context_size_high: float +class SearchFreeRequests(TypedDict): + count: ReadOnly[int] + period: ReadOnly[Literal["day", "month"]] + shared_group: ReadOnly[str] + + class AgenticLoopParams(TypedDict, total=False): """ Parameters passed to agentic loop hooks (e.g., WebSearch interception). @@ -436,6 +442,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None ) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x) + search_free_requests: ReadOnly[SearchFreeRequests | None] google_maps_grounding_cost_per_query: ReadOnly[float | None] citation_cost_per_token: float | None # Cost per citation token for Perplexity tiered_pricing: list[dict[str, Any]] | None # Tiered pricing structure for models like Dashscope diff --git a/litellm/utils.py b/litellm/utils.py index a6108a15bb3..6b16fc398bd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6878,6 +6878,7 @@ def get_model_info_helper( supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), web_search_billing_unit=_model_info.get("web_search_billing_unit", None), + search_free_requests=_model_info.get("search_free_requests", None), google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None), tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 903210c57e5..503877a9d4d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28948,6 +28948,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_token_batches": 1.5e-07, @@ -29062,6 +29063,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3-pro-image-preview": { @@ -29114,6 +29116,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { @@ -29164,6 +29167,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image-preview": { @@ -29207,6 +29211,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite-image": { @@ -29302,6 +29307,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29363,6 +29369,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -29422,6 +29429,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29508,6 +29516,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-08, "input_cost_per_token_batches": 5e-08, @@ -29562,6 +29571,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29608,6 +29618,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29795,6 +29806,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, @@ -29832,6 +29844,7 @@ "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29912,6 +29925,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29968,6 +29982,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3-flash-preview": { @@ -29992,6 +30007,7 @@ "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30074,6 +30090,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -30138,6 +30155,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30198,6 +30216,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30224,6 +30243,7 @@ "prompt_cache_min_tokens": 4096.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30369,6 +30389,7 @@ "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30455,6 +30476,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30782,6 +30804,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "supports_reasoning": false }, @@ -30826,6 +30849,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-lite-preview": { @@ -30887,6 +30911,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -31036,6 +31061,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_audio_token_priority": 1.8e-06, @@ -31148,6 +31174,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "supports_reasoning": true }, @@ -31212,6 +31239,7 @@ "output_cost_per_token_batches": 1.5e-06, "rpm": 1000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31362,6 +31390,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 1.5e-07, "input_cost_per_audio_token_priority": 5.4e-07, @@ -31567,6 +31596,7 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -31664,6 +31694,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07, @@ -31725,6 +31756,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -31781,6 +31813,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 5e-07, @@ -31844,6 +31877,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -31909,6 +31943,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -31970,6 +32005,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32031,6 +32067,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32126,6 +32163,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32189,6 +32227,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32217,6 +32256,7 @@ "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -32331,6 +32371,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -32395,6 +32436,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32455,6 +32497,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32480,6 +32523,7 @@ "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -52449,6 +52493,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -52511,6 +52556,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -52571,6 +52617,7 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index 925fbaf6f01..72da2bb48ee 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -105,3 +105,166 @@ async def test_delete_batch_survives_witnessed_lock_past_transaction_default(cap await database.disconnect() finally: setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema))) + + +async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does_not_refire_after_vacuum( + caplog: pytest.LogCaptureFixture, +) -> None: + """ + Deleting rows never shrinks the table file, so the size pass must stop on live bytes, not on + pg_total_relation_size, taking the oldest rows first and keeping the tool index on the same cutoff. Once + autovacuum has turned the deleted rows into free space the catalog no longer shows them as dead, and the + files still measure the same; the next check must read what the last run freed from its saved count and + leave the table alone instead of calling the whole file live again. + """ + caplog.set_level(logging.INFO, logger="LiteLLM Proxy") + schema: Final = f"integration_{uuid.uuid4().hex}" + url: Final = os.environ["DATABASE_URL"] + parsed: Final = urlsplit(url) + scoped_url: Final = urlunsplit( + parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})) + ) + logs: Final = sql.Identifier(schema, "LiteLLM_SpendLogs") + index: Final = sql.Identifier(schema, "LiteLLM_SpendLogToolIndex") + config: Final = sql.Identifier(schema, "LiteLLM_Config") + seeded: Final = 6000 + now: Final = datetime.now(timezone.utc).replace(tzinfo=None) + stats_sql: Final = ( + "SELECT sum(n_live_tup)::int AS live, sum(n_dead_tup)::int AS dead, sum(n_tup_ins)::int AS inserted " + "FROM pg_stat_user_tables " + "WHERE schemaname = %s AND relname IN ('LiteLLM_SpendLogs', 'LiteLLM_SpendLogToolIndex')" + ) + with psycopg.connect(url, autocommit=True) as setup: + setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) + try: + setup.execute( + sql.SQL( + 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL, payload text, ' + 'PRIMARY KEY (request_id, "startTime"))' + ).format(logs) + ) + setup.execute( + sql.SQL( + "CREATE TABLE {} (request_id text, tool_name text, start_time timestamp NOT NULL, " + "PRIMARY KEY (request_id, tool_name))" + ).format(index) + ) + setup.execute(sql.SQL("CREATE TABLE {} (param_name text PRIMARY KEY, param_value jsonb)").format(config)) + setup.execute( + sql.SQL( + "INSERT INTO {} SELECT 'req-' || lpad(i::text, 6, '0'), %s - (%s - i) * interval '1 minute', " + "repeat('x', 1000) FROM generate_series(1, %s) AS i" + ).format(logs), + (now, seeded, seeded), + ) + setup.execute( + sql.SQL( + "INSERT INTO {} SELECT request_id, 'tool', \"startTime\" FROM {} WHERE right(request_id, 1) = '0'" + ).format(index, logs) + ) + index_rows: Final = seeded // 10 + inserts_deadline: Final = time.monotonic() + 10 + while True: + # a backend only ships its pending row counters when it runs another command, so give the + # seeding connection one; the size rule rolls its count forward from n_tup_ins, which must + # show the seed inserts before ANALYZE sets the live estimate or the two add up to double + setup.execute("SELECT 1") + (inserts,) = read_rows(stats_sql, (schema,)) + if inserts["inserted"] == seeded + index_rows: + break + assert time.monotonic() < inserts_deadline, f"row statistics never saw the seed inserts: {inserts}" + await asyncio.sleep(0.1) + setup.execute(sql.SQL("ANALYZE {}").format(logs)) + setup.execute(sql.SQL("ANALYZE {}").format(index)) + stats_deadline: Final = time.monotonic() + 10 + while True: + setup.execute("SELECT 1") + (stats,) = read_rows(stats_sql, (schema,)) + if stats["live"] == seeded + index_rows: + break + assert time.monotonic() < stats_deadline, f"row statistics never caught up with ANALYZE: {stats}" + await asyncio.sleep(0.1) + (allocated,) = read_rows( + "SELECT (pg_total_relation_size(%s::regclass) + pg_total_relation_size(%s::regclass))::bigint " + "AS allocated", + (logs.as_string(setup), index.as_string(setup)), + ) + allocated_before: Final = int(str(allocated["allocated"])) + assert allocated_before > 0 + + cleanup: Final = SpendLogCleanup( + general_settings={ + # trigger below what is there, clean down to half of it + "maximum_spend_logs_size": str(allocated_before * 3 // 4), + "maximum_spend_logs_cleanup_target_size": str(allocated_before // 2), + "maximum_spend_logs_cleanup_batch_size": 100, + } + ) + cleanup.pod_lock_manager = None + + async def size_check() -> None: + database: Final = Prisma(datasource={"url": scoped_url}) + await database.connect() + try: + await cleanup.cleanup_spend_logs_by_size(PartitionConnection(database)) + finally: + await database.disconnect() + + await size_check() + + remaining: Final = read_rows( + f'SELECT request_id FROM {logs.as_string(setup)} ORDER BY "startTime"', + (), + ) + # half the bytes is half the rows: the oldest 3000 spend logs and their 300 index rows go, counted + # from the deletes themselves, so the stats collector's lag cannot make the run overshoot. A page the + # files grew by between this measurement and the run's probe costs at most one more batch + assert 2900 <= len(remaining) <= 3000, f"{len(remaining)} spend logs remain, expected 3000 less one batch" + assert remaining[0]["request_id"] == f"req-{seeded - len(remaining) + 1:06d}" + assert remaining[-1]["request_id"] == f"req-{seeded:06d}" + (orphans,) = read_rows( + f"SELECT count(*)::int AS n FROM {index.as_string(setup)} i " + f"WHERE NOT EXISTS (SELECT 1 FROM {logs.as_string(setup)} l WHERE l.request_id = i.request_id)", + (), + ) + assert orphans["n"] == 0 + (kept_index,) = read_rows(f"SELECT count(*)::int AS n FROM {index.as_string(setup)}", ()) + assert kept_index["n"] == len(remaining) // 10 + summary: Final = next( + record.getMessage() + for record in caplog.records + if "Spend log cleanup run finished" in record.getMessage() + ) + assert "target_reached=true" in summary + assert f"allocated_before={allocated_before}" in summary + # the file did not shrink: that is the whole reason the reclaimable count exists + assert f"allocated_after={allocated_before}" in summary + assert "VACUUM FULL or pg_repack" in summary + + deletes_deadline: Final = time.monotonic() + 10 + while True: + (after_deletes,) = read_rows(stats_sql, (schema,)) + if after_deletes["dead"] == 3300: + break + assert time.monotonic() < deletes_deadline, f"row statistics never saw the deletes: {after_deletes}" + await asyncio.sleep(0.1) + setup.execute(sql.SQL("VACUUM (ANALYZE) {}").format(logs)) + setup.execute(sql.SQL("VACUUM (ANALYZE) {}").format(index)) + vacuum_deadline: Final = time.monotonic() + 10 + while True: + (after_vacuum,) = read_rows(stats_sql, (schema,)) + if after_vacuum["dead"] == 0: + break + assert time.monotonic() < vacuum_deadline, f"row statistics never caught up with VACUUM: {after_vacuum}" + await asyncio.sleep(0.1) + caplog.clear() + + await size_check() + + (still_there,) = read_rows(f"SELECT count(*)::int AS n FROM {logs.as_string(setup)}", ()) + assert still_there["n"] == len(remaining), ( + "the second check deleted rows the first one had already brought under the limit" + ) + assert any("is under maximum_spend_logs_size" in record.getMessage() for record in caplog.records) + finally: + setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema))) diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 5dadeff6d7c..591205b4602 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -81,7 +81,7 @@ async def test_handle_count_tokens_request_looks_up_its_region_off_the_event_loo ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert result == {"input_tokens": 7} diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py index 2d21a3ddef5..a6b309c1b83 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -124,6 +124,6 @@ async def test_region_and_endpoint_lookups_leave_the_event_loop_free(monkeypatch ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert result == {"input_tokens": 2177} assert str(mantle.only_request().url) == MANTLE_COUNT_URL diff --git a/tests/unit/llms/bedrock/event_loop_probe.py b/tests/unit/llms/bedrock/event_loop_probe.py index cf6cd5975ed..8bf856582b5 100644 --- a/tests/unit/llms/bedrock/event_loop_probe.py +++ b/tests/unit/llms/bedrock/event_loop_probe.py @@ -16,7 +16,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 REFRESH_START_TIMEOUT_SECONDS: Final = 10.0 -LOOP_SERVE_TIMEOUT_SECONDS: Final = 2.0 class EventLoopProbe: @@ -63,26 +62,21 @@ class EventLoopProbe: class LoopProbingVault: - """A Vault holding none of its hosted secrets, whose read waits for the event loop to run a callback. + """A Vault with no hosted secrets that records reads on the event-loop thread.""" - A read made on the loop thread can never see that callback, so `stalled_the_loop` is set only then. - """ - - def __init__(self, loop: asyncio.AbstractEventLoop) -> None: - self._loop: Final = loop + def __init__(self) -> None: + self._loop_thread_id: Final = threading.get_ident() self.read: Final = threading.Event() - self.stalled_the_loop: Final = threading.Event() + self.read_on_loop: Final = threading.Event() def sync_read_secret(self, secret_name: str) -> None: self.read.set() - served: Final = threading.Event() - self._loop.call_soon_threadsafe(served.set) - if not served.wait(timeout=LOOP_SERVE_TIMEOUT_SECONDS): - self.stalled_the_loop.set() + if threading.get_ident() == self._loop_thread_id: + self.read_on_loop.set() def host_in_loop_probing_vault(monkeypatch: pytest.MonkeyPatch, *secret_names: str) -> LoopProbingVault: - vault: Final = LoopProbingVault(asyncio.get_running_loop()) + vault: Final = LoopProbingVault() monkeypatch.setattr(litellm, "secret_manager_client", vault) monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) monkeypatch.setattr( diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index ccea4ba44d7..2ac31b6887d 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -643,7 +643,7 @@ class TestBedrockRealtimeSessionLifecycle: ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert stub_aws_sdk_client["config_kwargs"]["region"] == "us-west-2" assert stub_aws_sdk_client["config_kwargs"]["endpoint_uri"] == "https://bedrock-runtime.us-west-2.amazonaws.com" diff --git a/tests/unit/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py index e43184df50e..2353a45139b 100644 --- a/tests/unit/llms/bedrock/test_claude_platform_provider.py +++ b/tests/unit/llms/bedrock/test_claude_platform_provider.py @@ -705,7 +705,7 @@ async def test_anthropic_messages_claude_platform_workspace_and_region_lookups_l ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert response["content"][0]["text"] == "ok" _assert_gateway_accepted(requests[0], sigv4_without_region) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 205c610a413..63cc45b29df 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1392,7 +1392,7 @@ async def test_bedrock_secret_lookups_leave_the_event_loop_free( _ = await call() assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert ".us-west-2." in upstream.calls.last.request.url.host diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 75d713028d4..69b3e113b8f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -8103,5 +8103,5 @@ async def test_guardrail_secret_lookups_leave_the_event_loop_free( ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert upstream.calls.last.request.url.host == "bedrock-runtime.us-west-2.amazonaws.com" diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 7e805a3ba07..f26e7868dde 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -29,7 +29,7 @@ from litellm.proxy.utils import ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, Usage +from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.mark.asyncio @@ -3279,3 +3279,92 @@ async def test_group_soft_alert_delivery_failure_does_not_suppress_other_groups( (call.kwargs["user_info"].key_alias, call.kwargs["user_info"].spend) for call in delivery.budget_alerts.await_args_list ] == [("broken-delivery", 350), ("healthy-delivery", 350)] + + +class _SharedCounter: + def __init__(self) -> None: + self.totals: dict[str, float] = {} + + async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: + self.totals[key] = self.totals.get(key, 0) + value + return self.totals[key] + + +def _grounded_gemini_kwargs(response_cost: float) -> dict: + return { + "model": "vertex_ai/gemini-test", + "call_type": "acompletion", + "response_cost": response_cost, + "litellm_params": {"metadata": {"user_api_key": "sk-test", "user_api_key_user_id": "user-1"}}, + "standard_logging_object": { + "response_cost": response_cost, + "request_tags": None, + "cache_hit": None, + "cost_breakdown": { + "input_cost": 0.0, + "output_cost": 0.0, + "tool_usage_cost": 0.5, + "total_cost": response_cost, + }, + "model_map_information": { + "model_map_key": "vertex_ai/gemini-test", + "model_map_value": { + "litellm_provider": "vertex_ai", + "mode": "chat", + "web_search_billing_unit": "per_query", + "search_context_cost_per_query": {"search_context_size_medium": 0.1}, + "search_free_requests": {"count": 100, "period": "month", "shared_group": "test"}, + }, + }, + }, + "stream": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("opted_in", [True, False]) +async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in): + monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False) + counter = _SharedCounter() + kwargs = _grounded_gemini_kwargs(response_cost=1.0) + response = ModelResponse( + usage=Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5), + ) + ) + + with ( + patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ), + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.internal_usage_cache.dual_cache.redis_cache = counter + + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, completion_response=response, start_time=datetime.now(), end_time=datetime.now() + ) + await asyncio.sleep(0) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + + breakdown = kwargs["standard_logging_object"]["cost_breakdown"] + if opted_in: + assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1) + assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1) + assert sum(counter.totals.values()) == 5 + else: + assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown + assert counter.totals == {} diff --git a/tests/unit/proxy/hooks/test_search_free_tier.py b/tests/unit/proxy/hooks/test_search_free_tier.py new file mode 100644 index 00000000000..3dc89c0dc72 --- /dev/null +++ b/tests/unit/proxy/hooks/test_search_free_tier.py @@ -0,0 +1,194 @@ +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest + +from litellm.proxy.hooks.search_free_tier import ( + apply_search_free_tier, + freed_units, + search_free_tier_key, + search_free_tier_ttl, +) +from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + +NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc) +RATE: Final = 0.01 + + +class _Counter: + def __init__(self) -> None: + self.totals: dict[str, float] = {} + self.ttls: dict[str, int | None] = {} + + async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: + self.totals[key] = self.totals.get(key, 0) + value + self.ttls[key] = ttl + return self.totals[key] + + +class _BrokenCounter: + async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: + raise ConnectionError("redis is down") + + +def _model_info(billing_unit: str, count: int, period: str = "month") -> dict: + return { + "litellm_provider": "vertex_ai", + "mode": "chat", + "web_search_billing_unit": billing_unit, + "search_context_cost_per_query": {"search_context_size_medium": RATE}, + "search_free_requests": {"count": count, "period": period, "shared_group": "test-group"}, + } + + +def _response(queries: int | None) -> ModelResponse: + details = PromptTokensDetailsWrapper(web_search_requests=queries) if queries is not None else None + return ModelResponse( + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=details) + ) + + +def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None) -> tuple[dict, dict]: + sl_object = { + "response_cost": response_cost, + "cache_hit": cache_hit, + "cost_breakdown": {"input_cost": 0.1, "output_cost": 0.2, "tool_usage_cost": 0.5, "total_cost": response_cost}, + "model_map_information": {"model_map_key": "gemini-test", "model_map_value": model_info}, + } + return {"response_cost": response_cost, "standard_logging_object": sl_object}, sl_object + + +@pytest.mark.parametrize( + ("allowance", "units", "total_after", "expected"), + [ + (5000, 5, 5, 5), + (5000, 5, 5000, 5), + (5000, 10, 5005, 5), + (5000, 5, 5010, 0), + (2, 5, 5, 2), + (0, 3, 3, 0), + ], +) +def test_freed_units_covers_only_what_landed_under_the_allowance(allowance, units, total_after, expected): + assert freed_units(allowance, units, total_after) == expected + + +def test_counter_key_is_per_group_and_utc_period(): + assert search_free_tier_key({"count": 1, "period": "month", "shared_group": "g3"}, NOW) == ( + "litellm:search_free_tier:g3:2026-10" + ) + assert search_free_tier_key({"count": 1, "period": "day", "shared_group": "g2"}, NOW) == ( + "litellm:search_free_tier:g2:2026-10-20" + ) + + +@pytest.mark.parametrize( + ("period", "boundary"), + [("month", datetime(2026, 11, 1, tzinfo=timezone.utc)), ("day", datetime(2026, 10, 21, tzinfo=timezone.utc))], +) +def test_counter_ttl_outlives_the_period_boundary_by_under_two_hours(period, boundary): + expiry = NOW + timedelta(seconds=search_free_tier_ttl({"count": 1, "period": period, "shared_group": "g"}, NOW)) + assert boundary < expiry < boundary + timedelta(hours=2) + + +@pytest.mark.asyncio +async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_crossing_call(): + counter = _Counter() + model_info = _model_info("per_query", count=7) + reduced = [] + for _ in range(3): + kwargs, sl_object = _call(model_info, response_cost=1.0) + cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + reduced.append((cost, kwargs["response_cost"], sl_object["response_cost"], sl_object["cost_breakdown"])) + (c1, k1, s1, b1), (c2, k2, s2, b2), (c3, k3, s3, b3) = reduced + assert c1 == k1 == s1 == pytest.approx(1.0 - 5 * RATE) + assert b1["original_cost"] == 1.0 and b1["discount_amount"] == pytest.approx(5 * RATE) + assert b1["total_cost"] == pytest.approx(c1) and b1["tool_usage_cost"] == 0.5 + assert c2 == k2 == s2 == pytest.approx(1.0 - 2 * RATE) + assert b2["discount_amount"] == pytest.approx(2 * RATE) + assert c3 == k3 == s3 == 1.0 and "discount_amount" not in b3 + assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 15} + assert counter.ttls["litellm:search_free_tier:test-group:2026-10"] > 0 + + +@pytest.mark.asyncio +async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_query_count(): + counter = _Counter() + model_info = _model_info("per_prompt", count=1, period="day") + kwargs, sl_object = _call(model_info, response_cost=1.0) + first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + kwargs, sl_object = _call(model_info, response_cost=1.0) + second = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + assert first == pytest.approx(1.0 - RATE) + assert second == 1.0 + assert counter.totals == {"litellm:search_free_tier:test-group:2026-10-20": 2} + + +@pytest.mark.asyncio +async def test_sidecar_shaped_response_mapping_is_counted_like_a_model_response(): + counter = _Counter() + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) + response = { + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + "prompt_tokens_details": {"web_search_requests": 3}, + } + } + assert await apply_search_free_tier(kwargs, sl_object, 1.0, response, counter, NOW) == pytest.approx(1.0 - 3 * RATE) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("counter", [None, _BrokenCounter()]) +async def test_full_price_when_no_counter_is_available(counter): + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 + assert kwargs["response_cost"] == 1.0 and "discount_amount" not in sl_object["cost_breakdown"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response", [_response(0), _response(None), ModelResponse()]) +async def test_requests_without_grounding_never_touch_the_counter(response): + counter = _Counter() + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, response, counter, NOW) == 1.0 + assert counter.totals == {} + + +@pytest.mark.asyncio +async def test_models_without_an_allowance_never_touch_the_counter(): + counter = _Counter() + model_info = {k: v for k, v in _model_info("per_query", count=100).items() if k != "search_free_requests"} + kwargs, sl_object = _call(model_info, response_cost=1.0) + assert await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) == 1.0 + assert counter.totals == {} + + +@pytest.mark.asyncio +async def test_a_second_callback_run_on_the_same_request_counts_it_once(): + counter = _Counter() + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=1.0) + first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) + second = await apply_search_free_tier(kwargs, sl_object, first, _response(5), counter, NOW) + assert first == second == pytest.approx(1.0 - 5 * RATE) + assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5} + + +@pytest.mark.asyncio +async def test_a_cache_hit_with_grounded_usage_is_not_counted_and_stays_at_zero(): + counter = _Counter() + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=0.0, cache_hit=True) + assert await apply_search_free_tier(kwargs, sl_object, 0.0, _response(2), counter, NOW) == 0.0 + assert counter.totals == {} + assert "discount_amount" not in sl_object["cost_breakdown"] + + +@pytest.mark.asyncio +async def test_freed_cost_never_exceeds_the_billed_cost(): + counter = _Counter() + billed = 3 * RATE + kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=billed) + assert await apply_search_free_tier(kwargs, sl_object, billed, _response(5), counter, NOW) == 0.0 + assert sl_object["cost_breakdown"]["discount_amount"] == pytest.approx(billed) + assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5} diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 66c52a76318..0186a33dbc4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2558,7 +2558,7 @@ async def test_bedrock_agent_runtime_region_lookup_leaves_the_event_loop_free(mo ) assert vault.read.is_set() - assert not vault.stalled_the_loop.is_set() + assert not vault.read_on_loop.is_set() assert retrieve.calls.last.request.headers["authorization"].startswith("AWS4-HMAC-SHA256") diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index a9bc82d141d..e70b93f7955 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -1681,7 +1681,7 @@ async def test_a_cleanup_delete_batch_renders_a_postgres_delete_span_for_its_tab assert await postgres_span_names() == ("postgres.delete LiteLLM_SpendLogs",) -# --- spend log storage measurement and the reclaimable-bytes state --------------------------------------- +# --- size rule (maximum_spend_logs_size) --------------------------------------------------------------- def _storage(allocated: int, live: int, dead: int, inserted: int | None = None, deleted: int | None = None) -> dict: @@ -1694,6 +1694,50 @@ def _storage(allocated: int, live: int, dead: int, inserted: int | None = None, } +def _wire_size_probes( + db, storage_rows: list[dict], cutoff: datetime, saved_state: SizeCleanupState | None = None +) -> list[SizeCleanupState]: + """ + Answer the size pass's catalog reads: the storage probe walks `storage_rows` one measurement per call + (repeating the last), the step-cutoff query returns `cutoff`, the remaining-rows probe returns 0, and the + LiteLLM_Config row holds `saved_state`. Writes of that row land in the returned list instead of + db.execute_raw, so DELETE counts stay what tests configure. + """ + measurements = iter(storage_rows) + last = [storage_rows[-1]] + saves: list[SizeCleanupState] = [] + deletes = db.execute_raw + + async def _query_raw(sql, *args): + if "pg_total_relation_size" in sql: + row = next(measurements, None) + if row is None: + row = last[0] + return [row] + if "OFFSET" in sql: + return [{"startTime": cutoff}] + if '"LiteLLM_Config"' in sql: + return [{"state_json": saved_state.model_dump_json()}] if saved_state is not None else [] + return [{"remaining": 0}] + + async def _execute_raw(sql, *args): + if '"LiteLLM_Config"' in sql: + saves.append(SizeCleanupState.model_validate_json(args[1])) + return 1 + return await deletes(sql, *args) + + db.query_raw = AsyncMock(side_effect=_query_raw) + db.execute_raw = _execute_raw + db.deletes = deletes + return saves + + +def _size_cleaner(settings: dict, **kwargs) -> SpendLogCleanup: + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_size": "1GB", **settings}, **kwargs) + cleaner.pod_lock_manager = None + return cleaner + + @pytest.mark.parametrize( ("value", "expected"), [ @@ -1719,6 +1763,160 @@ def test_size_in_bytes_rejects_what_cannot_be_a_size(value): size_in_bytes(value) +def test_an_unset_size_means_no_size_rule_and_a_bad_target_falls_back_to_the_size(caplog): + assert SpendLogCleanup(general_settings={})._size_rule() is None + rule = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "200MB"})._size_rule() + assert rule is not None + assert (rule.limit_bytes, rule.target_bytes) == (1024**3, 200 * 1024**2) + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + rule_without_room = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "2GB"})._size_rule() + assert rule_without_room is not None + assert rule_without_room.target_bytes == rule_without_room.limit_bytes == 1024**3 + assert any("must be below maximum_spend_logs_size" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_size_pass_under_the_limit_issues_no_delete(): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=0) + _wire_size_probes(db, [_storage(allocated=500 * 1024**2, live=1000, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + + await _size_cleaner({}).cleanup_spend_logs_by_size(client) + + assert db.deletes.await_count == 0 + + +@pytest.mark.asyncio +async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_is_under_the_target(caplog): + """ + Over the limit, the pass turns the next step of oldest rows into a cutoff and runs the retention DELETE + for the spend logs and their tool-index rows, re-measuring after each step; it stops on the first + measurement under the target, not when the file size (which deletes never shrink) moves. + """ + db = MagicMock() + _wire_tx(db) + # 3 GiB over 3000 rows is 1 MiB a row, so 2560 rows take it from 3 GiB to the 512 MiB target: one step + # of 3 batches; the tool index is cleaned on the same cutoff + db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0]) + step_cutoff = datetime(2026, 9, 1, tzinfo=timezone.utc) + allocated = 3 * 1024**3 + _wire_size_probes(db, [_storage(allocated, live=3000, dead=0)], cutoff=step_cutoff) + client = MagicMock() + client.db = db + cleaner = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "512MB"}) + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await cleaner.cleanup_spend_logs_by_size(client) + + deletes = [(call[0][0], call[0][1:]) for call in db.deletes.call_args_list] + assert len(deletes) == 5 + assert 'DELETE FROM "LiteLLM_SpendLogs"' in deletes[0][0] + assert 'WHERE "startTime" < $1' in deletes[0][0] + assert deletes[0][1][0] == step_cutoff + assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in deletes[3][0] + assert deletes[3][1][0] == step_cutoff + offset_query = next(call for call in db.query_raw.call_args_list if "OFFSET" in call[0][0]) + assert offset_query[0][1] == 3000 + summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) + assert "outcome=completed" in summary + assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=target_reached" in summary + assert "LiteLLM_SpendLogToolIndex (size rule): deleted=300 stop_reason=target_reached" in summary + assert f"allocated_before={allocated}" in summary + assert "estimated_live_before=3221225472" in summary + assert "estimated_live_after=0" in summary + assert "target_reached=true above_target=0" in summary + assert "VACUUM FULL or pg_repack" in summary + + +@pytest.mark.asyncio +async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_the_target_it_is(caplog): + db = MagicMock() + _wire_tx(db) + # 4 GiB over 4000 rows: 3584 rows must go to reach the 512 MiB target, the cap allows 3 batches + db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0]) + allocated = 4 * 1024**3 + saves = _wire_size_probes(db, [_storage(allocated, live=4000, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + cleaner = _size_cleaner( + {"maximum_spend_logs_cleanup_target_size": "512MB", "maximum_spend_logs_cleanup_max_batches": 3} + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await cleaner.cleanup_spend_logs_by_size(client) + + assert db.deletes.await_count == 4 + summary = next(r for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) + assert summary.levelno == logging.WARNING + assert "outcome=batch_cap_reached" in summary.getMessage() + assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=batch_cap_reached" in summary.getMessage() + # 1000 of 4000 rows left is 1 GiB live, 512 MiB above the target + assert f"target_reached=false above_target={512 * 1024**2}" in summary.getMessage() + # the next check starts from what this run freed, with the deletes already counted out of n_tup_del + assert saves == [ + SizeCleanupState( + allocated_bytes=allocated, + reclaimable_bytes=round(3000 * allocated / 4000), + bytes_per_row=allocated / 4000, + inserted_total=4000, + deleted_total=3000, + ) + ] + + +@pytest.mark.asyncio +async def test_size_pass_skips_on_stale_row_statistics_instead_of_trusting_an_estimate_of_zero(caplog): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=1000) + _wire_size_probes(db, [_storage(allocated=5 * 1024**3, live=0, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await _size_cleaner({}).cleanup_spend_logs_by_size(client) + + assert db.deletes.await_count == 0 + assert any("row statistics are stale" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_size_pass_does_not_refire_after_autovacuum_forgets_the_rows_it_deleted(): + """ + The run before this one deleted 15,000 of 20,000 rows and saved what it freed. Autovacuum has since + turned those dead tuples into free space, so pg_stat_user_tables reports 5,000 live and 0 dead rows while + the files still measure the full 32 MB. Trusting the live share would call all 32 MB live and empty the + table; the saved count says 8 MB are live, under the limit, so nothing is deleted. + """ + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=1000) + after_last_run = SizeCleanupState( + allocated_bytes=32_030_720, + reclaimable_bytes=24_015_000, + bytes_per_row=1601.0, + inserted_total=20_000, + deleted_total=15_000, + ) + saves = _wire_size_probes( + db, + [_storage(32_030_720, live=5000, dead=0, inserted=20_000, deleted=15_000)], + cutoff=datetime.now(timezone.utc), + saved_state=after_last_run, + ) + client = MagicMock() + client.db = db + cleaner = _size_cleaner({"maximum_spend_logs_size": "16MB", "maximum_spend_logs_cleanup_target_size": "10MB"}) + + await cleaner.cleanup_spend_logs_by_size(client) + + assert db.deletes.await_count == 0 + assert saves[0].reclaimable_bytes == 24_015_000 + + _RUN_STATE = SizeCleanupState( allocated_bytes=32_030_720, reclaimable_bytes=24_015_000, @@ -1765,6 +1963,58 @@ def test_advance_size_state_tracks_free_space_through_vacuum_inserts_and_other_d assert 0 <= probe["allocated_bytes"] - advanced.reclaimable_bytes <= probe["allocated_bytes"], name +@pytest.mark.asyncio +async def test_size_only_run_leaves_the_age_rules_to_the_retention_job(): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=0) + _wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + cleaner = _size_cleaner( + {"maximum_spend_logs_retention_period": "7d", "maximum_health_check_retention_period": "30d"} + ) + + await cleaner.cleanup_spend_logs_by_size(client) + + tables = {call[0][0].split('"')[1] for call in db.deletes.call_args_list} + assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("cron", "size_pass_expected"), [(None, False), ("0 3 * * *", True)]) +async def test_the_retention_run_carries_the_size_pass_only_when_one_cron_drives_both(cron, size_pass_expected): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=0) + _wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + cleaner = _size_cleaner({"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_cron": cron}) + + await cleaner.cleanup_old_spend_logs(client) + + measured = any("pg_total_relation_size" in call[0][0] for call in db.query_raw.call_args_list) + assert measured is size_pass_expected + + +@pytest.mark.parametrize( + ("settings", "retention_job", "size_job"), + [ + ({}, False, False), + ({"maximum_spend_logs_retention_period": "7d"}, True, False), + ({"maximum_spend_logs_size": "180GB"}, False, True), + ({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_retention_period": "7d"}, True, True), + ({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_cleanup_cron": "0 3 * * *"}, True, False), + ], +) +def test_which_scheduler_jobs_the_cleanup_settings_ask_for(settings, retention_job, size_job): + from litellm.proxy.proxy_server import _wants_retention_job, _wants_size_check_job + + assert _wants_retention_job(settings) is retention_job + assert _wants_size_check_job(settings) is size_job + + def test_after_deleting_counts_the_rows_and_any_file_growth_as_reclaimable_and_to_state_keeps_the_counters(): storage = SpendLogStorage(allocated_bytes=32_030_720, reclaimable_bytes=0, bytes_per_row=1601.0, live_rows=20_000) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 20b46e27213..628870259a4 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.fallback_generalizations import match_capability from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts +from litellm.types.utils import SearchFreeRequests REPO_ROOT = Path(__file__).parents[2] PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -721,6 +722,33 @@ def test_unregistered_provider_guard_flags_only_labels_nobody_registered(): ] +def _search_free_requests_rows(path: Path) -> Mapping[str, dict[str, JsonValue]]: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + return MappingProxyType( + {key: row for key, row in prices.items() if isinstance(row, dict) and "search_free_requests" in row} + ) + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_search_free_requests_rows_are_grounding_rows_with_one_allowance_per_group(path: Path) -> None: + rows = _search_free_requests_rows(path) + assert rows + allowances = { + key: TypeAdapter(SearchFreeRequests).validate_python(row["search_free_requests"]) for key, row in rows.items() + } + for key, row in rows.items(): + assert row.get("supports_web_search") is True, key + assert row.get("search_context_cost_per_query"), key + assert allowances[key]["count"] > 0, key + for group in {allowance["shared_group"] for allowance in allowances.values()}: + same_group = {(a["count"], a["period"]) for a in allowances.values() if a["shared_group"] == group} + assert len(same_group) == 1, (group, same_group) + + +def test_search_free_requests_backup_rows_match_the_main_catalog() -> None: + assert _search_free_requests_rows(BACKUP_PRICES_PATH) == _search_free_requests_rows(PRICES_PATH) + + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) def test_volcengine_doubao_seed_2_1_pro_260915_row_matches_published_metadata(path: Path): """Volcengine Ark lists doubao-seed-2-1-pro-260915 at 6 / 30 CNY per million input / output tokens with diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 40eb5c31fda..a1bdc4d5723 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1260,6 +1260,16 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "type": "string", "enum": ["per_prompt", "per_query"], }, + "search_free_requests": { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "period": {"type": "string", "enum": ["day", "month"]}, + "shared_group": {"type": "string"}, + }, + "required": ["count", "period", "shared_group"], + "additionalProperties": False, + }, "citation_cost_per_token": {"type": "number"}, "supported_modalities": { "type": "array", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 21ab374d3fb..3039ab73a15 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30072,11 +30072,26 @@ export interface components { * @description Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'. */ maximum_spend_logs_cleanup_run_budget?: string | null; + /** + * Maximum Spend Logs Cleanup Target Size + * @description Level a size-triggered cleanup run deletes down to (e.g. '20GB'). Must be below maximum_spend_logs_size. Unset means the run stops as soon as the tables are back under maximum_spend_logs_size. + */ + maximum_spend_logs_cleanup_target_size?: string | null; /** * Maximum Spend Logs Retention Period * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. */ maximum_spend_logs_retention_period?: string | null; + /** + * Maximum Spend Logs Size + * @description Size rule for spend logs (e.g. '180GB'): when LiteLLM_SpendLogs plus LiteLLM_SpendLogToolIndex (heap, indexes and TOAST, measured with pg_total_relation_size) hold more live data than this, the cleanup job deletes the oldest spend logs first until they are under maximum_spend_logs_cleanup_target_size. Independent of maximum_spend_logs_retention_period: a row is deleted when either rule says so. Unset means no size rule. + */ + maximum_spend_logs_size?: string | null; + /** + * Maximum Spend Logs Size Check Interval + * @description How often the size rule is checked (e.g. '10m'). Defaults to '1d'. Ignored when maximum_spend_logs_cleanup_cron is set, which then drives both the age and the size cleanup. + */ + maximum_spend_logs_size_check_interval?: string | null; /** * Mcp Advertised Versions * @description MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. Modern protocol serving requires explicit opt-in. Apps/Tasks remain disabled.