mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(spend): capture-rate check of LiteLLM spend against the OpenAI bill (#43044)
* feat(spend): capture-rate check of LiteLLM spend against the OpenAI bill * fix(spend): claim the alert lock after the check, NaN gauge on no rate, 180-day range cap, live settings, OpenAI adapter under llms * fix(spend): chart the capture-rate gauge in the all-metrics dashboard and clear it when the check is removed * fix(prometheus): record the capture-rate gauge when api_provider is an excluded label --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
27d1974e2f
commit
c183d810f3
18 changed files with 2012 additions and 433 deletions
|
|
@ -1697,6 +1697,63 @@
|
|||
"title": "litellm_video_duration_seconds_metric rate",
|
||||
"type": "timeseries"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"description": "Share of the provider's bill LiteLLM captured as spend over the scheduled capture-rate check's window (needs general_settings.spend_capture_rate_check); NaN while no rate is available",
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"color": {
|
||||
"mode": "palette-classic"
|
||||
},
|
||||
"custom": {
|
||||
"drawStyle": "line",
|
||||
"fillOpacity": 10,
|
||||
"lineWidth": 1,
|
||||
"showPoints": "never",
|
||||
"spanNulls": false
|
||||
},
|
||||
"unit": "percentunit"
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 12,
|
||||
"y": 107
|
||||
},
|
||||
"id": 111,
|
||||
"options": {
|
||||
"legend": {
|
||||
"calcs": [],
|
||||
"displayMode": "list",
|
||||
"placement": "bottom",
|
||||
"showLegend": true
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "max by (api_provider) (litellm_spend_capture_rate)",
|
||||
"legendFormat": "{{api_provider}}",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
}
|
||||
],
|
||||
"title": "litellm_spend_capture_rate",
|
||||
"type": "timeseries"
|
||||
},
|
||||
{
|
||||
"collapsed": false,
|
||||
"gridPos": {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# LiteLLM All Prometheus Metrics dashboard
|
||||
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (134 families across 95 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
|
||||
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
|
||||
|
||||
|
|
|
|||
|
|
@ -2122,6 +2122,14 @@ PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
|
|||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job"
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through"
|
||||
SPEND_CAPTURE_RATE_CHECK_JOB_ID: Final[str] = "spend_capture_rate_check_job"
|
||||
SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
SPEND_CAPTURE_RATE_MAX_RANGE_DAYS: Final[int] = 180
|
||||
SPEND_CAPTURE_RATE_DOCS_URL: Final[str] = "https://docs.litellm.ai/docs/proxy/spend_capture_rate"
|
||||
OPENAI_ORGANIZATION_COSTS_URL: Final[str] = "https://api.openai.com/v1/organization/costs"
|
||||
# Buckets per page the OpenAI costs endpoint allows (1 to 180, default 7), 2026-09-24
|
||||
OPENAI_ORGANIZATION_COSTS_PAGE_LIMIT: Final[int] = 180
|
||||
PROVIDER_BILLING_TIMEOUT_SECONDS: Final[float] = 30.0
|
||||
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
|
||||
# run's cutoff are stamped by different hosts, so clock skew between them must not let
|
||||
# one run delete a charge another just wrote. A stale row is hours old and a concurrent
|
||||
|
|
|
|||
|
|
@ -729,6 +729,15 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_zero_cost_requests_total"),
|
||||
)
|
||||
|
||||
self.litellm_spend_capture_rate = self._gauge_factory(
|
||||
"litellm_spend_capture_rate",
|
||||
(
|
||||
"Share of the provider's bill LiteLLM captured as spend over the scheduled check's window "
|
||||
"(captured spend / provider bill), by api_provider; NaN when the last check produced no rate"
|
||||
),
|
||||
labelnames=self.get_labels_for_metric("litellm_spend_capture_rate"),
|
||||
)
|
||||
|
||||
# Cache metrics
|
||||
self.litellm_cache_hits_metric = self._counter_factory(
|
||||
name="litellm_cache_hits_metric",
|
||||
|
|
@ -2028,6 +2037,15 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
self.litellm_zero_cost_requests_total.labels(**labels).inc()
|
||||
|
||||
def set_spend_capture_rate(self, api_provider: str, capture_rate: float | None) -> None:
|
||||
labels: Final = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric("litellm_spend_capture_rate"),
|
||||
enum_values=UserAPIKeyLabelValues(api_provider=api_provider),
|
||||
)
|
||||
gauge: Final = self.litellm_spend_capture_rate
|
||||
series: Final = gauge.labels(**labels) if labels else gauge
|
||||
series.set(math.nan if capture_rate is None else capture_rate)
|
||||
|
||||
@staticmethod
|
||||
def _get_remaining_from_v3_rate_limit_headers(
|
||||
standard_logging_payload: StandardLoggingPayload | None,
|
||||
|
|
|
|||
133
litellm/llms/openai/organization_costs.py
Normal file
133
litellm/llms/openai/organization_costs.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
"""OpenAI's organization costs endpoint: the USD the organization was billed per UTC day, read with an admin key."""
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm.constants import (
|
||||
OPENAI_ORGANIZATION_COSTS_PAGE_LIMIT,
|
||||
OPENAI_ORGANIZATION_COSTS_URL,
|
||||
PROVIDER_BILLING_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
OPENAI_ADMIN_KEY_ENV_VAR: Final = "OPENAI_ADMIN_KEY"
|
||||
|
||||
BillingHttpGet: TypeAlias = Callable[
|
||||
[str, Mapping[str, object], Mapping[str, str]], # mutable-ok: Callable parameter list is type syntax
|
||||
Awaitable[httpx.Response],
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OpenAICostsRequestFailed:
|
||||
detail: str
|
||||
|
||||
|
||||
class _OpenAICostAmount(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
value: float
|
||||
currency: Literal["usd"]
|
||||
|
||||
|
||||
class _OpenAICostResult(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
amount: _OpenAICostAmount
|
||||
|
||||
|
||||
class _OpenAICostBucket(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
start_time: int
|
||||
results: tuple[_OpenAICostResult, ...] = ()
|
||||
|
||||
|
||||
class _OpenAICostsPage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
data: tuple[_OpenAICostBucket, ...]
|
||||
has_more: bool = False
|
||||
next_page: str | None = None
|
||||
|
||||
|
||||
async def provider_billing_get(url: str, params: Mapping[str, object], headers: Mapping[str, str]) -> httpx.Response:
|
||||
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.ProviderBilling)
|
||||
return await client.get(
|
||||
url,
|
||||
params=dict(params), # mutable-ok: AsyncHTTPHandler.get takes dict params
|
||||
headers=dict(headers), # mutable-ok: AsyncHTTPHandler.get takes dict headers
|
||||
timeout=PROVIDER_BILLING_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _utc_midnight(day: date) -> int:
|
||||
return int(datetime(day.year, day.month, day.day, tzinfo=timezone.utc).timestamp())
|
||||
|
||||
|
||||
def _bucket_day(bucket: _OpenAICostBucket) -> str:
|
||||
return datetime.fromtimestamp(bucket.start_time, tz=timezone.utc).date().isoformat()
|
||||
|
||||
|
||||
async def fetch_openai_daily_costs(
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
*,
|
||||
admin_key: str,
|
||||
project_ids: Sequence[str] = (),
|
||||
http_get: BillingHttpGet = provider_billing_get,
|
||||
) -> Mapping[str, float] | OpenAICostsRequestFailed:
|
||||
"""USD billed by OpenAI per UTC day (ISO date) over the closed range, following pagination to the end."""
|
||||
scope: Final = (("project_ids[]", tuple(project_ids)),) if project_ids else ()
|
||||
window: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("start_time", _utc_midnight(start_date)),
|
||||
("end_time", _utc_midnight(end_date + timedelta(days=1))),
|
||||
("bucket_width", "1d"),
|
||||
("limit", OPENAI_ORGANIZATION_COSTS_PAGE_LIMIT),
|
||||
*scope,
|
||||
)
|
||||
}
|
||||
)
|
||||
headers: Final[Mapping[str, str]] = MappingProxyType({"Authorization": f"Bearer {admin_key}"})
|
||||
|
||||
async def fetch_from(page: str | None) -> tuple[_OpenAICostBucket, ...] | OpenAICostsRequestFailed:
|
||||
params: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{key: value for key, value in (*window.items(), ("page", page)) if value is not None}
|
||||
)
|
||||
try:
|
||||
response: Final = await http_get(OPENAI_ORGANIZATION_COSTS_URL, params, headers)
|
||||
except httpx.HTTPError as exc:
|
||||
return OpenAICostsRequestFailed(f"request failed: {exc}")
|
||||
if response.status_code != 200:
|
||||
return OpenAICostsRequestFailed(f"HTTP {response.status_code}: {response.text[:300]}")
|
||||
try:
|
||||
parsed: Final = _OpenAICostsPage.model_validate(response.json())
|
||||
except (ValueError, ValidationError) as exc:
|
||||
return OpenAICostsRequestFailed(f"unexpected response shape: {exc}")
|
||||
if not parsed.has_more or parsed.next_page is None:
|
||||
return parsed.data
|
||||
rest: Final = await fetch_from(parsed.next_page)
|
||||
return rest if isinstance(rest, OpenAICostsRequestFailed) else parsed.data + rest
|
||||
|
||||
buckets: Final = await fetch_from(None)
|
||||
if isinstance(buckets, OpenAICostsRequestFailed):
|
||||
return buckets
|
||||
days: Final = frozenset(_bucket_day(bucket) for bucket in buckets)
|
||||
return MappingProxyType(
|
||||
{
|
||||
day: sum(
|
||||
result.amount.value for bucket in buckets if _bucket_day(bucket) == day for result in bucket.results
|
||||
)
|
||||
for day in days
|
||||
}
|
||||
)
|
||||
|
|
@ -51,6 +51,7 @@ from litellm.types.proxy.carried_budget_state import (
|
|||
UserBudgetSnapshot,
|
||||
)
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings
|
||||
from litellm.types.router import RouterErrors, UpdateRouterConfig
|
||||
from litellm.types.router_weights import validate_router_settings_dict
|
||||
from litellm.types.secret_managers.main import KeyManagementSystem
|
||||
|
|
@ -779,6 +780,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/global/spend/provider",
|
||||
"/global/spend/tags",
|
||||
"/global/spend/all_tag_names",
|
||||
"/spend/capture_rate",
|
||||
]
|
||||
|
||||
public_routes = frozenset(
|
||||
|
|
@ -2948,6 +2950,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"every replica. On by default; set to tune the window, pin a job, or turn it off."
|
||||
),
|
||||
)
|
||||
spend_capture_rate_check: SpendCaptureRateCheckSettings | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"Daily check of the spend LiteLLM captured against the provider's own bill (OpenAI via OPENAI_ADMIN_KEY). "
|
||||
"Publishes litellm_spend_capture_rate per provider and alerts when the ratio over the lookback window "
|
||||
"falls under the threshold (default 0.9). Off unless set."
|
||||
),
|
||||
)
|
||||
maximum_spend_logs_retention_period: str | None = Field(
|
||||
None,
|
||||
description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.",
|
||||
|
|
|
|||
|
|
@ -293,6 +293,7 @@ from litellm.constants import (
|
|||
REALTIME_SESSION_FAILURE_LOGGED_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
|
||||
SPEND_CAPTURE_RATE_CHECK_JOB_ID,
|
||||
USER_SPEND_ALERTS_JOB_ID,
|
||||
WEEKLY_SPEND_REPORT_JOB_ID,
|
||||
)
|
||||
|
|
@ -758,6 +759,9 @@ from litellm.proxy.spend_tracking.budget_reservation import (
|
|||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
|
||||
run_scheduled_daily_global_spend_reconcile,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_capture_rate import (
|
||||
run_scheduled_spend_capture_rate_check,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
PendingSpendIncrement,
|
||||
active_spend_counter_batch,
|
||||
|
|
@ -856,6 +860,7 @@ from litellm.types.proxy.model_deprecation import (
|
|||
DEFAULT_DEPRECATION_WARN_DAYS,
|
||||
ModelDeprecationResponse,
|
||||
)
|
||||
from litellm.types.proxy.spend_capture_rate import SpendCaptureProvider, SpendCaptureRateCheckSettings
|
||||
from litellm.types.realtime import RealtimeQueryParams
|
||||
from litellm.types.router import (
|
||||
ClassifierPlugin,
|
||||
|
|
@ -5060,6 +5065,11 @@ def _bind_general_settings_store(settings: SettingsStore) -> None:
|
|||
general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings
|
||||
|
||||
|
||||
def _current_general_settings() -> Mapping[str, object]:
|
||||
"""The live ``general_settings``, whichever object a config reload has bound since the caller was created."""
|
||||
return general_settings
|
||||
|
||||
|
||||
@lru_cache(maxsize=4096)
|
||||
def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -10448,6 +10458,13 @@ class ProxyStartupEvent:
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
cls._initialize_spend_capture_rate_check_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
read_general_settings=_current_general_settings,
|
||||
)
|
||||
|
||||
### PTU DAILY ROLLUP ###
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
@ -10822,6 +10839,64 @@ class ProxyStartupEvent:
|
|||
next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _initialize_spend_capture_rate_check_job(
|
||||
cls,
|
||||
scheduler: AsyncIOScheduler,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
prisma_client: PrismaClient,
|
||||
read_general_settings: Callable[[], Mapping[str, object]],
|
||||
) -> None:
|
||||
"""The job always runs and re-reads ``spend_capture_rate_check`` each run; an absent setting clears the gauge."""
|
||||
cls._spend_capture_rate_check_settings(read_general_settings())
|
||||
|
||||
async def alert(message: str) -> None:
|
||||
await proxy_logging_obj.alerting_handler(
|
||||
message=message,
|
||||
level="High",
|
||||
alert_type=AlertType.failed_tracking_spend,
|
||||
)
|
||||
|
||||
def publish(provider: str, capture_rate: float | None) -> None:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
for logger in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=PrometheusLogger):
|
||||
if isinstance(logger, PrometheusLogger):
|
||||
logger.set_spend_capture_rate(api_provider=provider, capture_rate=capture_rate)
|
||||
|
||||
async def check() -> None:
|
||||
settings: Final = cls._spend_capture_rate_check_settings(read_general_settings())
|
||||
if settings is None:
|
||||
for provider in get_args(SpendCaptureProvider):
|
||||
publish(provider, None)
|
||||
return
|
||||
await run_scheduled_spend_capture_rate_check(
|
||||
prisma_client,
|
||||
settings,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
alert=alert,
|
||||
publish=publish,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
check,
|
||||
"cron",
|
||||
hour=1,
|
||||
minute=15,
|
||||
timezone="UTC",
|
||||
id=SPEND_CAPTURE_RATE_CHECK_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _spend_capture_rate_check_settings(
|
||||
general_settings: Mapping[str, object],
|
||||
) -> SpendCaptureRateCheckSettings | None:
|
||||
raw_settings: Final = general_settings.get("spend_capture_rate_check")
|
||||
return None if raw_settings is None else SpendCaptureRateCheckSettings.model_validate(raw_settings)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_slack_alerting_jobs(
|
||||
cls,
|
||||
|
|
|
|||
300
litellm/proxy/spend_tracking/spend_capture_rate.py
Normal file
300
litellm/proxy/spend_tracking/spend_capture_rate.py
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
"""Compare the spend LiteLLM captured for a provider against what that provider billed for the same UTC days.
|
||||
|
||||
LiteLLM's side is ``LiteLLM_DailyUserSpend``, summed over the ``custom_llm_provider`` values that land on the
|
||||
provider's bill. The provider's side is its billing API, read with the customer's own billing credential
|
||||
(OpenAI: the organization costs endpoint and an admin key in ``OPENAI_ADMIN_KEY``).
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
SPEND_CAPTURE_RATE_CHECK_JOB_ID,
|
||||
SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS,
|
||||
SPEND_CAPTURE_RATE_DOCS_URL,
|
||||
)
|
||||
from litellm.llms.openai.organization_costs import (
|
||||
OPENAI_ADMIN_KEY_ENV_VAR,
|
||||
BillingHttpGet,
|
||||
OpenAICostsRequestFailed,
|
||||
fetch_openai_daily_costs,
|
||||
provider_billing_get,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.proxy.spend_capture_rate import (
|
||||
CaptureRateDay,
|
||||
CaptureRateReport,
|
||||
SpendCaptureProvider,
|
||||
SpendCaptureRateCheckSettings,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
OPENAI_BILLED_LITELLM_PROVIDERS: Final = ("openai", "text-completion-openai")
|
||||
|
||||
CaptureRatePublisher: TypeAlias = Callable[[SpendCaptureProvider, float | None], None] # mutable-ok: Callable params
|
||||
|
||||
_CAPTURED_SPEND_BY_DAY_SQL: Final = """
|
||||
SELECT date, COALESCE(SUM(spend), 0)::float AS spend
|
||||
FROM "LiteLLM_DailyUserSpend"
|
||||
WHERE date >= $1 AND date <= $2 AND custom_llm_provider = ANY($3::text[])
|
||||
GROUP BY date
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderBillingCredentialMissing:
|
||||
provider: SpendCaptureProvider
|
||||
env_var: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderBillingRequestFailed:
|
||||
provider: SpendCaptureProvider
|
||||
detail: str
|
||||
|
||||
|
||||
ProviderBillingFailure: TypeAlias = ProviderBillingCredentialMissing | ProviderBillingRequestFailed
|
||||
CheckResult: TypeAlias = CaptureRateReport | ProviderBillingFailure
|
||||
|
||||
|
||||
class _CapturedSpendRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
date: str
|
||||
spend: float
|
||||
|
||||
|
||||
_CAPTURED_SPEND_ROWS: Final = TypeAdapter(tuple[_CapturedSpendRow, ...])
|
||||
|
||||
|
||||
async def captured_spend_by_day(
|
||||
prisma_client: "PrismaClient",
|
||||
*,
|
||||
litellm_providers: Sequence[str],
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
) -> Mapping[str, float]:
|
||||
"""LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values."""
|
||||
rows: Final = await prisma_client.db.query_raw(
|
||||
_CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers)
|
||||
)
|
||||
return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)})
|
||||
|
||||
|
||||
def _ratio(captured: float, billed: float) -> float | None:
|
||||
return None if billed <= 0 else captured / billed
|
||||
|
||||
|
||||
def _days(start_date: date, end_date: date) -> tuple[date, ...]:
|
||||
return tuple(start_date + timedelta(days=offset) for offset in range((end_date - start_date).days + 1))
|
||||
|
||||
|
||||
def compute_capture_rate(
|
||||
*,
|
||||
provider: SpendCaptureProvider,
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
captured_by_day: Mapping[str, float],
|
||||
billed_by_day: Mapping[str, float],
|
||||
threshold: float,
|
||||
) -> CaptureRateReport:
|
||||
days: Final = tuple(
|
||||
CaptureRateDay(
|
||||
date=day.isoformat(),
|
||||
captured_spend=captured_by_day.get(day.isoformat(), 0.0),
|
||||
provider_spend=billed_by_day.get(day.isoformat(), 0.0),
|
||||
capture_rate=_ratio(captured_by_day.get(day.isoformat(), 0.0), billed_by_day.get(day.isoformat(), 0.0)),
|
||||
)
|
||||
for day in _days(start_date, end_date)
|
||||
)
|
||||
captured: Final = sum(day.captured_spend for day in days)
|
||||
billed: Final = sum(day.provider_spend for day in days)
|
||||
rate: Final = _ratio(captured, billed)
|
||||
return CaptureRateReport(
|
||||
provider=provider,
|
||||
start_date=start_date.isoformat(),
|
||||
end_date=end_date.isoformat(),
|
||||
captured_spend=captured,
|
||||
provider_spend=billed,
|
||||
capture_rate=rate,
|
||||
threshold=threshold,
|
||||
below_threshold=rate is not None and rate < threshold,
|
||||
days=days,
|
||||
)
|
||||
|
||||
|
||||
async def capture_rate_report(
|
||||
prisma_client: "PrismaClient",
|
||||
*,
|
||||
provider: SpendCaptureProvider,
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
threshold: float,
|
||||
openai_project_ids: Sequence[str] = (),
|
||||
http_get: BillingHttpGet = provider_billing_get,
|
||||
) -> CheckResult:
|
||||
match provider:
|
||||
case "openai":
|
||||
admin_key: Final = get_secret_str(OPENAI_ADMIN_KEY_ENV_VAR)
|
||||
if admin_key is None:
|
||||
return ProviderBillingCredentialMissing(provider, OPENAI_ADMIN_KEY_ENV_VAR)
|
||||
billed: Final = await fetch_openai_daily_costs(
|
||||
start_date, end_date, admin_key=admin_key, project_ids=openai_project_ids, http_get=http_get
|
||||
)
|
||||
if isinstance(billed, OpenAICostsRequestFailed):
|
||||
return ProviderBillingRequestFailed(provider, billed.detail)
|
||||
captured: Final = await captured_spend_by_day(
|
||||
prisma_client,
|
||||
litellm_providers=OPENAI_BILLED_LITELLM_PROVIDERS,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
)
|
||||
return compute_capture_rate(
|
||||
provider=provider,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
captured_by_day=captured,
|
||||
billed_by_day=billed,
|
||||
threshold=threshold,
|
||||
)
|
||||
case _:
|
||||
assert_never(provider)
|
||||
|
||||
|
||||
def alert_message(result: CheckResult) -> str | None:
|
||||
"""The alert a check outcome warrants, or ``None`` when the capture rate is healthy."""
|
||||
match result:
|
||||
case ProviderBillingCredentialMissing(provider=provider, env_var=env_var):
|
||||
return (
|
||||
f"Spend capture-rate check: {env_var} is not set, so the {provider} bill cannot be read. "
|
||||
f"Set it or remove general_settings.spend_capture_rate_check. {SPEND_CAPTURE_RATE_DOCS_URL}"
|
||||
)
|
||||
case ProviderBillingRequestFailed(provider=provider, detail=detail):
|
||||
return f"Spend capture-rate check: could not read the {provider} bill ({detail}). {SPEND_CAPTURE_RATE_DOCS_URL}"
|
||||
case CaptureRateReport():
|
||||
if not result.below_threshold or result.capture_rate is None:
|
||||
return None
|
||||
return (
|
||||
f"Spend capture rate for {result.provider} is {result.capture_rate:.1%}, under the "
|
||||
f"{result.threshold:.0%} threshold: LiteLLM captured ${result.captured_spend:,.2f} of the "
|
||||
f"${result.provider_spend:,.2f} {result.provider} bill for {result.start_date} to {result.end_date}. "
|
||||
f"Requests reach {result.provider} outside LiteLLM or cost tracking is dropping spend. "
|
||||
f"{SPEND_CAPTURE_RATE_DOCS_URL}"
|
||||
)
|
||||
case _:
|
||||
assert_never(result)
|
||||
|
||||
|
||||
def _published_rate(result: CheckResult) -> float | None:
|
||||
"""The gauge value: the rate, or ``None`` (NaN on the gauge) when this window produced no rate."""
|
||||
return result.capture_rate if isinstance(result, CaptureRateReport) else None
|
||||
|
||||
|
||||
async def _check_every_provider(
|
||||
prisma_client: "PrismaClient",
|
||||
settings: SpendCaptureRateCheckSettings,
|
||||
*,
|
||||
publish: CaptureRatePublisher,
|
||||
today: date | None,
|
||||
http_get: BillingHttpGet,
|
||||
) -> tuple[CheckResult, ...]:
|
||||
"""Check every configured provider over the closed days before ``today`` and publish each outcome."""
|
||||
end_date: Final = (today or datetime.now(timezone.utc).date()) - timedelta(days=1)
|
||||
start_date: Final = end_date - timedelta(days=settings.lookback_days - 1)
|
||||
results: Final = tuple(
|
||||
[
|
||||
await capture_rate_report(
|
||||
prisma_client,
|
||||
provider=provider,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
threshold=settings.threshold,
|
||||
openai_project_ids=settings.openai_project_ids,
|
||||
http_get=http_get,
|
||||
)
|
||||
for provider in settings.providers
|
||||
]
|
||||
)
|
||||
for result in results:
|
||||
publish(result.provider, _published_rate(result))
|
||||
verbose_proxy_logger.info("Spend capture-rate check: %s", result)
|
||||
return results
|
||||
|
||||
|
||||
def _alert_messages(results: Sequence[CheckResult]) -> tuple[str, ...]:
|
||||
return tuple(message for message in map(alert_message, results) if message is not None)
|
||||
|
||||
|
||||
async def run_spend_capture_rate_check(
|
||||
prisma_client: "PrismaClient",
|
||||
settings: SpendCaptureRateCheckSettings,
|
||||
*,
|
||||
alert: Callable[[str], Awaitable[None]],
|
||||
publish: CaptureRatePublisher,
|
||||
today: date | None = None,
|
||||
http_get: BillingHttpGet = provider_billing_get,
|
||||
) -> tuple[CheckResult, ...]:
|
||||
"""Check every configured provider, publish each rate, and alert on every outcome that warrants one."""
|
||||
results: Final = await _check_every_provider(
|
||||
prisma_client, settings, publish=publish, today=today, http_get=http_get
|
||||
)
|
||||
for message in _alert_messages(results):
|
||||
await alert(message)
|
||||
return results
|
||||
|
||||
|
||||
async def run_scheduled_spend_capture_rate_check(
|
||||
prisma_client: "PrismaClient",
|
||||
settings: SpendCaptureRateCheckSettings,
|
||||
*,
|
||||
pod_lock_manager: "PodLockManager | None",
|
||||
alert: Callable[[str], Awaitable[None]],
|
||||
publish: CaptureRatePublisher,
|
||||
today: date | None = None,
|
||||
http_get: BillingHttpGet = provider_billing_get,
|
||||
) -> tuple[CheckResult, ...]:
|
||||
"""Every worker publishes its own gauge; the first replica whose finished check has an alert claims the window."""
|
||||
results: Final = await _check_every_provider(
|
||||
prisma_client, settings, publish=publish, today=today, http_get=http_get
|
||||
)
|
||||
messages: Final = _alert_messages(results)
|
||||
if not messages:
|
||||
return results
|
||||
if not await _claims_alert_window(pod_lock_manager):
|
||||
verbose_proxy_logger.info("Spend capture-rate check: another pod alerted this window")
|
||||
return results
|
||||
for message in messages:
|
||||
await alert(message)
|
||||
return results
|
||||
|
||||
|
||||
async def _claims_alert_window(pod_lock_manager: "PodLockManager | None") -> bool:
|
||||
"""The lock is left to expire, so every replica firing within its TTL of the winner stays quiet."""
|
||||
redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache
|
||||
if pod_lock_manager is None or redis_cache is None:
|
||||
return True
|
||||
acquired: Final = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=SPEND_CAPTURE_RATE_CHECK_JOB_ID, ttl=SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS
|
||||
)
|
||||
return acquired or not await _lock_is_held(pod_lock_manager, redis_cache)
|
||||
|
||||
|
||||
async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool:
|
||||
try:
|
||||
return bool(
|
||||
await redis_cache.async_get_cache(pod_lock_manager.get_redis_lock_key(SPEND_CAPTURE_RATE_CHECK_JOB_ID))
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # an unreadable lock must not silence the alert
|
||||
verbose_proxy_logger.warning("Spend capture-rate check: could not read the lock: %s", exc)
|
||||
return False
|
||||
|
|
@ -24,7 +24,7 @@ from typing import (
|
|||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import ReadOnly, assert_never
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -32,11 +32,17 @@ from litellm.constants import (
|
|||
EMPTY_MAPPING,
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
SPEND_CAPTURE_RATE_MAX_RANGE_DAYS,
|
||||
)
|
||||
from litellm.litellm_core_utils.classifier_logging import classifier_audit_fields, classifier_input_snapshot
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import ProviderBudgetResponse, ProviderBudgetResponseObject
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.spend_tracking.spend_capture_rate import (
|
||||
ProviderBillingCredentialMissing,
|
||||
ProviderBillingRequestFailed,
|
||||
capture_rate_report,
|
||||
)
|
||||
|
||||
# NOTE: Avoid module-level import from common_utils: proxy_server imports this
|
||||
# module while common_utils may pull proxy_server during init, which can leave
|
||||
|
|
@ -52,6 +58,7 @@ from litellm.repositories.team_repository import TeamRepository
|
|||
from litellm.repositories.verification_token_repository import (
|
||||
VerificationTokenRepository,
|
||||
)
|
||||
from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
|
@ -1183,6 +1190,84 @@ async def get_global_activity_exceptions(
|
|||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/spend/capture_rate",
|
||||
tags=["Budget & Spend Tracking"], # mutable-ok: FastAPI tags kwarg is list-typed
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=CaptureRateReport,
|
||||
)
|
||||
async def get_spend_capture_rate(
|
||||
start_date: Annotated[date, fastapi.Query(description="First UTC day of the range, YYYY-MM-DD")],
|
||||
end_date: Annotated[date, fastapi.Query(description="Last UTC day of the range, YYYY-MM-DD, inclusive")],
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
provider: Annotated[
|
||||
SpendCaptureProvider,
|
||||
fastapi.Query(description="Provider whose bill to compare against; needs OPENAI_ADMIN_KEY set on the proxy"),
|
||||
] = "openai",
|
||||
threshold: Annotated[
|
||||
float, fastapi.Query(gt=0, le=1, description="Ratio under which the report flags below_threshold")
|
||||
] = 0.9,
|
||||
project_ids: Annotated[
|
||||
list[str] | None,
|
||||
fastapi.Query(
|
||||
description=(
|
||||
"Scope the OpenAI bill to these project ids; omit to compare against the whole organization. Captured "
|
||||
"spend is never scoped, so pass every project LiteLLM's OpenAI keys belong to"
|
||||
)
|
||||
),
|
||||
] = None,
|
||||
) -> CaptureRateReport:
|
||||
"""
|
||||
Compare the spend LiteLLM captured for a provider against that provider's own bill, per UTC day.
|
||||
|
||||
Admin only. Reads the provider's billing API with the billing credential set on the proxy
|
||||
(OpenAI: `OPENAI_ADMIN_KEY`) and sums `LiteLLM_DailyUserSpend` for the same days.
|
||||
|
||||
Example:
|
||||
```
|
||||
curl -H "Authorization: Bearer sk-1234" \
|
||||
"http://localhost:4000/spend/capture_rate?provider=openai&start_date=2026-09-17&end_date=2026-09-23"
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if not _is_admin_view_safe(user_api_key_dict):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Only proxy admins can read the capture rate")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
if end_date < start_date:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="end_date must not be before start_date")
|
||||
if (end_date - start_date).days >= SPEND_CAPTURE_RATE_MAX_RANGE_DAYS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Date range too large; maximum is {SPEND_CAPTURE_RATE_MAX_RANGE_DAYS} days",
|
||||
)
|
||||
result: Final = await capture_rate_report(
|
||||
prisma_client,
|
||||
provider=provider,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
threshold=threshold,
|
||||
openai_project_ids=tuple(project_ids or ()),
|
||||
)
|
||||
match result:
|
||||
case ProviderBillingCredentialMissing(env_var=env_var):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail=f"{env_var} is not set on the proxy, so the {provider} bill cannot be read",
|
||||
)
|
||||
case ProviderBillingRequestFailed(detail=detail):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY, detail=f"Could not read the {provider} bill: {detail}"
|
||||
)
|
||||
case CaptureRateReport():
|
||||
return result
|
||||
case _:
|
||||
assert_never(result)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/global/spend/provider",
|
||||
tags=["Budget & Spend Tracking"],
|
||||
|
|
|
|||
|
|
@ -281,6 +281,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
"litellm_guardrail_errors_total",
|
||||
"litellm_guardrail_requests_total",
|
||||
"litellm_zero_cost_requests_total",
|
||||
"litellm_spend_capture_rate",
|
||||
# Cache metrics
|
||||
"litellm_cache_hits_metric",
|
||||
"litellm_cache_misses_metric",
|
||||
|
|
@ -600,6 +601,8 @@ class PrometheusMetricLabels:
|
|||
ZERO_COST_REASON_LABEL,
|
||||
)
|
||||
|
||||
litellm_spend_capture_rate = (UserAPIKeyLabelNames.API_PROVIDER.value,)
|
||||
|
||||
litellm_input_tokens_metric = [
|
||||
UserAPIKeyLabelNames.END_USER.value,
|
||||
UserAPIKeyLabelNames.API_KEY_HASH.value,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ class httpxSpecialProvider(str, Enum):
|
|||
Search = "search"
|
||||
MCP = "mcp"
|
||||
RAG = "rag"
|
||||
ProviderBilling = "provider_billing"
|
||||
A2AProvider = "a2a_provider"
|
||||
AgentHealthCheck = "agent_health_check"
|
||||
AgentKillSwitch = "agent_kill_switch"
|
||||
|
|
|
|||
54
litellm/types/proxy/spend_capture_rate.py
Normal file
54
litellm/types/proxy/spend_capture_rate.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
"""The captured-spend to provider-bill ratio: the share of a provider's bill that went through LiteLLM and was priced.
|
||||
|
||||
``capture_rate = captured_spend / provider_spend`` over the same UTC days. 1.0 means LiteLLM saw and priced every
|
||||
dollar the provider billed, lower means traffic reaches the provider outside LiteLLM or cost tracking drops spend,
|
||||
higher means LiteLLM prices above the bill. ``None`` means the provider billed nothing, so there is no ratio.
|
||||
"""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from litellm.constants import SPEND_CAPTURE_RATE_MAX_RANGE_DAYS
|
||||
|
||||
SpendCaptureProvider = Literal["openai"]
|
||||
|
||||
|
||||
class SpendCaptureRateCheckSettings(BaseModel):
|
||||
"""``general_settings.spend_capture_rate_check``: the daily check of captured spend against the provider bill."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
providers: tuple[SpendCaptureProvider, ...] = Field(("openai",), min_length=1)
|
||||
threshold: float = Field(0.9, gt=0, le=1)
|
||||
lookback_days: int = Field(7, ge=1, le=SPEND_CAPTURE_RATE_MAX_RANGE_DAYS)
|
||||
openai_project_ids: tuple[str, ...] = Field(
|
||||
(),
|
||||
description=(
|
||||
"Scope the OpenAI bill to these project ids; empty compares against the whole organization. Captured "
|
||||
"spend is never scoped, so list every project LiteLLM's OpenAI keys belong to"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CaptureRateDay(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
date: str
|
||||
captured_spend: float
|
||||
provider_spend: float
|
||||
capture_rate: float | None
|
||||
|
||||
|
||||
class CaptureRateReport(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
provider: SpendCaptureProvider
|
||||
start_date: str
|
||||
end_date: str
|
||||
captured_spend: float
|
||||
provider_spend: float
|
||||
capture_rate: float | None
|
||||
threshold: float
|
||||
below_threshold: bool
|
||||
days: tuple[CaptureRateDay, ...]
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
import math
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from prometheus_client import REGISTRY
|
||||
from prometheus_client.samples import Sample
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
METRIC: Final = "litellm_spend_capture_rate"
|
||||
|
||||
|
||||
def _clear_prometheus_registry() -> None:
|
||||
for collector in list(REGISTRY._collector_to_names.keys()):
|
||||
try:
|
||||
REGISTRY.unregister(collector)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _samples(metric_name: str) -> list[Sample]:
|
||||
return [sample for metric in REGISTRY.collect() for sample in metric.samples if sample.name == metric_name]
|
||||
|
||||
|
||||
def test_capture_rate_gauge_holds_the_latest_rate_per_provider_and_nan_when_there_is_none() -> None:
|
||||
_clear_prometheus_registry()
|
||||
try:
|
||||
logger: Final = PrometheusLogger()
|
||||
assert _samples(METRIC) == []
|
||||
|
||||
logger.set_spend_capture_rate(api_provider="openai", capture_rate=0.87)
|
||||
logger.set_spend_capture_rate(api_provider="openai", capture_rate=0.91)
|
||||
|
||||
samples: Final = _samples(METRIC)
|
||||
assert [(sample.labels, sample.value) for sample in samples] == [({"api_provider": "openai"}, 0.91)]
|
||||
|
||||
logger.set_spend_capture_rate(api_provider="openai", capture_rate=None)
|
||||
|
||||
(unavailable,) = _samples(METRIC)
|
||||
assert unavailable.labels == {"api_provider": "openai"} and math.isnan(unavailable.value)
|
||||
finally:
|
||||
_clear_prometheus_registry()
|
||||
|
||||
|
||||
def test_capture_rate_gauge_still_records_when_api_provider_is_an_excluded_label(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "prometheus_exclude_labels", ["api_provider"])
|
||||
_clear_prometheus_registry()
|
||||
try:
|
||||
logger: Final = PrometheusLogger()
|
||||
|
||||
logger.set_spend_capture_rate(api_provider="openai", capture_rate=0.42)
|
||||
|
||||
assert [(sample.labels, sample.value) for sample in _samples(METRIC)] == [({}, 0.42)]
|
||||
|
||||
logger.set_spend_capture_rate(api_provider="openai", capture_rate=None)
|
||||
|
||||
(unavailable,) = _samples(METRIC)
|
||||
assert unavailable.labels == {} and math.isnan(unavailable.value)
|
||||
finally:
|
||||
_clear_prometheus_registry()
|
||||
112
tests/test_litellm/llms/openai/test_organization_costs.py
Normal file
112
tests/test_litellm/llms/openai/test_organization_costs.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
from collections.abc import Mapping
|
||||
from datetime import date, datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.constants import OPENAI_ORGANIZATION_COSTS_URL
|
||||
from litellm.llms.openai.organization_costs import OpenAICostsRequestFailed, fetch_openai_daily_costs
|
||||
|
||||
_ADMIN_KEY: Final = "sk-admin-test"
|
||||
|
||||
|
||||
def _utc_midnight(day: str) -> int:
|
||||
return int(datetime.fromisoformat(day).replace(tzinfo=timezone.utc).timestamp())
|
||||
|
||||
|
||||
def _bucket(day: str, *amounts: float) -> dict[str, object]:
|
||||
return {
|
||||
"object": "bucket",
|
||||
"start_time": _utc_midnight(day),
|
||||
"end_time": _utc_midnight(day) + 86400,
|
||||
"results": [
|
||||
{"object": "organization.costs.result", "amount": {"value": amount, "currency": "usd"}, "line_item": None}
|
||||
for amount in amounts
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class _FakeCostsApi:
|
||||
"""Serves ``pages`` in order and records every request it saw."""
|
||||
|
||||
def __init__(self, *pages: dict[str, object] | httpx.Response | Exception) -> None:
|
||||
self._pages = list(pages)
|
||||
self.calls: list[tuple[str, Mapping[str, object], Mapping[str, str]]] = []
|
||||
|
||||
async def __call__(self, url: str, params: Mapping[str, object], headers: Mapping[str, str]) -> httpx.Response:
|
||||
self.calls.append((url, dict(params), dict(headers)))
|
||||
page = self._pages.pop(0)
|
||||
if isinstance(page, Exception):
|
||||
raise page
|
||||
if isinstance(page, httpx.Response):
|
||||
return page
|
||||
return httpx.Response(200, json=page)
|
||||
|
||||
|
||||
def _page(*buckets: dict[str, object], next_page: str | None = None) -> dict[str, object]:
|
||||
return {"object": "page", "data": list(buckets), "has_more": next_page is not None, "next_page": next_page}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_costs_are_summed_per_utc_day_across_pages_and_line_items():
|
||||
api = _FakeCostsApi(
|
||||
_page(_bucket("2026-09-20", 10.0, 2.5), _bucket("2026-09-21", 4.0), next_page="page_2"),
|
||||
_page(_bucket("2026-09-22", 1.0)),
|
||||
)
|
||||
|
||||
billed = await fetch_openai_daily_costs(date(2026, 9, 20), date(2026, 9, 22), admin_key=_ADMIN_KEY, http_get=api)
|
||||
|
||||
assert dict(billed) == {"2026-09-20": 12.5, "2026-09-21": 4.0, "2026-09-22": 1.0}
|
||||
first, second = api.calls
|
||||
assert first[0] == OPENAI_ORGANIZATION_COSTS_URL
|
||||
assert first[2] == {"Authorization": f"Bearer {_ADMIN_KEY}"}
|
||||
assert first[1]["start_time"] == _utc_midnight("2026-09-20")
|
||||
assert first[1]["end_time"] == _utc_midnight("2026-09-23")
|
||||
assert first[1]["bucket_width"] == "1d"
|
||||
assert "page" not in first[1]
|
||||
assert "project_ids[]" not in first[1]
|
||||
assert second[1] == {**first[1], "page": "page_2"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_costs_are_scoped_to_the_configured_projects():
|
||||
api = _FakeCostsApi(_page())
|
||||
|
||||
billed = await fetch_openai_daily_costs(
|
||||
date(2026, 9, 20), date(2026, 9, 20), admin_key=_ADMIN_KEY, project_ids=("proj_a", "proj_b"), http_get=api
|
||||
)
|
||||
|
||||
assert dict(billed) == {}
|
||||
assert api.calls[0][1]["project_ids[]"] == ("proj_a", "proj_b")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"page, detail_fragment",
|
||||
[
|
||||
(httpx.Response(401, json={"error": {"message": "Incorrect API key provided"}}), "HTTP 401"),
|
||||
(httpx.ConnectError("connection refused"), "request failed"),
|
||||
({"object": "page", "data": [{"start_time": "not-a-timestamp"}]}, "unexpected response shape"),
|
||||
],
|
||||
)
|
||||
async def test_an_unreadable_openai_bill_is_a_request_failure_not_an_exception(page, detail_fragment):
|
||||
api = _FakeCostsApi(page)
|
||||
|
||||
billed = await fetch_openai_daily_costs(date(2026, 9, 20), date(2026, 9, 20), admin_key=_ADMIN_KEY, http_get=api)
|
||||
|
||||
assert isinstance(billed, OpenAICostsRequestFailed)
|
||||
assert detail_fragment in billed.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failure_on_a_later_page_fails_the_whole_read():
|
||||
api = _FakeCostsApi(
|
||||
_page(_bucket("2026-09-20", 10.0), next_page="page_2"),
|
||||
httpx.Response(429, json={"error": {"message": "rate limited"}}),
|
||||
)
|
||||
|
||||
billed = await fetch_openai_daily_costs(date(2026, 9, 20), date(2026, 9, 21), admin_key=_ADMIN_KEY, http_get=api)
|
||||
|
||||
assert isinstance(billed, OpenAICostsRequestFailed)
|
||||
assert "HTTP 429" in billed.detail
|
||||
|
|
@ -0,0 +1,439 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import date, datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg.rows import dict_row
|
||||
from pydantic import ValidationError
|
||||
from pytest_postgresql import factories
|
||||
|
||||
from litellm.constants import (
|
||||
SPEND_CAPTURE_RATE_CHECK_JOB_ID,
|
||||
SPEND_CAPTURE_RATE_DOCS_URL,
|
||||
SPEND_CAPTURE_RATE_MAX_RANGE_DAYS,
|
||||
)
|
||||
from litellm.llms.openai.organization_costs import OPENAI_ADMIN_KEY_ENV_VAR
|
||||
from litellm.proxy.spend_tracking.spend_capture_rate import (
|
||||
ProviderBillingCredentialMissing,
|
||||
ProviderBillingRequestFailed,
|
||||
alert_message,
|
||||
captured_spend_by_day,
|
||||
compute_capture_rate,
|
||||
run_scheduled_spend_capture_rate_check,
|
||||
run_spend_capture_rate_check,
|
||||
)
|
||||
from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureRateCheckSettings
|
||||
|
||||
_ADMIN_KEY: Final = "sk-admin-test"
|
||||
|
||||
|
||||
def _utc_midnight(day: str) -> int:
|
||||
return int(datetime.fromisoformat(day).replace(tzinfo=timezone.utc).timestamp())
|
||||
|
||||
|
||||
def _bucket(day: str, *amounts: float) -> dict[str, object]:
|
||||
return {
|
||||
"object": "bucket",
|
||||
"start_time": _utc_midnight(day),
|
||||
"end_time": _utc_midnight(day) + 86400,
|
||||
"results": [
|
||||
{"object": "organization.costs.result", "amount": {"value": amount, "currency": "usd"}, "line_item": None}
|
||||
for amount in amounts
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class _FakeCostsApi:
|
||||
"""Serves ``pages`` in order and records every request it saw."""
|
||||
|
||||
def __init__(self, *pages: dict[str, object] | httpx.Response) -> None:
|
||||
self.responses = [
|
||||
page if isinstance(page, httpx.Response) else httpx.Response(200, json=page) for page in pages
|
||||
]
|
||||
self.calls: list[tuple[str, Mapping[str, object], Mapping[str, str]]] = []
|
||||
|
||||
async def __call__(self, url: str, params: Mapping[str, object], headers: Mapping[str, str]) -> httpx.Response:
|
||||
self.calls.append((url, dict(params), dict(headers)))
|
||||
return self.responses[len(self.calls) - 1]
|
||||
|
||||
|
||||
def _page(*buckets: dict[str, object], next_page: str | None = None) -> dict[str, object]:
|
||||
return {"object": "page", "data": list(buckets), "has_more": next_page is not None, "next_page": next_page}
|
||||
|
||||
|
||||
def _fake_prisma(rows: list[dict[str, object]]) -> MagicMock:
|
||||
prisma = MagicMock()
|
||||
prisma.db.query_raw = AsyncMock(return_value=rows)
|
||||
return prisma
|
||||
|
||||
|
||||
def test_capture_rate_covers_every_day_in_the_range_and_flags_the_threshold():
|
||||
report = compute_capture_rate(
|
||||
provider="openai",
|
||||
start_date=date(2026, 9, 20),
|
||||
end_date=date(2026, 9, 22),
|
||||
captured_by_day={"2026-09-20": 8.0, "2026-09-22": 1.0},
|
||||
billed_by_day={"2026-09-20": 10.0, "2026-09-21": 5.0},
|
||||
threshold=0.9,
|
||||
)
|
||||
|
||||
assert [day.date for day in report.days] == ["2026-09-20", "2026-09-21", "2026-09-22"]
|
||||
assert [day.capture_rate for day in report.days] == [0.8, 0.0, None]
|
||||
assert report.captured_spend == 9.0
|
||||
assert report.provider_spend == 15.0
|
||||
assert report.capture_rate == 0.6
|
||||
assert report.below_threshold is True
|
||||
|
||||
|
||||
def test_capture_rate_at_or_above_the_threshold_is_not_flagged_and_a_zero_bill_has_no_rate():
|
||||
healthy = compute_capture_rate(
|
||||
provider="openai",
|
||||
start_date=date(2026, 9, 20),
|
||||
end_date=date(2026, 9, 20),
|
||||
captured_by_day={"2026-09-20": 9.5},
|
||||
billed_by_day={"2026-09-20": 10.0},
|
||||
threshold=0.9,
|
||||
)
|
||||
over = compute_capture_rate(
|
||||
provider="openai",
|
||||
start_date=date(2026, 9, 20),
|
||||
end_date=date(2026, 9, 20),
|
||||
captured_by_day={"2026-09-20": 12.0},
|
||||
billed_by_day={"2026-09-20": 10.0},
|
||||
threshold=0.9,
|
||||
)
|
||||
unbilled = compute_capture_rate(
|
||||
provider="openai",
|
||||
start_date=date(2026, 9, 20),
|
||||
end_date=date(2026, 9, 20),
|
||||
captured_by_day={"2026-09-20": 3.0},
|
||||
billed_by_day={},
|
||||
threshold=0.9,
|
||||
)
|
||||
|
||||
assert (healthy.capture_rate, healthy.below_threshold) == (0.95, False)
|
||||
assert (over.capture_rate, over.below_threshold) == (1.2, False)
|
||||
assert (unbilled.capture_rate, unbilled.below_threshold) == (None, False)
|
||||
assert alert_message(healthy) is None
|
||||
assert alert_message(over) is None
|
||||
assert alert_message(unbilled) is None
|
||||
|
||||
|
||||
def test_alert_messages_name_the_cause_and_link_the_docs():
|
||||
below = compute_capture_rate(
|
||||
provider="openai",
|
||||
start_date=date(2026, 9, 14),
|
||||
end_date=date(2026, 9, 20),
|
||||
captured_by_day={"2026-09-14": 700.0},
|
||||
billed_by_day={"2026-09-14": 1000.0},
|
||||
threshold=0.9,
|
||||
)
|
||||
|
||||
below_message = alert_message(below)
|
||||
missing_message = alert_message(ProviderBillingCredentialMissing("openai", OPENAI_ADMIN_KEY_ENV_VAR))
|
||||
failed_message = alert_message(ProviderBillingRequestFailed("openai", "HTTP 401: nope"))
|
||||
|
||||
assert below_message is not None and "70.0%" in below_message and "90%" in below_message
|
||||
assert "$700.00" in below_message and "$1,000.00" in below_message
|
||||
assert "2026-09-14 to 2026-09-20" in below_message
|
||||
assert missing_message is not None and OPENAI_ADMIN_KEY_ENV_VAR in missing_message
|
||||
assert failed_message is not None and "HTTP 401: nope" in failed_message
|
||||
assert all(SPEND_CAPTURE_RATE_DOCS_URL in m for m in (below_message, missing_message, failed_message))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_reads_the_closed_window_before_today_and_publishes_the_rate(monkeypatch):
|
||||
monkeypatch.setenv(OPENAI_ADMIN_KEY_ENV_VAR, _ADMIN_KEY)
|
||||
api = _FakeCostsApi(_page(_bucket("2026-09-21", 100.0), _bucket("2026-09-22", 100.0)))
|
||||
prisma = _fake_prisma([{"date": "2026-09-21", "spend": 95.0}, {"date": "2026-09-22", "spend": 91.0}])
|
||||
alert = AsyncMock()
|
||||
publish = MagicMock()
|
||||
|
||||
results = await run_spend_capture_rate_check(
|
||||
prisma,
|
||||
SpendCaptureRateCheckSettings(lookback_days=2, threshold=0.9),
|
||||
alert=alert,
|
||||
publish=publish,
|
||||
today=date(2026, 9, 23),
|
||||
http_get=api,
|
||||
)
|
||||
|
||||
(report,) = results
|
||||
assert isinstance(report, CaptureRateReport)
|
||||
assert (report.start_date, report.end_date) == ("2026-09-21", "2026-09-22")
|
||||
assert report.capture_rate == 0.93
|
||||
publish.assert_called_once_with("openai", 0.93)
|
||||
alert.assert_not_awaited()
|
||||
assert api.calls[0][1]["start_time"] == _utc_midnight("2026-09-21")
|
||||
assert api.calls[0][1]["end_time"] == _utc_midnight("2026-09-23")
|
||||
sql, start, end, providers = prisma.db.query_raw.await_args.args
|
||||
assert (start, end) == ("2026-09-21", "2026-09-22")
|
||||
assert set(providers) == {"openai", "text-completion-openai"}
|
||||
assert '"LiteLLM_DailyUserSpend"' in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_alerts_under_the_threshold_and_still_publishes_the_rate(monkeypatch):
|
||||
monkeypatch.setenv(OPENAI_ADMIN_KEY_ENV_VAR, _ADMIN_KEY)
|
||||
api = _FakeCostsApi(_page(_bucket("2026-09-22", 200.0)))
|
||||
prisma = _fake_prisma([{"date": "2026-09-22", "spend": 50.0}])
|
||||
alert = AsyncMock()
|
||||
publish = MagicMock()
|
||||
|
||||
await run_spend_capture_rate_check(
|
||||
prisma,
|
||||
SpendCaptureRateCheckSettings(lookback_days=1),
|
||||
alert=alert,
|
||||
publish=publish,
|
||||
today=date(2026, 9, 23),
|
||||
http_get=api,
|
||||
)
|
||||
|
||||
publish.assert_called_once_with("openai", 0.25)
|
||||
alert.assert_awaited_once()
|
||||
assert "25.0%" in alert.await_args.args[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_alerts_on_a_missing_admin_key_and_publishes_nothing(monkeypatch):
|
||||
monkeypatch.delenv(OPENAI_ADMIN_KEY_ENV_VAR, raising=False)
|
||||
api = _FakeCostsApi()
|
||||
prisma = _fake_prisma([])
|
||||
alert = AsyncMock()
|
||||
publish = MagicMock()
|
||||
|
||||
(result,) = await run_spend_capture_rate_check(
|
||||
prisma, SpendCaptureRateCheckSettings(), alert=alert, publish=publish, today=date(2026, 9, 23), http_get=api
|
||||
)
|
||||
|
||||
assert result == ProviderBillingCredentialMissing("openai", OPENAI_ADMIN_KEY_ENV_VAR)
|
||||
publish.assert_called_once_with("openai", None)
|
||||
alert.assert_awaited_once()
|
||||
assert api.calls == []
|
||||
prisma.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_alerts_on_an_unreadable_bill_and_publishes_no_rate(monkeypatch):
|
||||
monkeypatch.setenv(OPENAI_ADMIN_KEY_ENV_VAR, _ADMIN_KEY)
|
||||
api = _FakeCostsApi(httpx.Response(401, json={"error": {"message": "Incorrect API key provided"}}))
|
||||
prisma = _fake_prisma([{"date": "2026-09-22", "spend": 5.0}])
|
||||
alert = AsyncMock()
|
||||
publish = MagicMock()
|
||||
|
||||
(result,) = await run_spend_capture_rate_check(
|
||||
prisma, SpendCaptureRateCheckSettings(), alert=alert, publish=publish, today=date(2026, 9, 23), http_get=api
|
||||
)
|
||||
|
||||
assert result == ProviderBillingRequestFailed("openai", "HTTP 401: " + api.responses[0].text[:300])
|
||||
publish.assert_called_once_with("openai", None)
|
||||
alert.assert_awaited_once()
|
||||
assert "HTTP 401" in alert.await_args.args[0]
|
||||
prisma.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_publishes_no_rate_when_the_provider_billed_nothing(monkeypatch):
|
||||
monkeypatch.setenv(OPENAI_ADMIN_KEY_ENV_VAR, _ADMIN_KEY)
|
||||
api = _FakeCostsApi(_page())
|
||||
prisma = _fake_prisma([{"date": "2026-09-22", "spend": 5.0}])
|
||||
publish = MagicMock()
|
||||
alert = AsyncMock()
|
||||
|
||||
(report,) = await run_spend_capture_rate_check(
|
||||
prisma,
|
||||
SpendCaptureRateCheckSettings(lookback_days=1),
|
||||
alert=alert,
|
||||
publish=publish,
|
||||
today=date(2026, 9, 23),
|
||||
http_get=api,
|
||||
)
|
||||
|
||||
assert isinstance(report, CaptureRateReport) and report.capture_rate is None
|
||||
publish.assert_called_once_with("openai", None)
|
||||
alert.assert_not_awaited()
|
||||
|
||||
|
||||
def _pod_lock(acquired: bool) -> MagicMock:
|
||||
lock = MagicMock()
|
||||
lock.redis_cache = MagicMock()
|
||||
lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod")
|
||||
lock.get_redis_lock_key = MagicMock(return_value="lock-key")
|
||||
lock.acquire_lock = AsyncMock(return_value=acquired)
|
||||
lock.release_lock = AsyncMock()
|
||||
return lock
|
||||
|
||||
|
||||
async def _scheduled_run(
|
||||
lock: MagicMock, monkeypatch, *, captured: float, prisma: MagicMock | None = None
|
||||
) -> tuple[AsyncMock, MagicMock]:
|
||||
monkeypatch.setenv(OPENAI_ADMIN_KEY_ENV_VAR, _ADMIN_KEY)
|
||||
alert = AsyncMock()
|
||||
publish = MagicMock()
|
||||
await run_scheduled_spend_capture_rate_check(
|
||||
prisma or _fake_prisma([{"date": "2026-09-22", "spend": captured}]),
|
||||
SpendCaptureRateCheckSettings(lookback_days=1),
|
||||
pod_lock_manager=lock,
|
||||
alert=alert,
|
||||
publish=publish,
|
||||
today=date(2026, 9, 23),
|
||||
http_get=_FakeCostsApi(_page(_bucket("2026-09-22", 200.0))),
|
||||
)
|
||||
return alert, publish
|
||||
|
||||
|
||||
async def _scheduled_run_under_threshold(lock: MagicMock, monkeypatch) -> tuple[AsyncMock, MagicMock]:
|
||||
return await _scheduled_run(lock, monkeypatch, captured=50.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_healthy_scheduled_check_publishes_and_never_touches_the_alert_lock(monkeypatch):
|
||||
lock = _pod_lock(acquired=False)
|
||||
|
||||
alert, publish = await _scheduled_run(lock, monkeypatch, captured=190.0)
|
||||
|
||||
publish.assert_called_once_with("openai", 0.95)
|
||||
alert.assert_not_awaited()
|
||||
lock.acquire_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_scheduled_check_that_fails_never_claims_the_alert_window(monkeypatch):
|
||||
lock = _pod_lock(acquired=True)
|
||||
prisma = MagicMock()
|
||||
prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("database gone"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="database gone"):
|
||||
await _scheduled_run(lock, monkeypatch, captured=0.0, prisma=prisma)
|
||||
|
||||
lock.acquire_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_check_publishes_but_stays_quiet_when_another_pod_holds_the_alert_window(monkeypatch):
|
||||
lock = _pod_lock(acquired=False)
|
||||
|
||||
alert, publish = await _scheduled_run_under_threshold(lock, monkeypatch)
|
||||
|
||||
publish.assert_called_once_with("openai", 0.25)
|
||||
alert.assert_not_awaited()
|
||||
lock.release_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_check_alerts_and_keeps_the_lock_until_it_expires_when_it_wins(monkeypatch):
|
||||
lock = _pod_lock(acquired=True)
|
||||
|
||||
alert, publish = await _scheduled_run_under_threshold(lock, monkeypatch)
|
||||
|
||||
lock.acquire_lock.assert_awaited_once_with(cronjob_id=SPEND_CAPTURE_RATE_CHECK_JOB_ID, ttl=900)
|
||||
lock.release_lock.assert_not_awaited()
|
||||
publish.assert_called_once_with("openai", 0.25)
|
||||
alert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_check_alerts_when_the_lock_cannot_be_acquired_or_read(monkeypatch):
|
||||
lock = _pod_lock(acquired=False)
|
||||
lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down"))
|
||||
|
||||
alert, publish = await _scheduled_run_under_threshold(lock, monkeypatch)
|
||||
|
||||
publish.assert_called_once_with("openai", 0.25)
|
||||
alert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_check_alerts_without_a_lock_manager(monkeypatch):
|
||||
lock = _pod_lock(acquired=False)
|
||||
lock.redis_cache = None
|
||||
|
||||
alert, publish = await _scheduled_run_under_threshold(lock, monkeypatch)
|
||||
|
||||
lock.acquire_lock.assert_not_awaited()
|
||||
publish.assert_called_once_with("openai", 0.25)
|
||||
alert.assert_awaited_once()
|
||||
|
||||
|
||||
def test_settings_reject_typos_and_out_of_range_values():
|
||||
with pytest.raises(ValidationError, match="threshhold"):
|
||||
SpendCaptureRateCheckSettings.model_validate({"threshhold": 0.9})
|
||||
with pytest.raises(ValidationError, match="threshold"):
|
||||
SpendCaptureRateCheckSettings.model_validate({"threshold": 1.5})
|
||||
with pytest.raises(ValidationError, match="providers"):
|
||||
SpendCaptureRateCheckSettings.model_validate({"providers": []})
|
||||
with pytest.raises(ValidationError, match="providers"):
|
||||
SpendCaptureRateCheckSettings.model_validate({"providers": ["anthropic"]})
|
||||
with pytest.raises(ValidationError, match="lookback_days"):
|
||||
SpendCaptureRateCheckSettings.model_validate({"lookback_days": SPEND_CAPTURE_RATE_MAX_RANGE_DAYS + 1})
|
||||
parsed = SpendCaptureRateCheckSettings.model_validate(
|
||||
json.loads('{"providers": ["openai"], "threshold": 0.8, "lookback_days": 3, "openai_project_ids": ["p"]}')
|
||||
)
|
||||
assert (parsed.threshold, parsed.lookback_days, parsed.openai_project_ids) == (0.8, 3, ("p",))
|
||||
|
||||
|
||||
_capture_postgresql_proc: Final = factories.postgresql_proc()
|
||||
_capture_postgresql: Final = factories.postgresql("_capture_postgresql_proc")
|
||||
|
||||
_DAILY_USER_SPEND_DDL: Final = """
|
||||
CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
id TEXT PRIMARY KEY,
|
||||
date TEXT NOT NULL,
|
||||
custom_llm_provider TEXT,
|
||||
spend DOUBLE PRECISION DEFAULT 0
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
class _PsycopgPrisma:
|
||||
"""``prisma_client.db.query_raw`` on a real connection, with ``$n`` placeholders converted for psycopg."""
|
||||
|
||||
def __init__(self, conn: psycopg.Connection) -> None:
|
||||
self.db = self
|
||||
self._conn = conn
|
||||
|
||||
async def query_raw(self, sql: str, *params: object) -> list[dict[str, object]]:
|
||||
converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql)
|
||||
with self._conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query
|
||||
{f"p{i}": list(v) if isinstance(v, tuple) else v for i, v in enumerate(params, start=1)},
|
||||
)
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_captured_spend_sums_only_the_openai_billed_providers_inside_the_window(
|
||||
_capture_postgresql: psycopg.Connection,
|
||||
):
|
||||
conn: Final = _capture_postgresql
|
||||
conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal
|
||||
rows: Final = (
|
||||
("2026-09-19", "openai", 1.0),
|
||||
("2026-09-20", "openai", 2.0),
|
||||
("2026-09-20", "openai", 3.0),
|
||||
("2026-09-20", "text-completion-openai", 0.5),
|
||||
("2026-09-20", "anthropic", 100.0),
|
||||
("2026-09-21", "azure", 100.0),
|
||||
("2026-09-22", "openai", 4.0),
|
||||
)
|
||||
for index, (day, provider, spend) in enumerate(rows):
|
||||
conn.execute(
|
||||
'INSERT INTO "LiteLLM_DailyUserSpend" (id, date, custom_llm_provider, spend) VALUES (%s, %s, %s, %s)',
|
||||
(f"row-{index}", day, provider, spend),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
captured = await captured_spend_by_day(
|
||||
_PsycopgPrisma(conn), # pyright: ignore[reportArgumentType] # duck-typed prisma for the raw query
|
||||
litellm_providers=("openai", "text-completion-openai"),
|
||||
start_date=date(2026, 9, 20),
|
||||
end_date=date(2026, 9, 21),
|
||||
)
|
||||
|
||||
assert dict(captured) == {"2026-09-20": 5.5}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -14,7 +14,7 @@ from datetime import datetime, timedelta, timezone
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock, MagicMock, create_autospec, mock_open, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch
|
||||
|
||||
import click
|
||||
import fastapi.routing
|
||||
|
|
@ -3524,9 +3524,7 @@ async def test_load_config_without_role_permissions_leaves_every_role_unrestrict
|
|||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
config_file: Final = tmp_path / "config.yaml"
|
||||
config_file.write_text(
|
||||
yaml.dump({"model_list": [], "general_settings": {"max_parallel_requests": 7}})
|
||||
)
|
||||
config_file.write_text(yaml.dump({"model_list": [], "general_settings": {"max_parallel_requests": 7}}))
|
||||
|
||||
_, _, settings = await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
|
||||
|
||||
|
|
@ -3556,7 +3554,9 @@ async def test_load_config_rejects_malformed_role_permissions(tmp_path):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_compiles_key_alias_pattern_at_startup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_load_config_compiles_key_alias_pattern_at_startup(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
monkeypatch.setattr(litellm, "key_alias_pattern", None)
|
||||
|
|
@ -3592,9 +3592,7 @@ def test_os_environ_resolution_leaves_the_config_layer_holding_the_reference(mon
|
|||
assert resolved["general_settings"]["coordination_redis"]["password"] == "sk-nested-value"
|
||||
assert resolved["general_settings"]["master_key"] == "sk-nested-value"
|
||||
assert proxy_config.settings.config_value("master_key") == "os.environ/PROOF_NESTED_SECRET"
|
||||
assert proxy_config.settings.config_value("coordination_redis") == {
|
||||
"password": "os.environ/PROOF_NESTED_SECRET"
|
||||
}
|
||||
assert proxy_config.settings.config_value("coordination_redis") == {"password": "os.environ/PROOF_NESTED_SECRET"}
|
||||
|
||||
|
||||
def test_os_environ_resolution_reaches_dicts_nested_in_a_list(monkeypatch):
|
||||
|
|
@ -5754,7 +5752,9 @@ async def test_boot_warns_that_a_shadowed_database_value_will_never_apply(tmp_pa
|
|||
config_path.write_text(
|
||||
yaml.safe_dump({"model_list": [], "general_settings": {"allowed_ips": ["1.2.3.4"], "max_file_size_mb": 5}})
|
||||
)
|
||||
db_row: Final = types.SimpleNamespace(param_value={"allowed_ips": ["1.2.3.4", "5.6.7.8"], "max_parallel_requests": 7})
|
||||
db_row: Final = types.SimpleNamespace(
|
||||
param_value={"allowed_ips": ["1.2.3.4", "5.6.7.8"], "max_parallel_requests": 7}
|
||||
)
|
||||
|
||||
async def read_config_row(_prisma_client, param_name):
|
||||
return db_row if param_name == "general_settings" else None
|
||||
|
|
@ -7957,7 +7957,9 @@ async def test_update_general_settings_clearing_a_db_override_falls_back_to_the_
|
|||
proxy_config.settings.load_yaml({"maximum_spend_logs_cleanup_run_budget": "90s"})
|
||||
|
||||
with patch("litellm.proxy.proxy_server.general_settings", proxy_config.settings):
|
||||
await proxy_config._update_general_settings(db_general_settings={"maximum_spend_logs_cleanup_run_budget": "30s"})
|
||||
await proxy_config._update_general_settings(
|
||||
db_general_settings={"maximum_spend_logs_cleanup_run_budget": "30s"}
|
||||
)
|
||||
await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
|
@ -8000,10 +8002,18 @@ async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to
|
|||
request.query_params = {}
|
||||
return request
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
|
||||
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch(
|
||||
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
|
||||
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
with settings, yaml_endpoints, initialize, master_key:
|
||||
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
|
||||
|
||||
|
|
@ -8050,10 +8060,18 @@ async def test_update_general_settings_db_pass_through_endpoint_cannot_override_
|
|||
request.headers = {}
|
||||
request.query_params = {}
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
|
||||
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch(
|
||||
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
|
||||
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
with settings, yaml_endpoints, initialize, master_key:
|
||||
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
|
||||
|
||||
|
|
@ -8087,11 +8105,19 @@ async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_servi
|
|||
prior_registry: Final = dict(_registered_pass_through_routes)
|
||||
|
||||
def live_routes() -> set[str]:
|
||||
return {route for route in InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() if path in route}
|
||||
return {
|
||||
route for route in InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() if path in route
|
||||
}
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", None) # test-quality-ok: module global holding the YAML endpoints; this case has none
|
||||
app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", None
|
||||
) # test-quality-ok: module global holding the YAML endpoints; this case has none
|
||||
app_routes: Final = patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists"
|
||||
) # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
try:
|
||||
with settings, yaml_endpoints, app_routes:
|
||||
pc = ProxyConfig()
|
||||
|
|
@ -8132,9 +8158,15 @@ async def test_a_stored_pass_through_row_never_disturbs_the_config_declared_rout
|
|||
registered: Final = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
|
||||
return {path for path in (config_path, db_path) if any(path in route for route in registered)}
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]) # test-quality-ok: module global holding the YAML endpoints the reload merges in
|
||||
app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]
|
||||
) # test-quality-ok: module global holding the YAML endpoints the reload merges in
|
||||
app_routes: Final = patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists"
|
||||
) # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
try:
|
||||
with settings, yaml_endpoints, app_routes:
|
||||
await initialize_pass_through_endpoints(pass_through_endpoints=[config_endpoint])
|
||||
|
|
@ -9263,9 +9295,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments():
|
|||
async def assert_reservation_not_finalized_yet(**kwargs):
|
||||
assert budget_reservation["finalized"] is False
|
||||
incremented_counters.append(kwargs["counter_key"])
|
||||
return ps.PendingSpendIncrement(
|
||||
counter_key=kwargs["counter_key"], increment=kwargs["increment"]
|
||||
)
|
||||
return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"])
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
|
|
@ -10854,9 +10884,15 @@ async def _lit6973_drive_realtime_session(
|
|||
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
|
||||
)
|
||||
ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws()
|
||||
can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
|
||||
pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
|
||||
route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
|
||||
can_call = patch.object(
|
||||
ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)
|
||||
) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
|
||||
pre = patch.object(
|
||||
ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call
|
||||
) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
|
||||
route = patch.object(
|
||||
ps, "route_request", new=AsyncMock(return_value=fake_llm_call())
|
||||
) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
|
||||
with can_call, pre, route:
|
||||
await ps.realtime_websocket_endpoint(
|
||||
websocket=ws,
|
||||
|
|
@ -10988,13 +11024,9 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
|
|||
from litellm.proxy.utils import InternalUsageCache
|
||||
|
||||
dual_cache: Final = DualCache()
|
||||
await dual_cache.async_set_cache(
|
||||
key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True
|
||||
)
|
||||
await dual_cache.async_set_cache(key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True)
|
||||
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
|
||||
stash: Final = RequestRateLimiterStash(
|
||||
parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
|
||||
)
|
||||
stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]})
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
stash_token: Final = _request_stash.set(stash)
|
||||
|
|
@ -11046,9 +11078,7 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_
|
|||
limiter's integer in-memory fallback, double-decrement the counter so the key
|
||||
admits more sessions than max_parallel_requests allows. With the success stamp
|
||||
present the route leaves the slot and the stash alone."""
|
||||
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
|
||||
backend_logged_success=True
|
||||
)
|
||||
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(backend_logged_success=True)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
|
||||
"slot-1": 1.0,
|
||||
|
|
@ -11094,8 +11124,12 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
|
|||
async def _record(counter_key: str) -> None:
|
||||
invalidated.append(counter_key)
|
||||
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
|
||||
sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
|
||||
failing_release = patch.object(
|
||||
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
|
||||
) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
|
||||
sink = patch.object(
|
||||
ps, "_invalidate_spend_counter", new=_record
|
||||
) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
|
||||
with failing_release, sink:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
||||
|
|
@ -11111,8 +11145,12 @@ async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback
|
|||
from litellm.proxy.spend_tracking import budget_reservation as br
|
||||
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch
|
||||
failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail
|
||||
failing_release = patch.object(
|
||||
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
|
||||
) # test-quality-ok: forces the fallback branch
|
||||
failing_invalidate = patch.object(
|
||||
br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))
|
||||
) # test-quality-ok: forces the fallback itself to fail
|
||||
|
||||
with failing_release, failing_invalidate:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
|
@ -12470,9 +12508,7 @@ async def test_update_config_general_settings_refuses_a_key_the_config_file_decl
|
|||
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="max_parallel_requests", field_value=999, config_type="general_settings"
|
||||
),
|
||||
data=ConfigFieldUpdate(field_name="max_parallel_requests", field_value=999, config_type="general_settings"),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
|
|
@ -14038,9 +14074,15 @@ async def test_moderations_response_carries_litellm_call_id_header():
|
|||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
|
||||
patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())
|
||||
), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj"
|
||||
) as mock_logging, # test-quality-ok: module global, no injection point
|
||||
):
|
||||
mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
|
||||
mock_logging.update_request_status = AsyncMock()
|
||||
|
|
@ -14077,9 +14119,15 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo
|
|||
verbose_proxy_logger.propagate = True
|
||||
try:
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))
|
||||
), # test-quality-ok: fakes the provider failure so the real route's error log is observable
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
|
|
@ -14112,7 +14160,9 @@ async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id():
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
await proxy_server_module.moderations(
|
||||
|
|
@ -14140,8 +14190,12 @@ async def test_moderations_already_shaped_failure_answers_with_the_callers_litel
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
await proxy_server_module.moderations(
|
||||
|
|
@ -14176,8 +14230,12 @@ async def test_audio_speech_already_shaped_failure_answers_with_the_callers_lite
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(type(exc)) as raised,
|
||||
):
|
||||
await proxy_server_module.audio_speech(
|
||||
|
|
@ -14955,14 +15013,18 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
|
|||
{
|
||||
"model_name": "self-hosted",
|
||||
"litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
|
||||
"model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}},
|
||||
"model_info": {
|
||||
"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}
|
||||
},
|
||||
}
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
response, took, lags = await timed_with_loop_lags(
|
||||
lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop"))
|
||||
lambda: proxy_server_module.token_counter(
|
||||
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
|
||||
)
|
||||
)
|
||||
|
||||
assert response.tokenizer_type == "huggingface_tokenizer"
|
||||
|
|
@ -15186,3 +15248,123 @@ async def test_initialize_jwt_auth_leaves_the_declared_jwtauth_mapping_unresolve
|
|||
|
||||
assert declared["team_id_jwt_field"] == "os.environ/JWT_TEAM_FIELD"
|
||||
assert proxy_server_module.jwt_handler.litellm_jwtauth.team_id_jwt_field == "resolved-team-field"
|
||||
|
||||
|
||||
def test_spend_capture_rate_check_job_validates_the_boot_settings_and_reads_them_again_on_every_run(monkeypatch):
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.constants import SPEND_CAPTURE_RATE_CHECK_JOB_ID
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
|
||||
scheduler = MagicMock()
|
||||
general_settings: dict[str, object] = {}
|
||||
seen_settings = []
|
||||
|
||||
async def fake_scheduled_check(prisma_client, settings, *, pod_lock_manager, alert, publish):
|
||||
seen_settings.append(settings)
|
||||
return ()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.run_scheduled_spend_capture_rate_check", fake_scheduled_check)
|
||||
ProxyStartupEvent._initialize_spend_capture_rate_check_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
prisma_client=MagicMock(),
|
||||
read_general_settings=lambda: general_settings,
|
||||
)
|
||||
scheduler.add_job.assert_called_once()
|
||||
assert scheduler.add_job.call_args.kwargs["id"] == SPEND_CAPTURE_RATE_CHECK_JOB_ID
|
||||
check = scheduler.add_job.call_args.args[0]
|
||||
|
||||
asyncio.run(check())
|
||||
assert seen_settings == []
|
||||
|
||||
general_settings["spend_capture_rate_check"] = {"providers": ["openai"], "threshold": 0.85}
|
||||
asyncio.run(check())
|
||||
general_settings["spend_capture_rate_check"] = {"threshold": 0.7, "lookback_days": 3}
|
||||
asyncio.run(check())
|
||||
assert [(s.threshold, s.lookback_days) for s in seen_settings] == [(0.85, 7), (0.7, 3)]
|
||||
|
||||
with pytest.raises(ValidationError, match="threshhold"):
|
||||
ProxyStartupEvent._initialize_spend_capture_rate_check_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
prisma_client=MagicMock(),
|
||||
read_general_settings=lambda: {"spend_capture_rate_check": {"threshhold": 0.85}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_capture_rate_check_job_publishes_to_prometheus_and_alerts(monkeypatch):
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
|
||||
scheduler = MagicMock()
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.alerting_handler = AsyncMock()
|
||||
proxy_logging.db_spend_update_writer.pod_lock_manager = None
|
||||
prometheus = MagicMock(spec=PrometheusLogger)
|
||||
monkeypatch.setattr(
|
||||
litellm.logging_callback_manager, "get_custom_loggers_for_type", lambda callback_type: [prometheus]
|
||||
)
|
||||
|
||||
async def fake_scheduled_check(prisma_client, settings, *, pod_lock_manager, alert, publish):
|
||||
publish("openai", 0.42)
|
||||
publish("openai", None)
|
||||
await alert("under the threshold")
|
||||
return ()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.run_scheduled_spend_capture_rate_check", fake_scheduled_check)
|
||||
ProxyStartupEvent._initialize_spend_capture_rate_check_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
prisma_client=MagicMock(),
|
||||
read_general_settings=lambda: {"spend_capture_rate_check": {}},
|
||||
)
|
||||
|
||||
await scheduler.add_job.call_args.args[0]()
|
||||
|
||||
assert prometheus.set_spend_capture_rate.call_args_list == [
|
||||
call(api_provider="openai", capture_rate=0.42),
|
||||
call(api_provider="openai", capture_rate=None),
|
||||
]
|
||||
proxy_logging.alerting_handler.assert_awaited_once()
|
||||
assert proxy_logging.alerting_handler.await_args.kwargs["message"] == "under the threshold"
|
||||
assert proxy_logging.alerting_handler.await_args.kwargs["level"] == "High"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is_removed(monkeypatch):
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
|
||||
scheduler = MagicMock()
|
||||
general_settings: dict[str, object] = {"spend_capture_rate_check": {}}
|
||||
prometheus = MagicMock(spec=PrometheusLogger)
|
||||
monkeypatch.setattr(
|
||||
litellm.logging_callback_manager, "get_custom_loggers_for_type", lambda callback_type: [prometheus]
|
||||
)
|
||||
scheduled_checks = []
|
||||
|
||||
async def fake_scheduled_check(prisma_client, settings, *, pod_lock_manager, alert, publish):
|
||||
scheduled_checks.append(settings)
|
||||
publish("openai", 0.97)
|
||||
return ()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.run_scheduled_spend_capture_rate_check", fake_scheduled_check)
|
||||
ProxyStartupEvent._initialize_spend_capture_rate_check_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
prisma_client=MagicMock(),
|
||||
read_general_settings=lambda: general_settings,
|
||||
)
|
||||
check = scheduler.add_job.call_args.args[0]
|
||||
|
||||
await check()
|
||||
del general_settings["spend_capture_rate_check"]
|
||||
await check()
|
||||
|
||||
assert len(scheduled_checks) == 1
|
||||
assert prometheus.set_spend_capture_rate.call_args_list == [
|
||||
call(api_provider="openai", capture_rate=0.97),
|
||||
call(api_provider="openai", capture_rate=None),
|
||||
]
|
||||
|
|
|
|||
134
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
134
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -14615,6 +14615,34 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/spend/capture_rate": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/**
|
||||
* Get Spend Capture Rate
|
||||
* @description Compare the spend LiteLLM captured for a provider against that provider's own bill, per UTC day.
|
||||
*
|
||||
* Admin only. Reads the provider's billing API with the billing credential set on the proxy
|
||||
* (OpenAI: `OPENAI_ADMIN_KEY`) and sums `LiteLLM_DailyUserSpend` for the same days.
|
||||
*
|
||||
* Example:
|
||||
* ```
|
||||
* curl -H "Authorization: Bearer sk-1234" "http://localhost:4000/spend/capture_rate?provider=openai&start_date=2026-09-17&end_date=2026-09-23"
|
||||
* ```
|
||||
*/
|
||||
get: operations["get_spend_capture_rate_spend_capture_rate_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/spend/keys": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -26859,6 +26887,41 @@ export interface components {
|
|||
*/
|
||||
threshold_step: number;
|
||||
};
|
||||
/** CaptureRateDay */
|
||||
CaptureRateDay: {
|
||||
/** Capture Rate */
|
||||
capture_rate: number | null;
|
||||
/** Captured Spend */
|
||||
captured_spend: number;
|
||||
/** Date */
|
||||
date: string;
|
||||
/** Provider Spend */
|
||||
provider_spend: number;
|
||||
};
|
||||
/** CaptureRateReport */
|
||||
CaptureRateReport: {
|
||||
/** Below Threshold */
|
||||
below_threshold: boolean;
|
||||
/** Capture Rate */
|
||||
capture_rate: number | null;
|
||||
/** Captured Spend */
|
||||
captured_spend: number;
|
||||
/** Days */
|
||||
days: components["schemas"]["CaptureRateDay"][];
|
||||
/** End Date */
|
||||
end_date: string;
|
||||
/**
|
||||
* Provider
|
||||
* @constant
|
||||
*/
|
||||
provider: "openai";
|
||||
/** Provider Spend */
|
||||
provider_spend: number;
|
||||
/** Start Date */
|
||||
start_date: string;
|
||||
/** Threshold */
|
||||
threshold: number;
|
||||
};
|
||||
/** ChangePasswordRequest */
|
||||
ChangePasswordRequest: {
|
||||
/** Current Password */
|
||||
|
|
@ -28390,6 +28453,8 @@ export interface components {
|
|||
reject_clientside_metadata_tags?: boolean | null;
|
||||
/** @description Spreads the proxy's scheduled background jobs (spend flushes, budget resets, config reloads, exports) across a window instead of firing them together on every replica. On by default; set to tune the window, pin a job, or turn it off. */
|
||||
scheduled_job_stagger?: components["schemas"]["ScheduledJobStaggerSettings"] | null;
|
||||
/** @description Daily check of the spend LiteLLM captured against the provider's own bill (OpenAI via OPENAI_ADMIN_KEY). Publishes litellm_spend_capture_rate per provider and alerts when the ratio over the lookback window falls under the threshold (default 0.9). Off unless set. */
|
||||
spend_capture_rate_check?: components["schemas"]["SpendCaptureRateCheckSettings"] | null;
|
||||
/**
|
||||
* Store Model In Db
|
||||
* @description If True, models and config are stored in and loaded from the database. Default is False.
|
||||
|
|
@ -42370,6 +42435,35 @@ export interface components {
|
|||
/** Model */
|
||||
model?: string | null;
|
||||
};
|
||||
/**
|
||||
* SpendCaptureRateCheckSettings
|
||||
* @description ``general_settings.spend_capture_rate_check``: the daily check of captured spend against the provider bill.
|
||||
*/
|
||||
SpendCaptureRateCheckSettings: {
|
||||
/**
|
||||
* Lookback Days
|
||||
* @default 7
|
||||
*/
|
||||
lookback_days: number;
|
||||
/**
|
||||
* Openai Project Ids
|
||||
* @description Scope the OpenAI bill to these project ids; empty compares against the whole organization. Captured spend is never scoped, so list every project LiteLLM's OpenAI keys belong to
|
||||
* @default []
|
||||
*/
|
||||
openai_project_ids: string[];
|
||||
/**
|
||||
* Providers
|
||||
* @default [
|
||||
* "openai"
|
||||
* ]
|
||||
*/
|
||||
providers: "openai"[];
|
||||
/**
|
||||
* Threshold
|
||||
* @default 0.9
|
||||
*/
|
||||
threshold: number;
|
||||
};
|
||||
/** SpendMetrics */
|
||||
SpendMetrics: {
|
||||
/**
|
||||
|
|
@ -65624,6 +65718,46 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
get_spend_capture_rate_spend_capture_rate_get: {
|
||||
parameters: {
|
||||
query: {
|
||||
/** @description First UTC day of the range, YYYY-MM-DD */
|
||||
start_date: string;
|
||||
/** @description Last UTC day of the range, YYYY-MM-DD, inclusive */
|
||||
end_date: string;
|
||||
/** @description Provider whose bill to compare against; needs OPENAI_ADMIN_KEY set on the proxy */
|
||||
provider?: "openai";
|
||||
/** @description Ratio under which the report flags below_threshold */
|
||||
threshold?: number;
|
||||
/** @description Scope the OpenAI bill to these project ids; omit to compare against the whole organization. Captured spend is never scoped, so pass every project LiteLLM's OpenAI keys belong to */
|
||||
project_ids?: string[] | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["CaptureRateReport"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
spend_key_fn_spend_keys_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue