Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/spend_tracking/spend_capture_rate.py: 37%
112 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""Compare the spend LiteLLM captured for a provider against what that provider billed for the same UTC days.
3LiteLLM's side is ``LiteLLM_DailyUserSpend``, summed over the ``custom_llm_provider`` values that land on the
4provider's bill. The provider's side is its billing API, read with the customer's own billing credential
5(OpenAI: the organization costs endpoint and an admin key in ``OPENAI_ADMIN_KEY``).
6"""
8from collections.abc import Awaitable, Callable, Mapping, Sequence
9from dataclasses import dataclass
10from datetime import date, datetime, timedelta, timezone
11from types import MappingProxyType
12from typing import TYPE_CHECKING, Final, TypeAlias
14from pydantic import BaseModel, ConfigDict, TypeAdapter
15from typing_extensions import assert_never
17from litellm._logging import verbose_proxy_logger
18from litellm.constants import (
19 SPEND_CAPTURE_RATE_CHECK_JOB_ID,
20 SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS,
21 SPEND_CAPTURE_RATE_DOCS_URL,
22)
23from litellm.llms.openai.organization_costs import (
24 OPENAI_ADMIN_KEY_ENV_VAR,
25 BillingHttpGet,
26 OpenAICostsRequestFailed,
27 fetch_openai_daily_costs,
28 provider_billing_get,
29)
30from litellm.secret_managers.main import get_secret_str
31from litellm.types.proxy.spend_capture_rate import (
32 CaptureRateDay,
33 CaptureRateReport,
34 SpendCaptureProvider,
35 SpendCaptureRateCheckSettings,
36)
38if TYPE_CHECKING: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true
39 from litellm.caching.redis_cache import RedisCache
40 from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
41 from litellm.proxy.utils import PrismaClient
43OPENAI_BILLED_LITELLM_PROVIDERS: Final = ("openai", "text-completion-openai")
45CaptureRatePublisher: TypeAlias = Callable[[SpendCaptureProvider, float | None], None] # mutable-ok: Callable params
47_CAPTURED_SPEND_BY_DAY_SQL: Final = """
48 SELECT date, COALESCE(SUM(spend), 0)::float AS spend
49 FROM "LiteLLM_DailyUserSpend"
50 WHERE date >= $1 AND date <= $2 AND custom_llm_provider = ANY($3::text[])
51 GROUP BY date
52"""
55@dataclass(frozen=True, slots=True)
56class ProviderBillingCredentialMissing:
57 provider: SpendCaptureProvider
58 env_var: str
61@dataclass(frozen=True, slots=True)
62class ProviderBillingRequestFailed:
63 provider: SpendCaptureProvider
64 detail: str
67ProviderBillingFailure: TypeAlias = ProviderBillingCredentialMissing | ProviderBillingRequestFailed
68CheckResult: TypeAlias = CaptureRateReport | ProviderBillingFailure
71class _CapturedSpendRow(BaseModel):
72 model_config = ConfigDict(frozen=True, extra="ignore")
74 date: str
75 spend: float
78_CAPTURED_SPEND_ROWS: Final = TypeAdapter(tuple[_CapturedSpendRow, ...])
81async def captured_spend_by_day(
82 prisma_client: "PrismaClient",
83 *,
84 litellm_providers: Sequence[str],
85 start_date: date,
86 end_date: date,
87) -> Mapping[str, float]:
88 """LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values."""
89 rows: Final = await prisma_client.db.query_raw(
90 _CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers)
91 )
92 return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)})
95def _ratio(captured: float, billed: float) -> float | None:
96 return None if billed <= 0 else captured / billed
99def _days(start_date: date, end_date: date) -> tuple[date, ...]:
100 return tuple(start_date + timedelta(days=offset) for offset in range((end_date - start_date).days + 1))
103def compute_capture_rate(
104 *,
105 provider: SpendCaptureProvider,
106 start_date: date,
107 end_date: date,
108 captured_by_day: Mapping[str, float],
109 billed_by_day: Mapping[str, float],
110 threshold: float,
111) -> CaptureRateReport:
112 days: Final = tuple(
113 CaptureRateDay(
114 date=day.isoformat(),
115 captured_spend=captured_by_day.get(day.isoformat(), 0.0),
116 provider_spend=billed_by_day.get(day.isoformat(), 0.0),
117 capture_rate=_ratio(captured_by_day.get(day.isoformat(), 0.0), billed_by_day.get(day.isoformat(), 0.0)),
118 )
119 for day in _days(start_date, end_date)
120 )
121 captured: Final = sum(day.captured_spend for day in days)
122 billed: Final = sum(day.provider_spend for day in days)
123 rate: Final = _ratio(captured, billed)
124 return CaptureRateReport(
125 provider=provider,
126 start_date=start_date.isoformat(),
127 end_date=end_date.isoformat(),
128 captured_spend=captured,
129 provider_spend=billed,
130 capture_rate=rate,
131 threshold=threshold,
132 below_threshold=rate is not None and rate < threshold,
133 days=days,
134 )
137async def capture_rate_report(
138 prisma_client: "PrismaClient",
139 *,
140 provider: SpendCaptureProvider,
141 start_date: date,
142 end_date: date,
143 threshold: float,
144 openai_project_ids: Sequence[str] = (),
145 http_get: BillingHttpGet = provider_billing_get,
146) -> CheckResult:
147 match provider:
148 case "openai": 148 ↛ 171line 148 didn't jump to line 171 because the pattern on line 148 always matched
149 admin_key: Final = get_secret_str(OPENAI_ADMIN_KEY_ENV_VAR)
150 if admin_key is None: 150 ↛ 152line 150 didn't jump to line 152 because the condition on line 150 was always true
151 return ProviderBillingCredentialMissing(provider, OPENAI_ADMIN_KEY_ENV_VAR)
152 billed: Final = await fetch_openai_daily_costs(
153 start_date, end_date, admin_key=admin_key, project_ids=openai_project_ids, http_get=http_get
154 )
155 if isinstance(billed, OpenAICostsRequestFailed):
156 return ProviderBillingRequestFailed(provider, billed.detail)
157 captured: Final = await captured_spend_by_day(
158 prisma_client,
159 litellm_providers=OPENAI_BILLED_LITELLM_PROVIDERS,
160 start_date=start_date,
161 end_date=end_date,
162 )
163 return compute_capture_rate(
164 provider=provider,
165 start_date=start_date,
166 end_date=end_date,
167 captured_by_day=captured,
168 billed_by_day=billed,
169 threshold=threshold,
170 )
171 case _:
172 assert_never(provider)
175def alert_message(result: CheckResult) -> str | None:
176 """The alert a check outcome warrants, or ``None`` when the capture rate is healthy."""
177 match result:
178 case ProviderBillingCredentialMissing(provider=provider, env_var=env_var):
179 return (
180 f"Spend capture-rate check: {env_var} is not set, so the {provider} bill cannot be read. "
181 f"Set it or remove general_settings.spend_capture_rate_check. {SPEND_CAPTURE_RATE_DOCS_URL}"
182 )
183 case ProviderBillingRequestFailed(provider=provider, detail=detail):
184 return f"Spend capture-rate check: could not read the {provider} bill ({detail}). {SPEND_CAPTURE_RATE_DOCS_URL}"
185 case CaptureRateReport():
186 if not result.below_threshold or result.capture_rate is None:
187 return None
188 return (
189 f"Spend capture rate for {result.provider} is {result.capture_rate:.1%}, under the "
190 f"{result.threshold:.0%} threshold: LiteLLM captured ${result.captured_spend:,.2f} of the "
191 f"${result.provider_spend:,.2f} {result.provider} bill for {result.start_date} to {result.end_date}. "
192 f"Requests reach {result.provider} outside LiteLLM or cost tracking is dropping spend. "
193 f"{SPEND_CAPTURE_RATE_DOCS_URL}"
194 )
195 case _:
196 assert_never(result)
199def _published_rate(result: CheckResult) -> float | None:
200 """The gauge value: the rate, or ``None`` (NaN on the gauge) when this window produced no rate."""
201 return result.capture_rate if isinstance(result, CaptureRateReport) else None
204async def _check_every_provider(
205 prisma_client: "PrismaClient",
206 settings: SpendCaptureRateCheckSettings,
207 *,
208 publish: CaptureRatePublisher,
209 today: date | None,
210 http_get: BillingHttpGet,
211) -> tuple[CheckResult, ...]:
212 """Check every configured provider over the closed days before ``today`` and publish each outcome."""
213 end_date: Final = (today or datetime.now(timezone.utc).date()) - timedelta(days=1)
214 start_date: Final = end_date - timedelta(days=settings.lookback_days - 1)
215 results: Final = tuple(
216 [
217 await capture_rate_report(
218 prisma_client,
219 provider=provider,
220 start_date=start_date,
221 end_date=end_date,
222 threshold=settings.threshold,
223 openai_project_ids=settings.openai_project_ids,
224 http_get=http_get,
225 )
226 for provider in settings.providers
227 ]
228 )
229 for result in results:
230 publish(result.provider, _published_rate(result))
231 verbose_proxy_logger.info("Spend capture-rate check: %s", result)
232 return results
235def _alert_messages(results: Sequence[CheckResult]) -> tuple[str, ...]:
236 return tuple(message for message in map(alert_message, results) if message is not None)
239async def run_spend_capture_rate_check(
240 prisma_client: "PrismaClient",
241 settings: SpendCaptureRateCheckSettings,
242 *,
243 alert: Callable[[str], Awaitable[None]],
244 publish: CaptureRatePublisher,
245 today: date | None = None,
246 http_get: BillingHttpGet = provider_billing_get,
247) -> tuple[CheckResult, ...]:
248 """Check every configured provider, publish each rate, and alert on every outcome that warrants one."""
249 results: Final = await _check_every_provider(
250 prisma_client, settings, publish=publish, today=today, http_get=http_get
251 )
252 for message in _alert_messages(results):
253 await alert(message)
254 return results
257async def run_scheduled_spend_capture_rate_check(
258 prisma_client: "PrismaClient",
259 settings: SpendCaptureRateCheckSettings,
260 *,
261 pod_lock_manager: "PodLockManager | None",
262 alert: Callable[[str], Awaitable[None]],
263 publish: CaptureRatePublisher,
264 today: date | None = None,
265 http_get: BillingHttpGet = provider_billing_get,
266) -> tuple[CheckResult, ...]:
267 """Every worker publishes its own gauge; the first replica whose finished check has an alert claims the window."""
268 results: Final = await _check_every_provider(
269 prisma_client, settings, publish=publish, today=today, http_get=http_get
270 )
271 messages: Final = _alert_messages(results)
272 if not messages:
273 return results
274 if not await _claims_alert_window(pod_lock_manager):
275 verbose_proxy_logger.info("Spend capture-rate check: another pod alerted this window")
276 return results
277 for message in messages:
278 await alert(message)
279 return results
282async def _claims_alert_window(pod_lock_manager: "PodLockManager | None") -> bool:
283 """The lock is left to expire, so every replica firing within its TTL of the winner stays quiet."""
284 redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache
285 if pod_lock_manager is None or redis_cache is None:
286 return True
287 acquired: Final = await pod_lock_manager.acquire_lock(
288 cronjob_id=SPEND_CAPTURE_RATE_CHECK_JOB_ID, ttl=SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS
289 )
290 return acquired or not await _lock_is_held(pod_lock_manager, redis_cache)
293async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool:
294 try:
295 return bool(
296 await redis_cache.async_get_cache(pod_lock_manager.get_redis_lock_key(SPEND_CAPTURE_RATE_CHECK_JOB_ID))
297 )
298 except Exception as exc: # noqa: BLE001 # an unreadable lock must not silence the alert
299 verbose_proxy_logger.warning("Spend capture-rate check: could not read the lock: %s", exc)
300 return False