scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

chore: merge litellm_internal_staging into litellm_techdebt_20260903

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-09-04 22:23:16 +00:00
commit 20bbdd28f8
136 changed files with 16332 additions and 1046 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
.circleci/config.yml
View file

@ -575,7 +575,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
@ -630,7 +630,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -737,7 +737,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -782,7 +782,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"assistants\""
@ -909,7 +909,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -999,7 +999,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1054,7 +1054,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 8 \
@ -1090,7 +1090,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1134,7 +1134,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1178,7 +1178,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1222,7 +1222,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1267,7 +1267,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1312,7 +1312,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1391,7 +1391,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1444,7 +1444,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
@ -1705,7 +1705,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -1794,7 +1794,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-s -v -x \
-s -v \
--junitxml=test-results/junit.xml \
-n 4 \
--durations=5"
@ -2012,7 +2012,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2092,7 +2092,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2195,7 +2195,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2266,7 +2266,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2350,7 +2350,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2446,7 +2446,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2516,7 +2516,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m

2
basedpyright-code-budget.json
View file

@ -93,7 +93,7 @@
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 24
"limit": 22
},
"reportUndefinedVariable": {
"limit": 0

1
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -161,6 +161,7 @@ class CheckBatchCost:
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION;
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "untracked_units" BIGINT NOT NULL DEFAULT 0;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1124,6 +1124,8 @@ model LiteLLM_DailyGuardrailUsageUnits {
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
cost Float? // USD for the priced share of units; null only on rows written before this column existed
untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out
created_at DateTime @default(now())
updated_at DateTime @updatedAt

16
litellm/integrations/cloudzero/database.py
View file

@ -94,8 +94,18 @@ class LiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *params)
# Convert the response to polars DataFrame with full schema inference
# This prevents schema mismatch errors when data types vary across rows
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as e:
raise Exception(f"Error retrieving usage data: {e}")

14
litellm/integrations/focus/database.py
View file

@ -96,7 +96,19 @@ class FocusLiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *query_params)
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc

8
litellm/litellm_core_utils/litellm_logging.py
View file

@ -10,7 +10,7 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime as dt_object
from functools import lru_cache
from types import MappingProxyType, TracebackType
@ -576,6 +576,7 @@ class Logging(LiteLLMLoggingBaseClass):
# enqueue closure here instead of firing it immediately.
self._defer_async_logging: bool = False
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
@ -1894,6 +1895,11 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
async def dispatch_failure_handlers(
self,
exception: Exception,

56
litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py
View file

@ -1,8 +1,8 @@
import math
from collections.abc import Mapping
from typing import Final
from typing import Annotated, Final
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
@ -30,6 +30,31 @@ class GuardrailCostEntry(BaseModel):
_GUARDRAIL_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailCostEntry]] = TypeAdapter(GuardrailCostEntry)
class GuardrailCostByUnitEntry(BaseModel):
"""The rollup-side view of a ``guardrail_information`` entry, validated apart from
``GuardrailCostEntry`` so a forged per-counter map can never zero the spend path."""
model_config = ConfigDict(extra="ignore", frozen=True)
guardrail_cost_by_unit: Mapping[str, Annotated[float, Field(ge=0, allow_inf_nan=False)] | None] | None = None
guardrail_cost_in_spend: bool | None = True
_GUARDRAIL_COST_BY_UNIT_ADAPTER: Final[TypeAdapter[GuardrailCostByUnitEntry]] = TypeAdapter(GuardrailCostByUnitEntry)
def billed_guardrail_cost_by_unit(raw: object) -> Mapping[str, float | None] | None:
"""Per-counter USD the daily rollup may record for one raw ``guardrail_information``
entry; None when the entry is unpriced, report-only, or malformed, and None per
counter the hook had no price for."""
try:
entry: Final = _GUARDRAIL_COST_BY_UNIT_ADAPTER.validate_python(raw)
except ValidationError as e:
verbose_logger.warning("Ignoring malformed guardrail_information entry for guardrail cost rollup: %s", e)
return None
return None if entry.guardrail_cost_in_spend is False else entry.guardrail_cost_by_unit
def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None:
regional_key: Final = f"bedrock/{aws_region_name}/guardrails" if aws_region_name else None
for key in (regional_key, BEDROCK_GUARDRAIL_PRICING_KEY):
@ -42,11 +67,32 @@ def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing
return None
def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float:
def _priced_units(units: int, price_per_unit: float | None) -> float | None:
return None if price_per_unit is None else units * price_per_unit
def bedrock_guardrail_cost_by_unit(
usage_units: Mapping[str, int], aws_region_name: str | None
) -> Mapping[str, float | None] | None:
"""USD per counter, keyed like ``usage_units``; None when no pricing entry exists,
and None for a counter the entry has no price for, since only an explicit 0.0 means free."""
pricing: Final = _bedrock_guardrail_pricing(aws_region_name)
if pricing is None:
return 0.0
return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items())
return None
return { # mutable-ok: stamped into guardrail_information, which safe_dumps only serializes as a plain dict
counter: _priced_units(units, pricing.guardrail_cost_per_unit.get(counter))
for counter, units in usage_units.items()
}
def guardrail_cost_total(cost_by_unit: Mapping[str, float | None] | None) -> float:
"""The scalar the spend path bills: unknown-priced counters count as 0 here, the
rollup keeps them unknown."""
return sum(cost for cost in cost_by_unit.values() if cost is not None) if cost_by_unit is not None else 0.0
def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float:
return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name))
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"

54
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Mapping, Sequence
from datetime import datetime
from typing import Any, Final, Protocol, runtime_checkable
@ -177,12 +177,6 @@ def _try_claim_detached_drain_slot() -> bool:
def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool:
"""After client detach the relay never reads the queue again, so drain it here.
The forwarded exception still sitting in the queue means the relay tore
down before re-raising it, so the proxy's failure handling never ran and
the caller must salvage spend itself.
"""
remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize()))
return any(item is exc for item in remaining)
@ -671,27 +665,41 @@ class BaseAnthropicMessagesStreamingIterator:
self,
queue: "asyncio.Queue[bytes | None | BaseException]",
client_detached: "asyncio.Event",
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks
exc: BaseException,
collected_chunks: Sequence[bytes],
exc: Exception,
) -> None:
"""Forward a provider error to a still-connected client, else salvage partial spend.
"""Log the request as failed with its partial usage, then make sure the proxy's failure hook runs once.
Handing the original exception to the client-facing generator lets it
re-raise so the proxy's failure handling keeps the provider status and
owns logging (no success-bill). If the client already went away, or
disconnects before ever consuming the queued exception, no failure hook
runs, so bill the partial instead of dropping the request.
A still-connected client gets the original exception through the queue,
the relay re-raises it, and the proxy's own failure handling records the
failed spend. When the client already left, or leaves before consuming
the queued exception, that handling never runs, so the detached-failure
hook the proxy armed on the logging object fires here instead.
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=self.litellm_logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
request_body=self.request_body,
raw_bytes=collected_chunks,
exception=exc,
)
if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc):
await client_detached.wait()
if not _exception_left_unconsumed(queue, exc):
return
verbose_proxy_logger.warning(
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
len(collected_chunks),
type(exc).__name__,
exc,
)
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
await self._fire_detached_failure_hook(exc)
async def _fire_detached_failure_hook(self, exc: Exception) -> None:
from litellm._logging import verbose_proxy_logger
on_detached_failure: Final = getattr(self.litellm_logging_obj, "_on_detached_stream_failure", None)
if on_detached_failure is None:
return
try:
await on_detached_failure(exc)
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not crash the detached pump
verbose_proxy_logger.warning(
"async_sse_wrapper detached failure hook raised: %s(%s)", type(hook_failure).__name__, hook_failure
)

2
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -29,7 +29,7 @@ from litellm.secret_managers.main import get_secret_str
BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE)
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws(?=/|$)", re.IGNORECASE)
def resolve_mantle_bearer_token(api_key: str | None) -> str | None:

57
litellm/llms/fireworks_ai/cost_calculator.py
View file

@ -2,6 +2,8 @@
For calculating cost of fireworks ai serverless inference models.
"""
import math
from datetime import datetime
from typing import Final
from litellm.constants import (
@ -10,9 +12,12 @@ from litellm.constants import (
FIREWORKS_AI_56_B_MOE,
FIREWORKS_AI_176_B_MOE,
)
from litellm.types.utils import Usage
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import get_model_info
NO_CACHE_READ_RATE: Final = float("nan")
# Extract the number of billion parameters from the model name
# only used for together_computer LLMs
@ -54,44 +59,50 @@ def get_base_model_for_pricing(model_name: str) -> str:
return "fireworks-ai-default"
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
def _resolve_model_info(model: str) -> ModelInfo:
try:
return get_model_info(model=model, custom_llm_provider="fireworks_ai")
except Exception:
base_model: Final = get_base_model_for_pricing(model_name=model)
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
swapping in the model's off_peak_pricing rates while one of its windows is open.
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- current_time: the moment the request is billed at; defaults to now, UTC
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
## check if model mapped, else use default pricing
try:
model_info = get_model_info(model=model, custom_llm_provider="fireworks_ai")
except Exception:
base_model: Final = get_base_model_for_pricing(model_name=model)
model_info: Final = _resolve_model_info(model)
standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
rates: Final = apply_off_peak_pricing(
model_info,
current_time,
TokenRates(
input_rate=model_info["input_cost_per_token"] or 0.0,
output_rate=model_info["output_cost_per_token"] or 0.0,
cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
cache_creation_rate=0.0,
reasoning_rate=None,
),
)
cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
## GET MODEL INFO
model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
## CALCULATE INPUT COST
prompt_tokens_details: Final = usage.prompt_tokens_details
cached_tokens: Final[int] = (
prompt_tokens_details.cached_tokens
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
else 0
)
input_cost_per_token: Final[float] = model_info["input_cost_per_token"] or 0.0
cache_read_input_token_cost: Final = model_info.get("cache_read_input_token_cost")
cache_read_cost_per_token: Final[float] = (
cache_read_input_token_cost if cache_read_input_token_cost is not None else input_cost_per_token
)
non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
prompt_cost: float = non_cached_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_cost_per_token
## CALCULATE OUTPUT COST
output_cost_per_token: Final[float] = model_info["output_cost_per_token"] or 0.0
completion_cost: Final[float] = usage.completion_tokens * output_cost_per_token
prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
return prompt_cost, completion_cost

23
litellm/llms/perplexity/cost_calculator.py
View file

@ -3,19 +3,23 @@ Helper util for handling perplexity-specific cost calculation
- e.g.: citation tokens, search queries
"""
from datetime import datetime
from typing import Final
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
from litellm.types.utils import Usage
from litellm.utils import get_model_info
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
The manual fallback swaps in the model's off_peak_pricing rates while one of its windows is open.
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing perplexity-specific usage information
- current_time: the moment the request is billed at; defaults to now, UTC
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@ -48,8 +52,21 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
except (ValueError, TypeError):
return default
rates: Final = apply_off_peak_pricing(
model_info,
current_time,
TokenRates(
input_rate=_safe_float_cast(model_info.get("input_cost_per_token")),
output_rate=_safe_float_cast(model_info.get("output_cost_per_token")),
cache_read_rate=0.0,
cache_creation_rate=0.0,
reasoning_rate=None,
),
)
input_cost_per_token: Final = rates.input_rate
output_cost_per_token: Final = rates.output_rate
## CALCULATE INPUT COST
input_cost_per_token: Final = _safe_float_cast(model_info.get("input_cost_per_token"))
prompt_cost: float = (usage.prompt_tokens or 0) * input_cost_per_token
## ADD CITATION TOKENS COST (if present)
@ -60,8 +77,6 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
prompt_cost += citation_tokens * citation_cost_per_token
## CALCULATE OUTPUT COST
output_cost_per_token: Final = _safe_float_cast(model_info.get("output_cost_per_token"))
reasoning_tokens = getattr(usage, "reasoning_tokens", 0) or 0
if reasoning_tokens == 0 and hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0

4914
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

3
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -928,6 +928,7 @@ def _resolve_openapi_tool_auth(
mcp_server_auth_headers,
alias=mcp_server.alias,
server_name=mcp_server.server_name,
access_groups=mcp_server.access_groups,
)
if mcp_server_auth_headers
else None
@ -3296,6 +3297,7 @@ class MCPServerManager:
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found
@ -5373,6 +5375,7 @@ class MCPServerManager:
mcp_server_auth_headers,
alias=mcp_server.alias,
server_name=mcp_server.server_name,
access_groups=mcp_server.access_groups,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found

7
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -257,7 +257,7 @@ if MCP_AVAILABLE:
)
def _get_server_auth_header(
server,
server: MCPServer,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
mcp_auth_header: str | None,
) -> dict[str, str] | str | None:
@ -269,8 +269,9 @@ if MCP_AVAILABLE:
if mcp_server_auth_headers:
server_auth: Final = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
if server_auth is not None:
return server_auth

6
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -1612,7 +1612,10 @@ if MCP_AVAILABLE:
)
server_headers: Final = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers, alias=server.alias, server_name=server.server_name
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
if isinstance(server_headers, str):
return bool(server_headers.strip())
@ -1712,6 +1715,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
extra_headers: dict[str, str] | None = None

49
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -8,11 +8,12 @@ import json
import os
import re
import typing
from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence
from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence, Sequence
from collections.abc import Set as AbstractSet
from typing import Any, Final, Protocol
from urllib.parse import quote
from litellm._logging import verbose_logger
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if typing.TYPE_CHECKING:
@ -169,34 +170,58 @@ def sanitize_mcp_alias_for_header(alias: str) -> str:
return sanitized.strip("_")
def _header_keys_for_identifier(identifier: str) -> tuple[str, ...]:
lowered: Final = identifier.lower()
sanitized: Final = sanitize_mcp_alias_for_header(identifier)
return (lowered,) if not sanitized or sanitized == lowered else (lowered, sanitized)
def _matching_header_key(normalized_headers: Mapping[str, object], identifier: str) -> str | None:
return next((key for key in _header_keys_for_identifier(identifier) if key in normalized_headers), None)
def lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers: Mapping[str, str | dict[str, str]],
*,
alias: str | None = None,
server_name: str | None = None,
access_groups: Sequence[str] | None = None,
) -> str | dict[str, str] | None:
"""
Resolve server-specific auth headers with case-insensitive matching.
Tries the raw alias/server_name (lowercased) and the header-safe sanitized
alias so dashboard clients using sanitize_mcp_alias_for_header() still match.
When no server-level header matches, an ``x-mcp-{access_group}-*`` header is
used as the default for every server in that group. If the server belongs to
several groups that each carry a different credential, nothing is returned so
a token is never forwarded to a server it may not have been meant for.
"""
if not mcp_server_auth_headers:
return None
normalized_headers: Final = {k.lower(): v for k, v in mcp_server_auth_headers.items()}
for identifier in (alias, server_name):
if not identifier:
continue
keys_to_try = [identifier.lower()]
sanitized = sanitize_mcp_alias_for_header(identifier)
if sanitized and sanitized not in keys_to_try:
keys_to_try.append(sanitized)
for key in keys_to_try:
if key in normalized_headers:
return normalized_headers[key]
return None
server_keys: Final = (
_matching_header_key(normalized_headers, identifier) for identifier in (alias, server_name) if identifier
)
server_key: Final = next((key for key in server_keys if key is not None), None)
if server_key is not None:
return normalized_headers[server_key]
group_keys: Final = (_matching_header_key(normalized_headers, group) for group in access_groups or ())
group_matches: Final = tuple(normalized_headers[key] for key in group_keys if key is not None)
if not group_matches:
return None
if any(match != group_matches[0] for match in group_matches[1:]):
verbose_logger.debug(
"Ambiguous MCP group auth headers for server alias=%s (groups=%s); not forwarding any group credential",
alias,
access_groups,
)
return None
return group_matches[0]
MCP_TOOL_ALLOWLIST_ENFORCED_KEY: Final = "tool_allowlist_enforced"

184
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -3151,6 +3151,17 @@
}
],
"title": "Team Id"
},
"user_email": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "User Email"
}
},
"title": "KeyMetadata",
@ -13039,6 +13050,59 @@
],
"title": "Avgscore"
},
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Cost"
},
"cost_by_key": {
"additionalProperties": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
]
},
"title": "Cost By Key",
"type": "object"
},
"cost_by_team": {
"additionalProperties": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
]
},
"title": "Cost By Team",
"type": "object"
},
"cost_by_unit": {
"additionalProperties": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
]
},
"title": "Cost By Unit",
"type": "object"
},
"description": {
"anyOf": [
{
@ -13089,6 +13153,13 @@
"title": "Type",
"type": "string"
},
"untracked_usage_units": {
"additionalProperties": {
"type": "integer"
},
"title": "Untracked Usage Units",
"type": "object"
},
"usage_units": {
"additionalProperties": {
"type": "integer"
@ -13140,7 +13211,12 @@
"usage_units",
"usage_units_daily",
"usage_units_by_team",
"usage_units_by_key"
"usage_units_by_key",
"cost",
"cost_by_unit",
"cost_by_team",
"cost_by_key",
"untracked_usage_units"
],
"title": "UsageDetailResponse",
"type": "object"
@ -13295,10 +13371,28 @@
"title": "Totalblocked",
"type": "integer"
},
"totalCost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Totalcost"
},
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUntrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totaluntrackedusageunits",
"type": "object"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
@ -13313,7 +13407,9 @@
"totalRequests",
"totalBlocked",
"passRate",
"totalUsageUnits"
"totalUsageUnits",
"totalCost",
"totalUntrackedUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -13342,6 +13438,18 @@
],
"title": "Avgscore"
},
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"description": "USD for the priced share of usageUnits over the window; null when no unit was priced",
"title": "Cost"
},
"failRate": {
"title": "Failrate",
"type": "number"
@ -13374,6 +13482,14 @@
"title": "Type",
"type": "string"
},
"untrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter",
"title": "Untrackedusageunits",
"type": "object"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
@ -13393,13 +13509,26 @@
"avgLatency",
"status",
"trend",
"usageUnits"
"usageUnits",
"cost",
"untrackedUsageUnits"
],
"title": "UsageOverviewRow",
"type": "object"
},
"UsageUnitsDailyPoint": {
"properties": {
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Cost"
},
"date": {
"title": "Date",
"type": "string"
@ -13414,7 +13543,8 @@
},
"required": [
"date",
"units"
"units",
"cost"
],
"title": "UsageUnitsDailyPoint",
"type": "object"
@ -28773,10 +28903,28 @@
"title": "Totalblocked",
"type": "integer"
},
"totalCost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Totalcost"
},
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUntrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totaluntrackedusageunits",
"type": "object"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
@ -28791,7 +28939,9 @@
"totalRequests",
"totalBlocked",
"passRate",
"totalUsageUnits"
"totalUsageUnits",
"totalCost",
"totalUntrackedUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -28820,6 +28970,18 @@
],
"title": "Avgscore"
},
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"description": "USD for the priced share of usageUnits over the window; null when no unit was priced",
"title": "Cost"
},
"failRate": {
"title": "Failrate",
"type": "number"
@ -28852,6 +29014,14 @@
"title": "Type",
"type": "string"
},
"untrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter",
"title": "Untrackedusageunits",
"type": "object"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
@ -28871,7 +29041,9 @@
"avgLatency",
"status",
"trend",
"usageUnits"
"usageUnits",
"cost",
"untrackedUsageUnits"
],
"title": "UsageOverviewRow",
"type": "object"

2
litellm/proxy/_types.py
View file

@ -670,6 +670,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_bulk_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/spend/by_user",
# gateway request counts (SGR); deployment-wide, admin-only
"/gateway/daily/activity",
# model
@ -832,6 +833,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/spend/by_user",
"/team/{team_id}/members/me",
"/model/new",
"/model/update",

19
litellm/proxy/auth/auth_checks.py
View file

@ -142,6 +142,8 @@ class _PrismaDictableRow(Protocol):
class _PrismaJWTKeyMappingRow(Protocol):
token: str
jwt_claim_name: str
jwt_claim_value: str
class _PrismaModelDumpRow(Protocol):
@ -3466,6 +3468,23 @@ async def _fetch_key_object_from_db_with_reconnect(
raise
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
"""Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping."""
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
@log_db_metrics
async def get_jwt_key_mapping_cache_keys_for_token(
hashed_token: str,
prisma_client: PrismaClient,
) -> tuple[str, ...]:
"""Cache keys of every JWT claim mapped to the given virtual key."""
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
where={"token": hashed_token}
)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings)
@log_db_metrics
async def get_jwt_key_mapping_object(
jwt_claim_name: str,

3
litellm/proxy/auth/user_api_key_auth.py
View file

@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import (
get_team_object,
get_user_object,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
resolve_and_validate_end_user_id,
)
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
@ -970,7 +971,7 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key: Final = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}"
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL:

33
litellm/proxy/common_request_processing.py
View file

@ -2520,6 +2520,11 @@ class ProxyBaseLLMRequestProcessing:
# This handles cases like websearch_interception agentic loop
# which returns a non-streaming dict even for streaming requests
if self._is_streaming_response(response):
self._arm_detached_stream_failure_hook(
logging_obj=logging_obj,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
selected_data_generator = ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
@ -2875,6 +2880,34 @@ class ProxyBaseLLMRequestProcessing:
),
)
def _arm_detached_stream_failure_hook(
self,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: "UserAPIKeyAuth",
proxy_logging_obj: ProxyLogging,
) -> None:
"""Let a stream that fails after the client left still reach ``post_call_failure_hook``.
The client-facing generator reports a mid-stream failure itself, but once
the client disconnects that generator is gone and the detached upstream
drain is the only code that sees the provider error. It fires this closure
so the failed spend is still written and the budget reservation released;
a replacement error the hook raises has no client left to reach.
"""
request_data: Final = self.data
async def _on_detached_stream_failure(exc: Exception) -> None:
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=exc,
request_data=request_data,
)
except HTTPException:
return
logging_obj._on_detached_stream_failure = _on_detached_stream_failure
def _is_streaming_response(self, response: Any) -> bool:
"""
Check if the response object is actually a streaming response by inspecting its type.

48
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -36,7 +36,10 @@ from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_rege
from litellm.litellm_core_utils.litellm_logging import (
_get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
bedrock_guardrail_cost_by_unit,
guardrail_cost_total,
)
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
@ -109,6 +112,7 @@ _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = (
_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3
_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5
_BEDROCK_WHITESPACE: Final = re.compile(r"\s")
_NO_TRACING_DETAIL: Final[GuardrailTracingDetail] = {}
# Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke"
# InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with
@ -2155,25 +2159,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
OTEL integration can expose it as a queryable span attribute
without re-parsing the redacted guardrail_response blob.
"""
tracing_detail: Final[GuardrailTracingDetail] = {}
violation_categories: Final = self._extract_violation_category_names(response)
if violation_categories:
tracing_detail["violation_categories"] = violation_categories
bedrock_action: Final = response.get("action")
if isinstance(bedrock_action, str):
tracing_detail["guardrail_action"] = bedrock_action
usage: Final = response.get("usage")
if isinstance(usage, dict):
usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream
key: value for key, value in usage.items() if isinstance(value, int)
}
if usage_units:
tracing_detail["guardrail_usage"] = usage_units
tracing_detail["guardrail_cost"] = bedrock_guardrail_cost(
usage_units=usage_units, aws_region_name=aws_region_name
)
categories_detail: Final[GuardrailTracingDetail] = {"violation_categories": violation_categories}
action_detail: Final[GuardrailTracingDetail] = {"guardrail_action": bedrock_action}
tracing_detail: Final[GuardrailTracingDetail] = {
**(categories_detail if violation_categories else _NO_TRACING_DETAIL),
**(action_detail if isinstance(bedrock_action, str) else _NO_TRACING_DETAIL),
**self._usage_tracing_detail(response.get("usage"), aws_region_name),
}
return tracing_detail
@staticmethod
def _usage_tracing_detail(
usage: BedrockGuardrailUsage | None, aws_region_name: str | None
) -> GuardrailTracingDetail:
if not isinstance(usage, dict):
return _NO_TRACING_DETAIL
usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream
key: value for key, value in usage.items() if isinstance(value, int)
}
if not usage_units:
return _NO_TRACING_DETAIL
cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage_units=usage_units, aws_region_name=aws_region_name)
priced_detail: Final[GuardrailTracingDetail] = {"guardrail_cost_by_unit": cost_by_unit}
usage_detail: Final[GuardrailTracingDetail] = {
"guardrail_usage": usage_units,
"guardrail_cost": guardrail_cost_total(cost_by_unit),
**(priced_detail if cost_by_unit is not None else _NO_TRACING_DETAIL),
}
return usage_detail
def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]:
"""
Flatten the BLOCKED assessments into a list of human-readable category

112
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -8,10 +8,10 @@ from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import date, datetime, timedelta, timezone
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, overload
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel
from pydantic import BaseModel, Field
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
@ -42,6 +42,8 @@ router: Final = APIRouter()
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
_T = TypeVar("_T")
_USAGE_MAX_RANGE_DAYS: Final = 366
@ -154,6 +156,16 @@ def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str:
return row.usage_unit
def _row_untracked_units(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> int:
"""A row written before the cost column carries NULL cost and is untracked in full."""
return int(row.units) if row.cost is None else int(row.untracked_units)
def _row_tracked_cost(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> float | None:
"""The row's cost when it prices at least one unit; None when every unit is untracked."""
return None if row.cost is None or _row_untracked_units(row) >= int(row.units) else row.cost
def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]:
ordered: Final = sorted(rows, key=_counter_name)
return MappingProxyType(
@ -161,12 +173,31 @@ def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsage
)
def _units_by(
def _sum_untracked_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]:
ordered: Final = sorted(rows, key=_counter_name)
per_counter: Final = tuple(
(name, sum(map(_row_untracked_units, group))) for name, group in groupby(ordered, key=_counter_name)
)
return MappingProxyType({name: units for name, units in per_counter if units})
def _sum_tracked_cost(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> float | None:
"""Sum over rows that price at least one unit; None when no row does."""
tracked: Final = tuple(cost for cost in map(_row_tracked_cost, rows) if cost is not None)
return sum(tracked) if tracked else None
def _by(
rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]",
key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]",
) -> Mapping[str, Mapping[str, int]]:
reduce: "Callable[[Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]], _T]",
) -> Mapping[str, _T]:
ordered: Final = sorted(rows, key=key_of)
return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)})
return MappingProxyType({key: reduce(group) for key, group in groupby(ordered, key=key_of)})
def _first_match(lookup_keys: Sequence[str], mapping: Mapping[str, _T], default: _T) -> _T:
return next((mapping[k] for k in lookup_keys if k in mapping), default)
# --- Response models ---
@ -218,6 +249,12 @@ class UsageOverviewRow(BaseModel):
status: str # healthy | warning | critical
trend: str # up | down | stable
usageUnits: Mapping[str, int]
cost: float | None = Field(
description="USD for the priced share of usageUnits over the window; null when no unit was priced"
)
untrackedUsageUnits: Mapping[str, int] = Field(
description="The share of usageUnits that cost leaves out: units recorded with no known price, per counter"
)
class UsageOverviewResponse(BaseModel):
@ -227,11 +264,26 @@ class UsageOverviewResponse(BaseModel):
totalBlocked: int
passRate: float
totalUsageUnits: Mapping[str, int]
totalCost: float | None
totalUntrackedUsageUnits: Mapping[str, int]
_EMPTY_OVERVIEW: Final = UsageOverviewResponse(
rows=[],
chart=[],
totalRequests=0,
totalBlocked=0,
passRate=100.0,
totalUsageUnits=_EMPTY_UNITS,
totalCost=None,
totalUntrackedUsageUnits=_EMPTY_UNITS,
)
class UsageUnitsDailyPoint(BaseModel):
date: str
units: Mapping[str, int]
cost: float | None
class UsageDetailResponse(BaseModel):
@ -251,6 +303,11 @@ class UsageDetailResponse(BaseModel):
usage_units_daily: Sequence[UsageUnitsDailyPoint]
usage_units_by_team: Mapping[str, Mapping[str, int]]
usage_units_by_key: Mapping[str, Mapping[str, int]]
cost: float | None
cost_by_unit: Mapping[str, float | None]
cost_by_team: Mapping[str, float | None]
cost_by_key: Mapping[str, float | None]
untracked_usage_units: Mapping[str, int]
class UsageLogEntry(BaseModel):
@ -367,6 +424,8 @@ def _guardrail_overview_rows(
agg: Mapping[str, _MetricTotals],
prev_agg: Mapping[str, float],
units_agg: Mapping[str, Mapping[str, int]],
cost_agg: Mapping[str, float | None],
untracked_agg: Mapping[str, Mapping[str, int]],
) -> list[UsageOverviewRow]:
rows: Final[list[UsageOverviewRow]] = []
covered_keys: Final[set[str]] = set()
@ -392,7 +451,6 @@ def _guardrail_overview_rows(
prev_fail = float(prev_agg.get(k, 0.0) or 0.0)
break
trend = _trend_from_comparison(fail_rate, prev_fail)
row_units: Mapping[str, int] = next((units_agg[k] for k in lookup_keys if k in units_agg), _EMPTY_UNITS)
rows.append(
UsageOverviewRow(
id=gid,
@ -405,7 +463,9 @@ def _guardrail_overview_rows(
avgLatency=None,
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=row_units,
usageUnits=_first_match(lookup_keys, units_agg, _EMPTY_UNITS),
cost=_first_match(lookup_keys, cost_agg, None),
untrackedUsageUnits=_first_match(lookup_keys, untracked_agg, _EMPTY_UNITS),
)
)
# Add rows for guardrails with metrics but not in guardrails table (e.g. MCP, config)
@ -429,6 +489,8 @@ def _guardrail_overview_rows(
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=units_agg.get(agg_key, _EMPTY_UNITS),
cost=cost_agg.get(agg_key),
untrackedUsageUnits=untracked_agg.get(agg_key, _EMPTY_UNITS),
)
)
return rows
@ -459,6 +521,8 @@ def _policy_overview_rows(
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=_EMPTY_UNITS,
cost=None,
untrackedUsageUnits=_EMPTY_UNITS,
)
)
return rows
@ -479,9 +543,7 @@ async def guardrails_usage_overview(
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return UsageOverviewResponse(
rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS
)
return _EMPTY_OVERVIEW
start, end = _resolve_usage_window(start_date, end_date)
@ -515,12 +577,14 @@ async def guardrails_usage_overview(
agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id")
prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id")
units_agg: Final = _units_by(units_rows, lambda r: r.guardrail_id)
units_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_counter_units)
cost_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_tracked_cost)
untracked_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_untracked_units)
chart: Final = _chart_from_metrics(metrics)
total_requests: Final = sum(a["requests"] for a in agg.values())
total_blocked: Final = sum(a["blocked"] for a in agg.values())
pass_rate: Final = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0
rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg)
rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg, cost_agg, untracked_agg)
return UsageOverviewResponse(
rows=rows,
chart=chart,
@ -528,6 +592,8 @@ async def guardrails_usage_overview(
totalBlocked=total_blocked,
passRate=round(pass_rate, 1),
totalUsageUnits=_sum_counter_units(units_rows),
totalCost=_sum_tracked_cost(units_rows),
totalUntrackedUsageUnits=_sum_untracked_units(units_rows),
)
except Exception as e:
from litellm.proxy.utils import handle_exception_on_proxy
@ -618,8 +684,11 @@ async def guardrails_usage_detail(
litellm_params: Final = _to_dict(_get_guardrail_field(guardrail, "litellm_params"))
guardrail_info: Final = _to_dict(_get_guardrail_field(guardrail, "guardrail_info"))
_guardrail_name: Final = _get_guardrail_field(guardrail, "guardrail_name")
daily_unit_sums: Final = sorted(_units_by(units_rows, lambda r: r.date).items())
units_daily: Final = tuple(UsageUnitsDailyPoint(date=d, units=units) for d, units in daily_unit_sums)
daily_unit_sums: Final = sorted(_by(units_rows, lambda r: r.date, _sum_counter_units).items())
daily_cost: Final = _by(units_rows, lambda r: r.date, _sum_tracked_cost)
units_daily: Final = tuple(
UsageUnitsDailyPoint(date=d, units=units, cost=daily_cost.get(d)) for d, units in daily_unit_sums
)
return UsageDetailResponse(
guardrail_id=guardrail_id,
@ -636,8 +705,13 @@ async def guardrails_usage_detail(
time_series=time_series,
usage_units=_sum_counter_units(units_rows),
usage_units_daily=units_daily,
usage_units_by_team=_units_by(units_rows, lambda r: r.team_id),
usage_units_by_key=_units_by(units_rows, lambda r: r.api_key),
usage_units_by_team=_by(units_rows, lambda r: r.team_id, _sum_counter_units),
usage_units_by_key=_by(units_rows, lambda r: r.api_key, _sum_counter_units),
cost=_sum_tracked_cost(units_rows),
cost_by_unit=_by(units_rows, _counter_name, _sum_tracked_cost),
cost_by_team=_by(units_rows, lambda r: r.team_id, _sum_tracked_cost),
cost_by_key=_by(units_rows, lambda r: r.api_key, _sum_tracked_cost),
untracked_usage_units=_sum_untracked_units(units_rows),
)
@ -857,9 +931,7 @@ async def policies_usage_overview(
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return UsageOverviewResponse(
rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS
)
return _EMPTY_OVERVIEW
start, end = _resolve_usage_window(start_date, end_date)
@ -891,6 +963,8 @@ async def policies_usage_overview(
totalBlocked=total_blocked,
passRate=round(pass_rate, 1),
totalUsageUnits=_EMPTY_UNITS,
totalCost=None,
totalUntrackedUsageUnits=_EMPTY_UNITS,
)
except Exception as e:
from litellm.proxy.utils import handle_exception_on_proxy

73
litellm/proxy/guardrails/usage_tracking.py
View file

@ -6,7 +6,7 @@ insert into SpendLogGuardrailIndex when spend logs are written.
import asyncio
import json
from collections import defaultdict
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from functools import partial
from itertools import groupby
@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guardrail_cost_by_unit
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import (
@ -44,6 +45,20 @@ class _UsageUnitKey(NamedTuple):
usage_unit: str
class _UsageUnitIncrement(NamedTuple):
units: int
cost: float
"""USD for the priced share of units."""
untracked_units: int
"""Units recorded with no known price, the share cost leaves out."""
def _usage_unit_increment(units: int, cost: float | None) -> _UsageUnitIncrement:
if cost is None:
return _UsageUnitIncrement(units=units, cost=0.0, untracked_units=units)
return _UsageUnitIncrement(units=units, cost=cost, untracked_units=0)
class _MetricsKey(NamedTuple):
guardrail_id: str
date: str
@ -67,22 +82,37 @@ class PendingRollups:
def __init__(self) -> None:
self.lock: Final = asyncio.Lock()
self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({})
self.units: Mapping[_UsageUnitKey, int] = MappingProxyType({})
self.units: Mapping[_UsageUnitKey, _UsageUnitIncrement] = MappingProxyType({})
_PENDING_ROLLUPS: Final = PendingRollups()
_NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({})
_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0, untracked_units=0)
def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]:
return (*base, *(key for key in extra if key not in base))
def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitIncrement:
materialized: Final = tuple(increments)
return _UsageUnitIncrement(
units=sum(i.units for i in materialized),
cost=sum(i.cost for i in materialized),
untracked_units=sum(i.untracked_units for i in materialized),
)
def _merged_unit_rows(
base: Mapping[_UsageUnitKey, int], extra: Mapping[_UsageUnitKey, int]
) -> Mapping[_UsageUnitKey, int]:
return MappingProxyType({key: base.get(key, 0) + extra.get(key, 0) for key in _merged_keys(base, extra)})
base: Mapping[_UsageUnitKey, _UsageUnitIncrement], extra: Mapping[_UsageUnitKey, _UsageUnitIncrement]
) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]:
return MappingProxyType(
{
key: _summed_increments((base.get(key, _NO_INCREMENT), extra.get(key, _NO_INCREMENT)))
for key in _merged_keys(base, extra)
}
)
def _merged_metric_rows(
@ -209,7 +239,9 @@ def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None:
return None
def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]:
def _iter_usage_unit_increments(
logs_to_process: Sequence[Mapping[str, Any]],
) -> Iterator[tuple[_UsageUnitKey, _UsageUnitIncrement]]:
for payload in logs_to_process:
start_time = _parse_payload_start_time(payload)
if not payload.get("request_id") or start_time is None:
@ -222,26 +254,38 @@ def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) ->
usage = entry.get("guardrail_usage")
if not guardrail_id or not isinstance(usage, dict):
continue
cost_by_unit = billed_guardrail_cost_by_unit(entry)
for unit_name, units in usage.items():
if isinstance(units, int) and not isinstance(units, bool) and units > 0:
yield _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)), units
key = _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name))
cost = cost_by_unit.get(str(unit_name)) if cost_by_unit is not None else None
yield key, _usage_unit_increment(units=units, cost=cost)
def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]:
def _sum_usage_unit_increments(
logs_to_process: Sequence[Mapping[str, Any]],
) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]:
ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=itemgetter(0))
return MappingProxyType(
{key: sum(units for _, units in group) for key, group in groupby(ordered, key=itemgetter(0))}
{
key: _summed_increments(increment for _, increment in group)
for key, group in groupby(ordered, key=itemgetter(0))
}
)
async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey, units: int) -> None:
async def _upsert_usage_unit_row(
prisma_client: PrismaClient, key: _UsageUnitKey, increment: _UsageUnitIncrement
) -> None:
row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = {
"guardrail_id": key.guardrail_id,
"date": key.date,
"team_id": key.team_id,
"api_key": key.api_key,
"usage_unit": key.usage_unit,
"units": units,
"units": increment.units,
"cost": increment.cost,
"untracked_units": increment.untracked_units,
}
where: Final[_UsageUnitWhereUnique] = {
"guardrail_id_date_team_id_api_key_usage_unit": {
@ -252,9 +296,14 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey
"usage_unit": key.usage_unit,
}
}
# A row written before the cost column has NULL cost, and NULL + x stays NULL, so it keeps reading as unknown
data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = {
"create": row,
"update": {"units": {"increment": units}},
"update": {
"units": {"increment": increment.units},
"cost": {"increment": increment.cost},
"untracked_units": {"increment": increment.untracked_units},
},
}
await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data)

46
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -115,6 +115,29 @@ _CONFIG_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset(
)
def _request_inherits_config_credentials(
config_params: Mapping[str, object],
request_params: Mapping[str, object],
allow_client_side_credentials: bool,
) -> bool:
"""Whether the configuration's credentials are this request's to be probed with.
The configuration reached here by matching the request's model string, which
also matches wildcard routes and unrelated deployments that merely serve the
same model, so a request naming a stored credential of its own has already
said where its credentials come from and does not borrow that one's. A blank
name is no name: ``load_credentials_from_list`` resolves nothing from it, so
it must not cost the request the credentials it would otherwise be probed
with.
"""
requested_credential: Final = request_params.get("litellm_credential_name")
if requested_credential and requested_credential != config_params.get("litellm_credential_name"):
return False
if allow_client_side_credentials:
return True
return not any(param in request_params for param in _BANNED_REQUEST_BODY_PARAMS)
def _config_base_for_health_check(
config_params: Mapping[str, object],
request_params: Mapping[str, object],
@ -122,25 +145,19 @@ def _config_base_for_health_check(
) -> dict[str, object]:
"""Return the configured parameters to merge under a connection-test request.
A request that sets its own connection fields describes a connection of its
own, so the configuration's credentials are not carried into it: they belong
to the endpoint the configuration names. Anything the request does not set
still comes from the configuration, which is what lets a request name a
configured model and test it as configured.
A request that sets its own connection fields, or names its own stored
credential, describes a connection of its own, so the configuration's
credentials are not carried into it: they belong to the endpoint the
configuration names. Anything the request does not set still comes from the
configuration, which is what lets a request name a configured model and test
it as configured.
``litellm_credential_name`` is dropped alongside the literal credential
fields: it names a stored credential that ``load_credentials_from_list``
resolves into the same secrets further down the call, so leaving it in place
would reintroduce them by reference.
``general_settings.allow_client_side_credentials`` is the existing proxy-wide
opt-in for callers supplying their own connection parameters. Where an admin
has enabled it, a request may pair its own endpoint with the configured
credentials, as it could before.
"""
if allow_client_side_credentials:
return dict(config_params)
if not any(param in request_params for param in _BANNED_REQUEST_BODY_PARAMS):
if _request_inherits_config_credentials(config_params, request_params, allow_client_side_credentials):
return dict(config_params)
return {key: value for key, value in config_params.items() if key not in _CONFIG_CONNECTION_FIELDS}
@ -1959,6 +1976,9 @@ async def test_model_connection(
Note:
- If the model is configured in proxy_config.yaml, credentials (api_key, api_base, etc.)
will be automatically loaded from the config (with resolved environment variables).
- A request naming a stored credential (`litellm_credential_name`) that the configuration
does not name is probed with that credential instead, and inherits no credentials
from the configuration its model string happened to match.
- You can override specific params by including them in the request.
- You can use `os.environ/VARIABLE_NAME` syntax to reference environment variables,
which will be resolved automatically (same as in proxy_config.yaml).

58
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -4518,12 +4518,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
statuses=statuses,
)
def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]:
usage: Final = source.get("combined_usage_object")
if not isinstance(usage, Usage) or (usage.completion_tokens or 0) <= 0:
return 0, 0, 0
billable_input, completion_tokens, _ = self._resolve_io_token_reconcile_usage(usage)
return (
self._get_total_tokens_from_usage(usage=usage, rate_limit_type=self.get_rate_limit_type()),
billable_input,
completion_tokens,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
On failure: decrement max_parallel_requests and refund the upfront
TPM reservation only against the scopes the reservation actually
charged. Unreserved scopes were never incremented at pre-call, so
refunding them would drive their counter negative.
refunding them would drive their counter negative. A failed stream
whose partial usage was recovered settles the reservation at that
usage instead of refunding it.
"""
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
@ -4552,31 +4565,31 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if stash is None or stash.reservation_released
else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(kwargs)
if stash is not None and reserved_tokens > 0:
verbose_proxy_logger.debug("Releasing reserved TPM tokens on failure: %s", reserved_tokens)
# Refund only against the scopes the reservation actually
# charged. _build_reservation_aware_tpm_ops with
# actual_tokens=0 emits -reserved on reserved scopes and 0
# on unreserved (skipped), so unreserved scopes can't drift
# negative.
verbose_proxy_logger.debug(
"Settling reserved TPM tokens on failure: reserved=%s actual=%s", reserved_tokens, tpm_actual
)
# Settle only against the scopes the reservation actually
# charged: unreserved scopes were never incremented, so a
# refund there would drive their counter negative.
pipeline_operations.extend(
self._build_reservation_aware_tpm_ops(
targets=list(stash.reserved_scopes),
reserved_scopes=stash.reserved_scopes,
actual_tokens=0,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
)
)
# Refund project ITPM/OTPM reservations the same way -- full
# refund, since a failed call has no billable usage to reconcile
# against.
# Settle project ITPM/OTPM reservations the same way: at the
# recovered partial usage, or a full refund when there is none.
itpm_operations: Final = (
self._build_project_reservation_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
reservation_window_identities=stash.itpm_reserved_window_identities,
)
@ -4584,7 +4597,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
)
if stash is not None and itpm_reserved > 0
@ -4595,7 +4608,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
reservation_window_identities=stash.otpm_reserved_window_identities,
)
@ -4603,7 +4616,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
)
if stash is not None and otpm_reserved > 0
@ -4742,7 +4755,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
removal is a no-op ZREM on a second run), and the TPM/ITPM/OTPM
refund is guarded by the stash's ``reservation_released`` flag — if
both this hook and async_log_failure_event end up running in the same
flow, only the first release/refund applies.
flow, only the first release/refund applies. A mid-stream failure
relayed here with recovered partial usage settles the reservation at
that usage instead of refunding it.
"""
try:
stash: Final = get_request_stash()
@ -4769,12 +4784,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
otpm_reserved: Final = stash.otpm_reserved_tokens
if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0:
return
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data)
combined_ops: Final = (
self._build_reservation_aware_tpm_ops(
targets=tuple(stash.reserved_scopes),
reserved_scopes=stash.reserved_scopes,
actual_tokens=0,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
)
if reserved_tokens > 0
@ -4784,7 +4800,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
reservation_window_identities=stash.itpm_reserved_window_identities,
)
@ -4792,7 +4808,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
)
if itpm_reserved > 0
@ -4802,7 +4818,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
reservation_window_identities=stash.otpm_reserved_window_identities,
)
@ -4810,7 +4826,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
)
if otpm_reserved > 0

4
litellm/proxy/logging_endpoints/callback_logs_endpoints.py
View file

@ -88,8 +88,10 @@ class CallbackLogsReplayer:
)
metadata: Final[dict[str, Any]] = payload.get("metadata") or {}
user_api_key_hash: Final = metadata.get("user_api_key_hash")
litellm_metadata: Final[dict[str, Any]] = {
"user_api_key": metadata.get("user_api_key_hash"),
"user_api_key": user_api_key_hash,
"user_api_key_hash": user_api_key_hash,
"user_api_key_alias": metadata.get("user_api_key_alias"),
"user_api_key_user_id": metadata.get("user_api_key_user_id"),
"user_api_key_team_id": metadata.get("user_api_key_team_id"),

84
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -2,15 +2,19 @@ import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import HTTPException, status
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
recover_double_hashed_key_metadata,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
@ -111,8 +115,19 @@ class DailySpendRecord(Protocol):
class _KeyMetadataDict(TypedDict, total=False):
key_alias: str | None
team_id: str | None
key_alias: ReadOnly[str | None]
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
meta: Final = api_key_metadata.get(api_key, {})
return KeyMetadata(
key_alias=meta.get("key_alias"),
team_id=meta.get("team_id"),
user_email=meta.get("user_email"),
)
_WhereValue = str | dict[str, object]
@ -283,10 +298,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.models[model_key].api_key_breakdown:
breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics,
@ -310,10 +322,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown:
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics,
@ -335,10 +344,7 @@ def update_breakdown_metrics(
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
)
@ -363,10 +369,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
@ -388,10 +391,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
@ -403,10 +403,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.api_keys:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
), # Add any api_key-specific metadata here
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
@ -426,10 +423,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
@ -442,17 +436,23 @@ def update_breakdown_metrics(
async def get_api_key_metadata(
prisma_client: PrismaClient,
api_keys: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
) -> Mapping[str, _KeyMetadataDict]:
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
This ensures that key_alias and team_id are preserved in historical activity logs
even after a key is deleted or regenerated.
even after a key is deleted or regenerated. Also recovers aliases for api_key
values that were double-hashed by the v1.99 spend-log provenance gate.
"""
key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
where={"token": {"in": list(api_keys)}}
)
result: Final[dict[str, _KeyMetadataDict]] = {
k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records
k.token: {
"key_alias": k.key_alias,
"team_id": k.team_id,
"user_id": getattr(k, "user_id", None),
}
for k in key_records
}
# For any keys not found in the active table, check the deleted keys table
@ -471,6 +471,7 @@ async def get_api_key_metadata(
result[k.token] = {
"key_alias": k.key_alias,
"team_id": k.team_id,
"user_id": getattr(k, "user_id", None),
}
except Exception as e:
verbose_proxy_logger.warning(
@ -479,7 +480,13 @@ async def get_api_key_metadata(
e,
)
return result
still_missing: Final = api_keys - frozenset(result)
combined: Final = (
result
if not still_missing
else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))})
)
return await attach_user_emails(prisma_client, combined)
def _adjust_dates_for_timezone(
@ -951,11 +958,6 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
)
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
meta: Final = api_key_metadata.get(api_key, {})
return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id"))
def _aggregate_grouping_sets_records_sync(
*,
records: Sequence[_GroupingSetsRow],

4
litellm/proxy/management_endpoints/common_utils.py
View file

@ -22,7 +22,7 @@ def validate_finite_spend(spend: float | None) -> None:
)
def validate_budget_duration(budget_duration: str | None) -> None:
def validate_budget_duration(budget_duration: str | None, status_code: int = 400) -> None:
"""Reject budget durations that can't be parsed, are non-positive, or
overflow date math, so a bad value can't be persisted and later crash the
budget reset job.
@ -44,7 +44,7 @@ def validate_budget_duration(budget_duration: str | None) -> None:
get_budget_reset_time(budget_duration=budget_duration)
except (ValueError, OverflowError):
raise HTTPException(
status_code=400,
status_code=status_code,
detail={
"error": f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'."
},

30
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
View file

@ -13,7 +13,9 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
hash_token,
)
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.repositories.table_repositories import JWTKeyMappingRepository
@ -118,9 +120,8 @@ async def create_jwt_key_mapping(
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
# Invalidate cache
cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(new_mapping)
except HTTPException:
@ -169,17 +170,20 @@ async def update_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data)
if updated_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
# Invalidate new cache key if claim fields changed
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
# Evict only after the write commits: a concurrent request between an
# early eviction and the commit would re-cache the old mapping and keep
# it authorized until TTL.
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
new_cache_key: Final = jwt_key_mapping_cache_key(
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
)
cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key)
await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
return _to_response(updated_mapping)
except HTTPException:
@ -219,10 +223,12 @@ async def delete_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
await _mapping_table(prisma_client).delete(where={"id": data.id})
# Evict only after the row is gone, else a concurrent request can
# re-cache the deleted mapping and keep it authorized until TTL.
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return {"status": "success"}
except HTTPException:
raise

11
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -54,6 +54,7 @@ from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken,
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
can_team_access_model,
get_jwt_key_mapping_cache_keys_for_token,
get_org_object,
get_project_object,
get_team_object,
@ -65,6 +66,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
evict_and_broadcast,
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
@ -4975,6 +4977,13 @@ async def _execute_virtual_key_regeneration(
update_data.update(non_default_values)
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
# Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash,
# but their cached jwt_key_mapping entries still point at the old token (LIT-5379).
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_token(
hashed_token=hashed_api_key,
prisma_client=prisma_client,
)
# If grace period set, insert deprecated key so old key remains valid
await _insert_deprecated_key(
prisma_client=prisma_client,
@ -5000,6 +5009,8 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
# After credential invalidation, so a failure here can never keep the old key alive.
await sync_key_regeneration_access_group_membership(
prisma_client=prisma_client,

18
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -41,6 +41,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import get_daily_a
from litellm.proxy.management_endpoints.common_utils import (
_set_object_metadata_field,
_user_has_admin_view,
validate_budget_duration,
)
from litellm.proxy.management_helpers.object_permission_utils import (
handle_update_object_permission_common,
@ -312,7 +313,7 @@ def handle_nested_budget_structure_in_organization_update_request(
# Extract valid budget fields and merge into top level
budget_fields: Final = LiteLLM_BudgetTable.model_fields.keys()
for key, value in budget_data.items():
if key in budget_fields and value is not None:
if key in budget_fields:
transformed_data[key] = value
return transformed_data
@ -708,9 +709,8 @@ async def update_organization(
existing_organization_row=existing_organization_row,
)
# Handle budget updates if budget fields are provided
budget_fields: Final = {
k: v for k, v in data.model_dump().items() if k in LiteLLM_BudgetTable.model_fields and v is not None
k: v for k, v in data.model_dump().items() if k in _BUDGET_SETTABLE_FIELDS and k in data.model_fields_set
}
if budget_fields and existing_organization_row.budget_id:
@ -764,7 +764,6 @@ async def handle_update_object_permission(
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
response_model=LiteLLM_OrganizationTableWithMembers,
include_in_schema=False,
)
async def update_organization_v2(
organization_id: str,
@ -807,6 +806,17 @@ async def update_organization_v2(
status_code=422,
detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
)
for limit_name, limit_value in (
("tpm_limit", data.tpm_limit),
("rpm_limit", data.rpm_limit),
("max_parallel_requests", data.max_parallel_requests),
):
if limit_value is not None and limit_value < 0:
raise HTTPException(
status_code=422,
detail={"error": f"{limit_name} must be non-negative. Received: {limit_value}"},
)
validate_budget_duration(data.budget_duration, status_code=422)
if data.model_max_budget:
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_model_max_budget,

123
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -170,6 +170,8 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
TeamMemberAddResult,
TeamMemberInfoResponse,
TeamMetadataSchemaResponse,
TeamUserSpendResponse,
TeamUserSpendRow,
UpdateTeamMemberPermissionsRequest,
)
@ -6231,3 +6233,124 @@ async def get_team_daily_activity_aggregated(
timezone_offset_minutes=timezone,
include_entity_breakdown=True,
)
def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str:
team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count))
user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else ""
return f"""
SELECT
sl.team_id,
sl."user" AS user_id,
u.user_email,
u.user_alias,
SUM(sl.spend)::float AS spend,
SUM(sl.prompt_tokens)::bigint AS prompt_tokens,
SUM(sl.completion_tokens)::bigint AS completion_tokens,
SUM(sl.total_tokens)::bigint AS total_tokens,
COUNT(*)::bigint AS api_requests,
COUNT(*) FILTER (WHERE sl.status IS DISTINCT FROM 'failure')::bigint AS successful_requests,
COUNT(*) FILTER (WHERE sl.status = 'failure')::bigint AS failed_requests
FROM "LiteLLM_SpendLogs" sl
LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = sl."user"
WHERE sl."startTime" >= $1::timestamp
AND sl."startTime" < $2::timestamp + INTERVAL '1 day'
AND sl.team_id IN ({team_placeholders}){user_clause}
GROUP BY sl.team_id, sl."user", u.user_email, u.user_alias
ORDER BY spend DESC, sl.team_id, sl."user"
"""
class _TeamUserSpendDbRow(TypedDict):
team_id: ReadOnly[str]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
user_alias: ReadOnly[str | None]
spend: ReadOnly[float]
prompt_tokens: ReadOnly[int]
completion_tokens: ReadOnly[int]
total_tokens: ReadOnly[int]
api_requests: ReadOnly[int]
successful_requests: ReadOnly[int]
failed_requests: ReadOnly[int]
@router.get(
"/team/spend/by_user",
response_model=TeamUserSpendResponse,
tags=["team management"], # mutable-ok: fastapi route tags must be a list
)
async def get_team_spend_by_user(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
team_ids: str | None = None,
start_date: str | None = None,
end_date: str | None = None,
) -> TeamUserSpendResponse:
"""
Spend per user within the given teams, attributed per request from spend logs.
Proxy admins may query any team. Team admins and members holding the
`/team/daily/activity` permission see every user of the requested teams;
other members only see their own row.
"""
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value)
range_error: Final = _aggregated_date_range_error(start_date, end_date)
if range_error is not None or start_date is None or end_date is None:
raise _daily_activity_error(status_code=400, message=range_error or "Please provide start_date and end_date")
if not team_ids:
raise _daily_activity_error(status_code=400, message="Please provide team_ids")
scope: Final = await _resolve_team_daily_activity_scope(
team_ids=team_ids,
exclude_team_ids=None,
api_key=None,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
scoped_team_ids: Final = tuple(scope.team_ids or ())
if not scoped_team_ids:
return TeamUserSpendResponse(start_date=start_date, end_date=end_date, results=())
own_user_only: Final = scope.api_key_filter is not None
user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw(
_team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
start_date,
end_date,
*scoped_team_ids,
*user_param,
)
results: Final = tuple(
TeamUserSpendRow(
team_id=row["team_id"],
team_alias=_team_alias_or_none(scope.team_alias_metadata.get(row["team_id"])),
user_id=row["user_id"] or "",
user_email=row["user_email"],
user_alias=row["user_alias"],
spend=row["spend"],
prompt_tokens=row["prompt_tokens"],
completion_tokens=row["completion_tokens"],
total_tokens=row["total_tokens"],
api_requests=row["api_requests"],
successful_requests=row["successful_requests"],
failed_requests=row["failed_requests"],
)
for row in rows
)
return TeamUserSpendResponse(start_date=start_date, end_date=end_date, results=results)
def _team_alias_or_none(metadata: Mapping[str, object] | None) -> str | None:
alias: Final = metadata.get("team_alias") if metadata is not None else None
return alias if isinstance(alias, str) else None

227
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -37,6 +37,7 @@ from litellm.types.utils import (
Message,
ModelResponse,
TextCompletionResponse,
Usage,
)
if TYPE_CHECKING:
@ -148,6 +149,144 @@ class AnthropicPassthroughLoggingHandler:
return model_group.removeprefix("passthrough/")
return model
@staticmethod
def _resolve_logged_model(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
) -> str:
request_model: Final = request_body.get("model")
logged_model: Final = (
request_model
if isinstance(request_model, str) and request_model
else str(litellm_logging_obj.model_call_details.get("model") or "")
)
if logged_model and logged_model != "unknown":
return logged_model
return AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks) or logged_model
@staticmethod
def _usage_only_response_or_none(
all_chunks: Sequence[str | bytes], model: str, speed: str | None
) -> ModelResponse | None:
try:
return AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks, model=model, speed=speed
)
except Exception as e: # noqa: BLE001 # the usage-only fallback must never raise out of failure logging
verbose_proxy_logger.warning("Anthropic passthrough: usage-only fallback failed (model=%s): %s", model, e)
return None
@staticmethod
def _assemble_streaming_response(
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None,
) -> ModelResponse | TextCompletionResponse | None:
try:
assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
except Exception as e: # noqa: BLE001 # any assembly error falls back to usage-only cost
verbose_proxy_logger.warning(
"Anthropic passthrough: stream assembly raised (model=%s): %s; falling "
"back to usage-only cost from raw SSE events.",
model,
e,
)
return AnthropicPassthroughLoggingHandler._usage_only_response_or_none(all_chunks, model, speed)
if assembled is not None:
return assembled
return AnthropicPassthroughLoggingHandler._usage_only_response_or_none(all_chunks, model, speed)
@staticmethod
def _build_streaming_response_for_logging(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
model: str,
) -> ModelResponse | TextCompletionResponse | None:
response: Final = AnthropicPassthroughLoggingHandler._assemble_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body),
)
if response is None:
return None
AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens(
response=response, all_chunks=all_chunks, model=model
)
return response
@staticmethod
def record_partial_usage_for_failure(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
) -> None:
if not all_chunks:
return
model: Final = AnthropicPassthroughLoggingHandler._resolve_logged_model(
litellm_logging_obj, request_body, all_chunks
)
partial_response: Final = AnthropicPassthroughLoggingHandler._build_streaming_response_for_logging(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=all_chunks, model=model
)
usage: Final = cast(Usage | None, getattr(partial_response, "usage", None))
if partial_response is None or usage is None:
return
litellm_logging_obj.record_partial_usage_for_failure(
usage=usage,
response_cost=AnthropicPassthroughLoggingHandler._cost_partial_stream_or_zero(
partial_response=partial_response, model=model, logging_obj=litellm_logging_obj
),
)
@staticmethod
def _cost_partial_stream_or_zero(
partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj
) -> float:
try:
return AnthropicPassthroughLoggingHandler._compute_response_cost(
litellm_model_response=partial_response,
model=AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj),
logging_obj=logging_obj,
)
except Exception as e: # noqa: BLE001 # an uncostable partial stream still bills its tokens, at zero cost
verbose_proxy_logger.warning(
"Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e
)
return 0.0
@staticmethod
def _compute_response_cost(
litellm_model_response: ModelResponse | TextCompletionResponse,
model: str,
logging_obj: LiteLLMLoggingObj,
) -> float:
if logging_obj.model_call_details.get("cache_hit") is True:
return 0.0
custom_llm_provider: Final = logging_obj.model_call_details.get("custom_llm_provider")
model_for_cost: Final = (
f"{custom_llm_provider}/{model}"
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/")
else model
)
return litellm.completion_cost(
completion_response=litellm_model_response,
model=model_for_cost,
custom_llm_provider=custom_llm_provider,
custom_pricing=use_custom_pricing_for_model(
litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None)
),
router_model_id=logging_obj.get_router_model_id(),
)
@staticmethod
def _extract_message_start_field(
all_chunks: Sequence[str | bytes],
@ -278,31 +417,9 @@ class AnthropicPassthroughLoggingHandler:
if logging_obj.model_call_details.get("stream") is True:
logging_obj.model_call_details["complete_streaming_response"] = litellm_model_response
try:
# Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic)
custom_llm_provider: Final = logging_obj.model_call_details.get("custom_llm_provider")
model = AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj)
# Prepend custom_llm_provider to model if not already present
model_for_cost = model
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
model_for_cost = f"{custom_llm_provider}/{model}"
router_model_id: Final = logging_obj.get_router_model_id()
custom_pricing: Final = use_custom_pricing_for_model(
litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None)
)
response_cost: Final = (
0.0
if logging_obj.model_call_details.get("cache_hit") is True
else litellm.completion_cost(
completion_response=litellm_model_response,
model=model_for_cost,
custom_llm_provider=custom_llm_provider,
custom_pricing=custom_pricing,
router_model_id=router_model_id,
)
response_cost: Final = AnthropicPassthroughLoggingHandler._compute_response_cost(
litellm_model_response=litellm_model_response, model=model, logging_obj=logging_obj
)
kwargs["response_cost"] = response_cost
@ -356,57 +473,12 @@ class AnthropicPassthroughLoggingHandler:
- Logs in litellm callbacks
"""
speed: Final = AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body)
model = request_body.get("model", "")
# Check if it's available in the logging object
if (
not model
and hasattr(litellm_logging_obj, "model_call_details")
and litellm_logging_obj.model_call_details.get("model")
):
model = cast(str, litellm_logging_obj.model_call_details.get("model"))
if not model or model == "unknown":
chunk_model: Final = AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks)
if chunk_model:
model = chunk_model
try:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
except Exception as e:
# stream_chunk_builder re-raises assembly failures (as litellm.APIError)
# on large agentic tool-use / thinking streams; treat that the same as a
# None result so the usage-only fallback below still recovers cost
verbose_proxy_logger.warning(
"Anthropic passthrough: stream assembly raised (model=%s): %s; falling "
"back to usage-only cost from raw SSE events.",
model,
e,
)
complete_streaming_response = None
if complete_streaming_response is None:
# stream_chunk_builder cannot always reassemble large agentic streams, but
# Anthropic still emits token usage in the message_start / message_delta SSE
# events regardless of content shape; recover usage-only so cost is tracked.
# Guard it too: a raise here would defeat the point and drop the request
try:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks,
model=model,
speed=speed,
)
except Exception as e:
verbose_proxy_logger.warning(
"Anthropic passthrough: usage-only fallback failed (model=%s): %s",
model,
e,
)
complete_streaming_response = None
model: Final = AnthropicPassthroughLoggingHandler._resolve_logged_model(
litellm_logging_obj, request_body, all_chunks
)
complete_streaming_response: Final = AnthropicPassthroughLoggingHandler._build_streaming_response_for_logging(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=all_chunks, model=model
)
if complete_streaming_response is None:
verbose_proxy_logger.error(
"Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..."
@ -415,11 +487,6 @@ class AnthropicPassthroughLoggingHandler:
"result": None,
"kwargs": {},
}
AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens(
response=complete_streaming_response,
all_chunks=all_chunks,
model=model,
)
kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=complete_streaming_response,
model=model,

81
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -870,6 +870,35 @@ async def _log_passthrough_upstream_failure(
)
async def _relay_reporting_failures(
stream: AsyncGenerator[bytes, None],
upstream_status: int,
user_api_key_dict: UserAPIKeyAuth,
request_payload: dict, # mutable-ok: post_call_failure_hook lifts fields onto request_data in place
) -> AsyncGenerator[bytes, None]:
from litellm.proxy.proxy_server import proxy_logging_obj
try:
async for chunk in stream:
yield chunk
except Exception as e:
if upstream_status >= 400:
raise
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=request_payload,
traceback_str=traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG),
)
except Exception: # noqa: BLE001 - a failing logging callback must never mask the upstream error
verbose_proxy_logger.warning(
"pass_through_endpoint: post_call_failure_hook raised for a mid-stream upstream error",
exc_info=True,
)
raise
from litellm.passthrough.timeout_utils import (
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, # noqa: F401 - re-exported for backward compat
resolve_llm_passthrough_timeout, # noqa: F401 - re-exported for backward compat
@ -1293,14 +1322,24 @@ async def pass_through_request(
return StreamingResponse(
wrap_passthrough_sse_bytes_with_keepalive_pings(
stream=_own_streamed_managed_ids(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
stream=_relay_reporting_failures(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
),
upstream_status=response.status_code,
user_api_key_dict=user_api_key_dict,
request_payload=_build_passthrough_failure_request_payload(
parsed_body=_parsed_body,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
),
),
managed_id_provider=_managed_id_provider,
request=request,
@ -1374,14 +1413,24 @@ async def pass_through_request(
return StreamingResponse(
wrap_passthrough_sse_bytes_with_keepalive_pings(
stream=_own_streamed_managed_ids(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
stream=_relay_reporting_failures(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
),
upstream_status=response.status_code,
user_api_key_dict=user_api_key_dict,
request_payload=_build_passthrough_failure_request_payload(
parsed_body=_parsed_body,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
),
),
managed_id_provider=_managed_id_provider,
request=request,

110
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -1,5 +1,7 @@
from collections.abc import Coroutine
from datetime import datetime
import traceback
from collections.abc import Coroutine, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final, Protocol
import httpx
@ -12,7 +14,7 @@ from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
from litellm.types.utils import StandardPassThroughResponseObject
from litellm.types.utils import StandardPassThroughResponseObject, Usage
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -44,12 +46,85 @@ class RouteStreamingLogging(Protocol):
) -> Coroutine[None, None, None]: ...
@dataclass(frozen=True, slots=True)
class PassThroughStreamContext:
passthrough_success_handler_obj: PassThroughEndpointLogging
url_route: str
start_time: datetime
class PassThroughStreamingHandler:
@staticmethod
def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None:
if litellm_logging_obj.completion_start_time is None:
litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now())
@staticmethod
def schedule_stream_failure_logging(
litellm_logging_obj: LiteLLMLoggingObj,
endpoint_type: EndpointType,
request_body: dict[str, object],
raw_bytes: Sequence[bytes],
exception: Exception,
stream_context: PassThroughStreamContext | None = None,
) -> None:
PassThroughStreamingHandler._record_partial_usage_for_failure(
litellm_logging_obj=litellm_logging_obj,
endpoint_type=endpoint_type,
request_body=request_body,
raw_bytes=raw_bytes,
stream_context=stream_context,
)
try:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=litellm_logging_obj.dispatch_failure_handlers(
exception, traceback.format_exc(), prefer_async_handlers=True
)
)
except Exception as e:
verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e)
@staticmethod
def _record_partial_usage_for_failure(
litellm_logging_obj: LiteLLMLoggingObj,
endpoint_type: EndpointType,
request_body: dict[str, object],
raw_bytes: Sequence[bytes],
stream_context: PassThroughStreamContext | None,
) -> None:
if endpoint_type == EndpointType.ANTHROPIC:
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes
)
return
if stream_context is None or not raw_bytes:
return
try:
partial_response, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=stream_context.passthrough_success_handler_obj,
url_route=stream_context.url_route,
request_body=request_body,
endpoint_type=endpoint_type,
start_time=stream_context.start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(timezone.utc),
model=None,
)
except Exception as e:
verbose_proxy_logger.warning(
"Could not recover the partial usage of a failed %s pass-through stream: %s", endpoint_type.value, e
)
return
usage: Final = getattr(partial_response, "usage", None)
if not isinstance(usage, Usage):
return
response_cost: Final = kwargs.get("response_cost")
litellm_logging_obj.record_partial_usage_for_failure(
usage=usage,
response_cost=float(response_cost) if isinstance(response_cost, (int, float)) else 0.0,
)
@staticmethod
async def chunk_processor(
response: httpx.Response,
@ -65,13 +140,14 @@ class PassThroughStreamingHandler:
route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
)
raw_bytes: Final[list[bytes]] = []
resolved_request_body: Final[dict[str, object]] = request_body or {}
def _build_logging_coroutine() -> Coroutine[None, None, None]:
return resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
request_body=resolved_request_body,
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
@ -132,9 +208,9 @@ class PassThroughStreamingHandler:
# coroutine on logging_obj instead of enqueueing now, so
# ProxyLogging._fire_deferred_stream_logging fires it after
# guardrail end-of-stream blocks populate guardrail_information.
# Disconnect/exception paths skip this and fall through to the
# immediate enqueue in ``finally`` to keep partial billing
# (LIT-2642).
# Disconnect paths skip this and fall through to the immediate
# enqueue in ``finally`` to keep partial billing (LIT-2642);
# upstream exceptions log a failure instead (LIT-3798).
if (
getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None
and raw_bytes
@ -144,6 +220,20 @@ class PassThroughStreamingHandler:
litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
except Exception as e:
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
if response.status_code < 400:
logging_scheduled = True
PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=litellm_logging_obj,
endpoint_type=endpoint_type,
request_body=resolved_request_body,
raw_bytes=raw_bytes,
exception=e,
stream_context=PassThroughStreamContext(
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
start_time=start_time,
),
)
raise
finally:
# GeneratorExit (raised on client disconnect) is not caught by
@ -168,7 +258,7 @@ class PassThroughStreamingHandler:
request_body: dict,
endpoint_type: EndpointType,
start_time: datetime,
raw_bytes: list[bytes],
raw_bytes: Sequence[bytes],
end_time: datetime,
model: str | None = None,
):
@ -218,7 +308,7 @@ class PassThroughStreamingHandler:
request_body: dict,
endpoint_type: EndpointType,
start_time: datetime,
raw_bytes: list[bytes],
raw_bytes: Sequence[bytes],
end_time: datetime,
model: str | None,
) -> tuple[PassThroughEndpointLoggingResultValues, dict]:
@ -336,7 +426,7 @@ class PassThroughStreamingHandler:
return None
@staticmethod
def _convert_raw_bytes_to_str_lines(raw_bytes: list[bytes]) -> list[str]:
def _convert_raw_bytes_to_str_lines(raw_bytes: Sequence[bytes]) -> list[str]:
"""
Converts a list of raw bytes into a list of string lines, similar to aiter_lines()

1
litellm/proxy/proxy_server.py
View file

@ -15212,6 +15212,7 @@ async def async_queue_request(
# extra_body); see above for the same guard upstream.
data["metadata"] = {}
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
_headers: Final = _safe_get_request_headers(request).copy()
_headers.pop("authorization", None) # do not store the original `sk-..` api key in the db

4
litellm/proxy/public_endpoints/autorouter_presets.json
View file

@ -1,12 +1,12 @@
{
"1m_context": {
"label": "1M Context",
"description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Opus 5 for complex, Opus 5 at high thinking for reasoning.",
"description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Sol for complex, Opus 5 at high thinking for reasoning.",
"complexity_router_config": {
"tiers": {
"SIMPLE": ["gpt-5.6-luna"],
"MEDIUM": ["gpt-5.6-terra"],
"COMPLEX": ["claude-opus-5"],
"COMPLEX": ["gpt-5.6-sol"],
"REASONING": ["claude-opus-5"]
},
"tier_model_configs": {

2
litellm/proxy/schema.prisma
View file

@ -1124,6 +1124,8 @@ model LiteLLM_DailyGuardrailUsageUnits {
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
cost Float? // USD for the priced share of units; null only on rows written before this column existed
untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out
created_at DateTime @default(now())
updated_at DateTime @updatedAt

236
litellm/proxy/spend_tracking/key_metadata_recovery.py Normal file
View file

@ -0,0 +1,236 @@
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from types import MappingProxyType
from typing import Final, TypeVar
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy.utils import PrismaClient
from litellm.repositories.user_repository import UserRepository
_T = TypeVar("_T")
_ACTIVE_TOKEN_DIGEST_SQL: Final = """
SELECT encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
FROM "LiteLLM_VerificationToken"
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
"""
_DELETED_TOKEN_DIGEST_SQL: Final = """
SELECT DISTINCT ON (token)
encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
FROM "LiteLLM_DeletedVerificationToken"
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
ORDER BY token, deleted_at DESC
"""
class KeyMetadataDict(TypedDict, total=False):
key_alias: ReadOnly[str | None]
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
class _TokenDigestRow(BaseModel):
digest: str
key_alias: str | None = None
team_id: str | None = None
user_id: str | None = None
_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...])
_EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({})
_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({})
async def _db_or_empty(
load: Callable[[], Awaitable[_T]],
warning: str,
count: int,
) -> _T | None:
from prisma.errors import PrismaError
try:
return await load()
except PrismaError as e:
verbose_proxy_logger.warning(warning, count, e)
return None
async def _reverse_hash_key_metadata(
prisma_client: PrismaClient,
sql: str,
wanted: AbstractSet[str],
*,
warning: str,
) -> Mapping[str, KeyMetadataDict]:
rows: Final = await _db_or_empty(
lambda: prisma_client.db.query_raw(sql, sorted(wanted)),
warning,
len(wanted),
)
if rows is None:
return _EMPTY_KEY_METADATA
return MappingProxyType(
{
row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id)
for row in _TOKEN_DIGEST_ROWS.validate_python(rows)
if row.digest in wanted
}
)
async def _emails_for_user_ids(
prisma_client: PrismaClient,
user_ids: AbstractSet[str],
) -> Mapping[str, str]:
if not user_ids:
return _EMPTY_EMAILS
users: Final = await _db_or_empty(
lambda: UserRepository(prisma_client).table.find_many(
where={"user_id": {"in": list(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict
),
"Failed user_email recovery for %d user ids: %s",
len(user_ids),
)
if users is None:
return _EMPTY_EMAILS
return MappingProxyType(
{
user.user_id: user.user_email
for user in users
if getattr(user, "user_id", None) and getattr(user, "user_email", None)
}
)
def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict:
if meta.get("user_email"):
return meta
user_id: Final = meta.get("user_id")
if not isinstance(user_id, str) or user_id not in emails:
return meta
updated: Final[KeyMetadataDict] = {**meta, "user_email": emails[user_id]}
return updated
async def attach_user_emails(
prisma_client: PrismaClient,
recovered: Mapping[str, KeyMetadataDict],
) -> Mapping[str, KeyMetadataDict]:
needing_email: Final = frozenset(
user_id
for meta in recovered.values()
for user_id in (meta.get("user_id"),)
if isinstance(user_id, str) and user_id and not meta.get("user_email")
)
emails: Final = await _emails_for_user_ids(prisma_client, needing_email)
if not emails:
return recovered
return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()})
async def recover_double_hashed_key_metadata(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
) -> Mapping[str, KeyMetadataDict]:
"""
Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that
were double-hashed by the v1.99 spend-log provenance gate.
Those rows store hash(VerificationToken.token) instead of the token, so the
exact join misses. Postgres hashes the token column itself, one pass over
active keys and one over deleted keys, so no key row crosses the wire.
"""
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
if not sha_missing:
return _EMPTY_KEY_METADATA
from_active: Final = await _reverse_hash_key_metadata(
prisma_client,
_ACTIVE_TOKEN_DIGEST_SQL,
sha_missing,
warning="Failed reverse-hash recovery against active keys for %d missing keys: %s",
)
still_missing: Final = sha_missing - frozenset(from_active)
if not still_missing:
return from_active
from_deleted: Final = await _reverse_hash_key_metadata(
prisma_client,
_DELETED_TOKEN_DIGEST_SQL,
still_missing,
warning="Failed reverse-hash recovery against deleted keys for %d missing keys: %s",
)
return MappingProxyType({**from_active, **from_deleted})
def _row_with_recovered_fields(
row: Mapping[str, object],
recovered: Mapping[str, KeyMetadataDict],
*,
api_key_field: str,
alias_field: str,
team_id_field: str,
user_email_field: str,
) -> Mapping[str, object]:
api_key: Final = row.get(api_key_field)
if not isinstance(api_key, str) or api_key not in recovered:
return row
meta: Final = recovered[api_key]
return MappingProxyType(
{
**row,
alias_field: meta.get("key_alias") or row.get(alias_field),
team_id_field: meta.get("team_id") or row.get(team_id_field),
user_email_field: row.get(user_email_field) or meta.get("user_email"),
}
)
async def fill_missing_api_key_aliases(
prisma_client: PrismaClient,
rows: Sequence[Mapping[str, object]],
*,
api_key_field: str = "api_key",
alias_field: str = "api_key_alias",
team_id_field: str = "team_id",
user_email_field: str = "user_email",
) -> tuple[Mapping[str, object], ...]:
"""
Fill null api_key_alias / team_id / user_email on export rows whose api_key
was double-hashed.
Used by CloudZero and Focus, which join DailyUserSpend.api_key to
VerificationToken.token and otherwise export null aliases for those rows.
"""
missing_keys: Final = frozenset(
key
for row in rows
for key in (row.get(api_key_field),)
if isinstance(key, str) and key and row.get(alias_field) in (None, "")
)
if not missing_keys:
return tuple(rows)
recovered: Final = await attach_user_emails(
prisma_client,
await recover_double_hashed_key_metadata(prisma_client, missing_keys),
)
if not recovered:
return tuple(rows)
return tuple(
_row_with_recovered_fields(
row,
recovered,
api_key_field=api_key_field,
alias_field=alias_field,
team_id_field=team_id_field,
user_email_field=user_email_field,
)
for row in rows
)

2
litellm/router.py
View file

@ -7513,7 +7513,7 @@ class Router:
# Check retry policy FIRST, before should_retry_this_error
# This allows retry policies to override the healthy deployments check
_retry_policy_applies = False
if self.retry_policy is not None or model_group_retry_policy is not None:
if request_num_retries != 0 and (self.retry_policy is not None or model_group_retry_policy is not None):
# get num_retries from retry policy
# Use the model_group captured at the start of the function, or get it from metadata
# kwargs.get("model") at this point is the deployment model, not the model_group

9
litellm/router_strategy/complexity_router/README.md
View file

@ -275,6 +275,15 @@ model_list:
keep the classifier deployment or provider default, or set a supported value such as `none` or
`low` to override that call.
Classifier calls have a one-attempt hard deadline. After a timeout, the router opens a process-local
circuit for that classifier and sends every session through `classifier_fallback` for
`classifier_llm_config.circuit_breaker_cooldown_seconds` (30 seconds by default). When the cooldown
expires, one request probes the classifier while concurrent requests continue through the fallback.
A successful probe closes the circuit; a failed probe restarts the cooldown. The circuit breaker is
on by default; set `classifier_llm_config.circuit_breaker_enabled: false` to disable it. The default
fallback is the local heuristic scorer, so a classifier outage does not repeat its timeout across
every turn or session handled by the router process.
A request short-circuits, meaning it routes on the scorer's own tier with no classifier call, when
two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least
one signal. Everything else goes to the classifier, which then decides as it normally would.

361
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -18,8 +18,10 @@ from __future__ import annotations
import asyncio
import random
import re
from collections.abc import Iterator, Mapping, Sequence
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import accumulate, islice, takewhile
from threading import Lock
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
@ -33,7 +35,10 @@ from litellm.constants import (
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
@ -311,6 +316,8 @@ _TRUNCATION_MARKER: Final = "..."
_TRUNCATION_HEAD_FRACTION: Final = 0.3
_MIN_QUOTED_TURN_CHARS: Final = 120
_CLASSIFIER_CIRCUIT_OPEN_SIGNAL: Final = "classifier-circuit-open"
_CJK_CHARACTER: Final = re.compile("[぀-ヿㇰ-ㇿ㐀-䶿一-鿿豈-﫿ｦ-ﾝ\U00020000-\U0003ffff]")
@ -755,6 +762,17 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
image), not what the session's traffic looks like, and pinning it would hold every following
text turn on the vision-capable model the image forced. A modality pin override is the same
fact on a session that already holds a pin, so it must not overwrite the pin it displaced.
An open classifier circuit is the shortest-lived state of all: the fallback ran because the
breaker skipped the classifier, not because the request got classified, and the cooldown is
seconds against a TTL of an hour that every later turn refreshes. Its cause is whatever the
fallback path reports, so the circuit signal is what marks the decision, and leaving it
unpinned lets the session classify again as soon as the breaker closes.
A health failover describes the fleet's state right now, not the session's traffic, and it can
displace decisions that were themselves unpinnable (a housekeeping call, a modality escalation).
Pinning it would hold the session on the substitute long after the displaced group recovers; the
gate re-fires per request, so leaving it unpinned costs nothing but the classifier call.
"""
return decision is None or (
decision.get("cause")
@ -764,8 +782,10 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
"housekeeping",
"modality_escalation",
"modality_pin_override",
"health_failover",
)
and not decision.get("context_escalated")
and _CLASSIFIER_CIRCUIT_OPEN_SIGNAL not in (decision.get("signals") or ())
)
@ -816,6 +836,81 @@ class ClassificationOutcome(NamedTuple):
classifier_cost: float | None = None
def _with_signal(outcome: ClassificationOutcome, signal: str | None) -> ClassificationOutcome:
return outcome if signal is None else outcome._replace(signals=(*outcome.signals, signal))
class _ClassifierCircuitBreaker:
"""Process-local timeout breaker for one complexity-router classifier.
The router instance serves every session assigned to that auto-router deployment, so the
breaker prevents one unhealthy classifier from charging the same timeout to each session.
Exactly one request becomes the recovery probe after the cooldown; the lock makes that state
transition atomic even when several request tasks arrive together.
"""
CLOSED: Final = "closed"
OPEN: Final = "open"
HALF_OPEN: Final = "half_open"
def __init__(self, cooldown_seconds: float, clock: Callable[[], float] = time.monotonic) -> None:
self._cooldown_seconds = cooldown_seconds
self._clock = clock
self._state = self.CLOSED
self._opened_at: float | None = None
self._generation = 0
self._lock = Lock()
def acquire_permit(self) -> int | None:
"""Return a generation-scoped permit, or deny the call while the circuit is open.
Calls admitted together while closed share a generation. The first timeout advances it,
making every other in-flight completion stale so it cannot erase the new cooldown.
"""
with self._lock:
if self._state == self.CLOSED:
return self._generation
if self._state == self.HALF_OPEN:
return None
opened_at: Final = self._opened_at
if opened_at is not None and self._clock() - opened_at >= self._cooldown_seconds:
self._state = self.HALF_OPEN
return self._generation
return None
def record_success(self, permit: int) -> None:
"""Close only when the current half-open recovery probe succeeds."""
with self._lock:
if self._state != self.HALF_OPEN or permit != self._generation:
return
self._state = self.CLOSED
self._opened_at = None
def record_failure(self, permit: int, *, is_timeout: bool) -> None:
"""Open on a normal timeout, or reopen when the single recovery probe fails."""
with self._lock:
if permit != self._generation:
return
if self._state == self.CLOSED:
if not is_timeout:
return
elif self._state != self.HALF_OPEN:
return
self._generation += 1
self._state = self.OPEN
self._opened_at = self._clock()
def _is_classifier_timeout(exc: BaseException) -> bool:
# asyncio.TimeoutError became an alias of the built-in TimeoutError in Python 3.11.
# LiteLLM still supports 3.10, where they are distinct exception classes.
if isinstance(exc, (TimeoutError, asyncio.TimeoutError)):
return True
from litellm.exceptions import Timeout as LiteLLMTimeout
return isinstance(exc, LiteLLMTimeout)
def _allowed(models: tuple[str, ...], fit_filter: frozenset[str] | None) -> tuple[str, ...]:
return models if fit_filter is None else tuple(model for model in models if model in fit_filter)
@ -993,6 +1088,15 @@ class ComplexityRouter(CustomLogger):
if llm_classifier_configured
else None
)
self._classifier_circuit_breaker: _ClassifierCircuitBreaker | None = (
_ClassifierCircuitBreaker(self.config.classifier_llm_config.circuit_breaker_cooldown_seconds)
if (
llm_classifier_configured
and self.config.classifier_llm_config is not None
and self.config.classifier_llm_config.circuit_breaker_enabled
)
else None
)
self._tier_success_predictor: TierSuccessPredictor | None = (
TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact))
if self.config.classifier_type == "heuristic_v2"
@ -1474,8 +1578,20 @@ class ComplexityRouter(CustomLogger):
`scored` is the heuristic outcome the caller already computed, which only "heuristic_first"
has. It is handed to the failure path so a classifier error does not re-run the scorer.
"""
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._classifier_failure_outcome(
"LLM classifier circuit is open",
prompt,
system_prompt,
scored,
signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL,
)
try:
tier, classifier_cost = await self._classify_with_llm(prompt, system_prompt, request_kwargs, messages)
if breaker is not None and permit is not None:
breaker.record_success(permit)
return ClassificationOutcome(
tier=tier,
score=None,
@ -1483,7 +1599,13 @@ class ComplexityRouter(CustomLogger):
cause="llm_classifier",
classifier_cost=classifier_cost,
)
except asyncio.CancelledError:
if breaker is not None and permit is not None:
breaker.record_failure(permit, is_timeout=False)
raise
except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path
if breaker is not None and permit is not None:
breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e))
return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored)
def _classifier_failure_outcome(
@ -1492,6 +1614,7 @@ class ComplexityRouter(CustomLogger):
prompt: str,
system_prompt: str | None,
scored: ClassificationOutcome | None = None,
signal: str | None = None,
) -> ClassificationOutcome:
"""The outcome when the LLM classifier or classifier plugin produced no usable tier:
fallback_tier on a custom tier set, classifier_fallback otherwise.
@ -1501,21 +1624,24 @@ class ComplexityRouter(CustomLogger):
fallback_tier: Final = self.config.fallback_tier
if fallback_tier is not None:
verbose_router_logger.warning("ComplexityRouter: %s, routing to fallback_tier %s", reason, fallback_tier)
return ClassificationOutcome(
tier=fallback_tier,
score=None,
signals=(f"classifier-fallback:{fallback_tier}",),
cause="classifier_fallback",
return _with_signal(
ClassificationOutcome(
tier=fallback_tier,
score=None,
signals=(f"classifier-fallback:{fallback_tier}",),
cause="classifier_fallback",
),
signal,
)
verbose_router_logger.warning(
"ComplexityRouter: %s, falling back to %s", reason, self.config.classifier_fallback
)
if self.config.classifier_fallback == "default_model":
return self._default_model_fallback_outcome()
return _with_signal(self._default_model_fallback_outcome(), signal)
if scored is not None:
return scored
return _with_signal(scored, signal)
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
return _with_signal(ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause), signal)
async def _classify_with_plugin(
self,
@ -1694,16 +1820,23 @@ class ComplexityRouter(CustomLogger):
}
}
response: Final[ModelResponse] = await self.litellm_router_instance.acompletion(
model=llm_config.model,
messages=messages_for_call,
response_format=response_format,
timeout=llm_config.timeout_ms / 1000,
metadata=metadata,
proxy_server_request=proxy_server_request,
turn_off_message_logging=turn_off_message_logging,
**classifier_call_params,
**_parent_session_kwargs(request_kwargs),
classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000
response: Final[ModelResponse] = await asyncio.wait_for(
self.litellm_router_instance.acompletion(
model=llm_config.model,
messages=messages_for_call,
stream=False,
response_format=response_format,
timeout=classifier_timeout_s,
num_retries=0,
disable_fallbacks=True,
metadata=metadata,
proxy_server_request=proxy_server_request,
turn_off_message_logging=turn_off_message_logging,
**classifier_call_params,
**_parent_session_kwargs(request_kwargs),
),
timeout=classifier_timeout_s,
)
content: Final = response.choices[0].message.content
if not content:
@ -2526,6 +2659,150 @@ class ComplexityRouter(CustomLogger):
and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None
)
async def _model_group_can_serve(
self,
model_name: str,
messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the router's own probe
input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim
request_kwargs: dict, # mutable-ok: same shape the hook receives
) -> bool:
"""Whether the router would find a deployment for this group ON THIS REQUEST.
Asks the same owner the routing path itself will ask, with the same prompt arguments it
will pass, so every filter that decides a deployment's eligibility applies here exactly
as it applies downstream: cooldowns, admin pause, team scoping, model access groups, tag
routing, routing plugins, RPM limits, and the context-window pre-call check. Re-deriving
any subset of that list is how a substitute gets chosen that the pipeline then rejects,
and dropping `input` would silently skip the window check on the Responses API surface,
where the prompt never arrives as messages.
Probed on a COPY of request_kwargs because the owner pops routing bookkeeping off the
dict it is handed (`_target_order`, `_excluded_deployment_ids`), and this is a
speculative question about a model that may never be picked.
Every way the owner says "nothing here can serve this" is a negative verdict: no healthy
deployment for the group at all (BadRequestError, which ContextWindowExceededError
subclasses), every deployment filtered out (RouterRateLimitError), and every deployment
over its RPM (RouterRateLimitErrorBasic). Anything else is unknown rather than negative,
so it reads as capacity: absent information must never decide the verdict.
"""
from litellm.exceptions import BadRequestError
from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic
probe_kwargs: Final = dict(request_kwargs) # mutable-ok: the owner pops routing keys off the dict it is handed
try:
deployments: Final = await self.litellm_router_instance.async_get_healthy_deployments(
model=model_name,
request_kwargs=probe_kwargs,
messages=messages,
input=input,
parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs),
)
except (RouterRateLimitError, RouterRateLimitErrorBasic, BadRequestError):
return False
except Exception as exc: # noqa: BLE001 # a speculative eligibility read must fail open on unknown faults
verbose_router_logger.debug(
"ComplexityRouter: eligibility probe for %s failed, treating the group as live: %s", model_name, exc
)
return True
return bool(deployments)
async def _gate_response_health(
self,
response: PreRoutingHookResponse,
messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick
input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim
resolved_messages: Sequence[Mapping[str, object]] | None,
request_kwargs: dict, # mutable-ok: same shape the hook receives
) -> PreRoutingHookResponse:
"""Replace a decided model group that has no serving capacity with a live peer in the same tier.
Applied to the decided response at the hook's exits, so every arm that can place a request
is covered by one owner: a fresh classification, a replayed or escalated session pin, a
plan-mode floor, a context-window escalation, an adaptive pick, and whatever arm is added
next. Peers come from the DECIDED tier only; climbing to another tier is deliberately not
done here, since a higher tier costs more than the classifier asked for.
Serving capacity is one question asked of one owner (`_model_group_can_serve`), so the
substitute is only ever a group the pipeline would actually accept for this request. The
pick then runs through `_pick_model_for_tier`, so routing plugins decide the substitute
exactly as they decided the original.
Fails open everywhere it cannot be sure: an unreadable eligibility view, a decision
carrying no tier (default_model), or a tier whose every peer is unusable too. It fails
CLOSED on a plugin that empties the pool, leaving the original decision to fail rather
than serving a model the plugin excluded.
"""
decision: Final = response.routing_decision
decided_tier: Final = decision.get("tier") if decision is not None else None
if decision is None or not isinstance(decided_tier, str):
return response
peers: Final = tuple(self._tier_pools().get(decided_tier, ()))
if len(peers) < 2:
return response
if await self._model_group_can_serve(response.model, messages, input, request_kwargs):
return response
eligible: Final = (
self._modality_eligible_models()
if self.config.modality_routing and resolved_messages and request_contains_image_content(resolved_messages)
else None
)
candidates: Final = tuple(
peer for peer in peers if peer != response.model and (eligible is None or peer in eligible)
)
if not candidates:
return response
servable: Final = await asyncio.gather(
*(self._model_group_can_serve(peer, messages, input, request_kwargs) for peer in candidates)
)
live: Final = tuple(peer for peer, can_serve in zip(candidates, servable) if can_serve)
if not live:
return response
repick_messages: Final = (
list(resolved_messages) if resolved_messages else None # mutable-ok: the pick's param is list-typed
)
try:
new_model: Final = await self._pick_model_for_tier(
decided_tier if self.config.has_custom_tiers else ComplexityTier(decided_tier),
messages,
repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them
request_kwargs,
allowed_models=live,
)
except ValueError as exc:
verbose_router_logger.debug(
"ComplexityRouter: health failover found no candidate the routing plugins allow: %s", exc
)
return response
self._restamp_adaptive_choice(request_kwargs, response.model, new_model)
verbose_router_logger.info(
"ComplexityRouter: routing decision cause=health_failover, routed_model=%s, displaced=%s",
new_model,
response.model,
)
new_decision: Final = self._build_routing_decision(
routed_model=new_model,
cause="health_failover",
tier=decision.get("tier"),
score=decision.get("score"),
signals=(*(decision.get("signals") or ()), f"health_displaced:{response.model}"),
matched_keyword=decision.get("matched_keyword"),
escalation_keyword=decision.get("escalation_keyword"),
escalated=bool(decision.get("escalated", False)),
classifier_model=decision.get("classifier_model"),
classifier_cost=decision.get("classifier_cost"),
conversation_continuing=bool(decision.get("conversation_continuing", True)),
tier_litellm_params=self._litellm_params_for_model(decided_tier, new_model),
context_escalation_original_tier=decision.get("context_escalation_original_tier"),
)
return response.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict
"model": new_model,
"litellm_params": self._litellm_params_for_model(decided_tier, new_model),
"routing_decision": new_decision,
}
)
def _placed_default_model(self) -> str:
"""The default_model behind a usable-default verdict; the raise is the type-level
proof, not a reachable path."""
@ -2923,24 +3200,30 @@ class ComplexityRouter(CustomLogger):
session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model)
has_original_messages: Final = messages is not None and len(messages) > 0
return self._with_session_deployment_affinity(
await self._gate_response_modality(
PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
litellm_params=session_tier_litellm_params,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=routed_pin_tier,
matched_keyword=pin_plan_sentinel if plan_floored else None,
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
tier_litellm_params=session_tier_litellm_params,
context_escalation_original_tier=pin_context_original_tier,
await self._gate_response_health(
await self._gate_response_modality(
PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
litellm_params=session_tier_litellm_params,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=routed_pin_tier,
matched_keyword=pin_plan_sentinel if plan_floored else None,
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
tier_litellm_params=session_tier_litellm_params,
context_escalation_original_tier=pin_context_original_tier,
),
),
messages,
resolved_messages,
request_kwargs,
),
messages,
input,
resolved_messages,
request_kwargs,
)
@ -2956,7 +3239,13 @@ class ComplexityRouter(CustomLogger):
resolved_messages=resolved_messages,
)
response: Final = (
await self._gate_response_modality(routed_response, messages, resolved_messages, request_kwargs)
await self._gate_response_health(
await self._gate_response_modality(routed_response, messages, resolved_messages, request_kwargs),
messages,
input,
resolved_messages,
request_kwargs,
)
if routed_response is not None
else None
)

17
litellm/router_strategy/complexity_router/config.py
View file

@ -444,6 +444,23 @@ class ClassifierLLMConfig(BaseModel):
default=3000,
description="Timeout budget for the classification call, in milliseconds",
)
circuit_breaker_enabled: bool = Field(
default=True,
description=(
"Whether one classifier timeout temporarily sends requests through classifier_fallback. "
"Enabled by default so an unhealthy classifier cannot repeat its timeout across sessions."
),
)
circuit_breaker_cooldown_seconds: float = Field(
default=30.0,
gt=0.0,
description=(
"How long to skip this router's LLM classifier after a classification call times out. "
"Requests use classifier_fallback during the cooldown. When it expires, one request "
"probes the classifier while concurrent requests keep using the fallback; a successful "
"probe closes the circuit and a failed probe restarts the cooldown."
),
)
classification_rubric: ClassificationRubric | None = Field(
default=None,
description=(

1
litellm/types/proxy/management_endpoints/common_daily_activity.py
View file

@ -43,6 +43,7 @@ class KeyMetadata(BaseModel):
key_alias: str | None = None
team_id: str | None = None
user_email: str | None = None
class KeyMetricWithMetadata(MetricBase):

21
litellm/types/proxy/management_endpoints/team_endpoints.py
View file

@ -143,3 +143,24 @@ class TeamMetadataSchemaResponse(BaseModel):
"""Response for GET /team/metadata_schema; ``fields`` is empty when no schema is configured."""
fields: tuple[TeamMetadataFieldSchema, ...]
class TeamUserSpendRow(BaseModel):
team_id: str
team_alias: str | None = None
user_id: str
user_email: str | None = None
user_alias: str | None = None
spend: float = 0.0
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
api_requests: int = 0
successful_requests: int = 0
failed_requests: int = 0
class TeamUserSpendResponse(BaseModel):
start_date: str
end_date: str
results: tuple[TeamUserSpendRow, ...]

11
litellm/types/utils.py
View file

@ -2886,6 +2886,10 @@ RoutingDecisionCause = Literal[
# carries an image the pinned model cannot accept. The stored pin is untouched, so the next
# text turn replays it. Distinct from "modality_escalation", which never displaces a pin.
"modality_pin_override",
# Every deployment behind the decided model group was in cooldown, so a healthy peer in the
# same tier served instead. The displaced group rides in signals. Reported even on a kept
# session pin, since the pinned model did not serve the request.
"health_failover",
"session_affinity_pin",
"session_affinity_escalation",
# classification_mode 'user_turn': the request is an agent loop's continuation turn (no new
@ -3151,6 +3155,12 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False):
provider hook. Summed into the request's ``response_cost`` so it counts against
spend and budgets like token cost, unless ``guardrail_cost_in_spend`` is False."""
guardrail_cost_by_unit: ReadOnly[Mapping[str, float | None] | None]
"""``guardrail_cost`` split per ``guardrail_usage`` counter, so the daily
per-counter usage rollup can carry cost at its own grain. Absent when the
hook had no pricing for the invocation; a counter is None when the pricing
entry has no price for it, which the rollup stores as unknown rather than $0."""
guardrail_cost_in_spend: ReadOnly[bool | None]
"""Whether ``guardrail_cost`` participates in the request's ``response_cost`` and
the spend/budget aggregates built from it. Absent, None, or True keeps the default
@ -3202,6 +3212,7 @@ class GuardrailTracingDetail(TypedDict, total=False):
guardrail_action: str | None
guardrail_usage: ReadOnly[Mapping[str, int] | None]
guardrail_cost: ReadOnly[float | None]
guardrail_cost_by_unit: ReadOnly[Mapping[str, float | None] | None]
guardrail_cost_in_spend: ReadOnly[bool | None]

4914
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

20
model_prices_and_context_window.schema.json
View file

@ -79,6 +79,11 @@
"minimum": 0,
"description": "USD per token written to the provider's prompt cache."
},
"cache_creation_input_token_cost_above_128k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_1hr": {
"type": "number",
"minimum": 0,
@ -94,6 +99,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_272k_tokens": {
"type": "number",
"minimum": 0,
@ -128,6 +138,11 @@
"minimum": 0,
"description": "USD per prompt token served from the provider's prompt cache."
},
"cache_read_input_token_cost_above_128k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_200k_tokens": {
"type": "number",
"minimum": 0,
@ -138,6 +153,11 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_read_input_token_cost_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_272k_tokens": {
"type": "number",
"minimum": 0,

2
schema.prisma
View file

@ -1124,6 +1124,8 @@ model LiteLLM_DailyGuardrailUsageUnits {
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
cost Float? // USD for the priced share of units; null only on rows written before this column existed
untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out
created_at DateTime @default(now())
updated_at DateTime @updatedAt

1
terraform/provider/tools/endpointaudit/coverage_allowlist.txt
View file

@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics
GET /tag/wau
GET /team/daily/activity
GET /team/daily/activity/aggregated
GET /team/spend/by_user
GET /team/spend/report
GET /user/daily/activity
GET /user/daily/activity/aggregated

58
tests/proxy_behavior/management/test_team_spend_by_user.py Normal file
View file

@ -0,0 +1,58 @@
import pytest
from .actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# GET /team/spend/by_user shares the team-scope resolver with
# /team/daily/activity, so the membership matrix must hold here too. team_ids
# is mandatory on this route (a per-user rollup with no team is meaningless),
# so the bare query is 400 for everyone instead of defaulting to own teams.
_MEMBERS = {
"alpha": {
Actor.TEAM_ADMIN,
Actor.INTERNAL_USER,
Actor.OWNER,
Actor.UNRELATED_SAME_ORG,
Actor.SERVICE_ACCOUNT,
},
"beta": {Actor.CROSS_ORG_USER},
}
def _expected(actor: Actor, team: str) -> int:
if team == "none":
return 400
if actor == Actor.PROXY_ADMIN:
return 200
return 200 if actor in _MEMBERS.get(team, set()) else 404
_CASES = [
(f"{team}/{actor.value}", actor, team, _expected(actor, team))
for team in ("none", "alpha", "beta")
for actor in Actor
]
_DATES = "start_date=2024-01-01&end_date=2024-12-31"
@pytest.mark.parametrize(
"actor,team,expected_status",
[(a, t, s) for (_id, a, t, s) in _CASES],
ids=[c[0] for c in _CASES],
)
async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, proxy_client, world):
team_id = {"alpha": world.team_alpha_id, "beta": world.team_beta_id}.get(team)
query = _DATES if team_id is None else f"{_DATES}&team_ids={team_id}"
resp = await proxy_client.get(
f"/team/spend/by_user?{query}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}"
if expected_status == 200:
body = resp.json()
assert (body["start_date"], body["end_date"]) == ("2024-01-01", "2024-12-31")
assert all(row["team_id"] == team_id for row in body["results"])

43
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -2445,6 +2445,7 @@ class TestBatchCostAttribution:
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key"] == "hash-alice"
assert metadata["user_api_key_hash"] == "hash-alice"
assert metadata["user_api_key_user_id"] == "alice"
assert metadata["user_api_key_team_id"] == "team-alpha"
assert metadata["user_api_key_alias"] == "prod-key"
@ -2553,6 +2554,48 @@ class TestBatchCostAttribution:
assert metadata["user_api_key_alias"] == "prod-key"
@pytest.mark.asyncio
async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self):
"""
CheckBatchCost stores the VerificationToken hash on the managed object. The
spend-log writer must receive matching user_api_key_hash provenance so it
does not re-hash that value; otherwise DailyUserSpend.api_key no longer joins
VerificationToken and Usage shows key-hash-... with a null alias/email.
"""
from datetime import datetime, timezone
from types import SimpleNamespace
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import hash_token
token_hash = hash_token("sk-batch-creator-key")
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key"),
user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
)
metadata = await instance._build_creator_attribution_metadata(
self._job(api_key=token_hash), "batch-1"
)
assert metadata["user_api_key"] == token_hash
assert metadata["user_api_key_hash"] == token_hash
payload = get_logging_payload(
kwargs={
"model": "gpt-4o",
"call_type": "aretrieve_batch",
"litellm_params": {"metadata": metadata},
},
response_obj={
"id": "batch_123",
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
start_time=datetime.now(timezone.utc),
end_time=datetime.now(timezone.utc),
)
assert payload["api_key"] == token_hash
assert payload["api_key"] != hash_token(token_hash)
class TestPollPageStarvation:
"""LIT-5462 regression: a row that can never be costed used to keep its slot in the

83
tests/proxy_unit_tests/test_jwt_key_mapping.py
View file

@ -1333,3 +1333,86 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate():
virtual_key_claim_field="new_field",
)
assert auth.virtual_key_claim_field == "new_field"
# ──────────────────────────────────────────────
# Tests: cache eviction must happen AFTER the DB write commits
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_delete_evicts_cache_after_row_is_gone():
"""A JWT request racing the delete must not keep the removed mapping authorized.
The DB delete simulates a concurrent request re-caching the mapping mid-write.
If the endpoint evicts before the delete commits, that repopulated entry
survives until TTL and the deleted mapping stays usable.
"""
from litellm.proxy._types import DeleteJWTKeyMappingRequest
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
cache_key = jwt_key_mapping_cache_key("email", "user@example.com")
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token")
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
async def concurrent_reader_repopulates(**kwargs):
await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token")
return _mock_mapping()
mock_prisma.db.litellm_jwtkeymapping.delete.side_effect = concurrent_reader_repopulates
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point
):
result = await delete_jwt_key_mapping(
data=DeleteJWTKeyMappingRequest(id="mapping-1"),
user_api_key_dict=_make_admin_auth(),
)
assert result == {"status": "success"}
assert await user_api_key_cache.async_get_cache(cache_key) is None
@pytest.mark.asyncio
async def test_update_evicts_old_and_new_cache_keys_after_write():
"""Renaming a mapping's claim must leave neither claim serving stale cache.
The DB update simulates a concurrent request re-caching the OLD mapping
mid-write. Both the old claim's entry (would restore the pre-rename token)
and the new claim's __NO_MAPPING__ sentinel (would 403 the renamed claim)
must be gone once the endpoint returns.
"""
from litellm.proxy._types import UpdateJWTKeyMappingRequest
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
old_cache_key = jwt_key_mapping_cache_key("email", "user@example.com")
new_cache_key = jwt_key_mapping_cache_key("email", "renamed@example.com")
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token")
await user_api_key_cache.async_set_cache(key=new_cache_key, value="__NO_MAPPING__")
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
async def concurrent_reader_repopulates(**kwargs):
await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token")
return _mock_mapping(claim_value="renamed@example.com")
mock_prisma.db.litellm_jwtkeymapping.update.side_effect = concurrent_reader_repopulates
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point
):
result = await update_jwt_key_mapping(
data=UpdateJWTKeyMappingRequest(id="mapping-1", jwt_claim_value="renamed@example.com"),
user_api_key_dict=_make_admin_auth(),
)
assert result.jwt_claim_value == "renamed@example.com"
assert await user_api_key_cache.async_get_cache(old_cache_key) is None
assert await user_api_key_cache.async_get_cache(new_cache_key) is None

10
tests/router_unit_tests/test_router_index_management.py
View file

@ -237,10 +237,12 @@ class TestRouterIndexManagement:
- model_name_to_deployment_indices for O(1) + O(k) model_name lookups
"""
# Methods that are allowed to iterate through self.model_list
ALLOWED_METHODS = [
"_get_deployment_by_litellm_model", # Edge case: lookup by litellm_params.model (not indexed)
"_finalize_adaptive_router_if_configured", # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match)
]
ALLOWED_METHODS = {
"_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed",
"_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match',
"config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)",
"heuristic_v2_router_limit_violation": "counts heuristic_v2 routers across the whole list; admin path only (auto-router init/upsert)",
}
# Get path to router.py
router_file = os.path.join(

16
tests/rust-python-harness/shared/native_build.py
View file

@ -73,6 +73,15 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]:
return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
def trace_bridge_error() -> str | None:
bridge: Final = get_native_bridge()
if bridge is None:
return "native Rust bridge is not importable"
if getattr(bridge, "_trace", None) is None:
return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature"
return None
def ensure_trace_bridge(repo_root: Path) -> str | None:
native_path: Final = _native_module_path()
native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
@ -84,9 +93,4 @@ def ensure_trace_bridge(repo_root: Path) -> str | None:
if not succeeded:
return f"native Rust bridge rebuild failed:\n{output}"
_drop_imported_bridge()
bridge: Final = get_native_bridge()
if bridge is None:
return "native Rust bridge is not importable"
if getattr(bridge, "_trace", None) is None:
return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature"
return None
return trace_bridge_error()

13
tests/store_model_in_db_tests/test_openai_error_handling.py
View file

@ -106,15 +106,22 @@ def test_missing_model_parameter_curl(curl_command):
# Run the curl command and capture the output
key = generate_key_sync()
curl_command = curl_command.replace("sk-1234", key)
result = subprocess.run(curl_command, shell=True, capture_output=True, text=True)
result = subprocess.run(
f'{curl_command} -s -w "\\n%{{http_code}}"',
shell=True,
capture_output=True,
text=True,
)
body, _, status_code = result.stdout.rpartition("\n")
# Parse the JSON response
response = json.loads(result.stdout)
response = json.loads(body)
# Check that we got an error response
assert "error" in response
print("error in response", json.dumps(response, indent=4))
assert "litellm.BadRequestError" in response["error"]["message"]
assert status_code == "400", f"expected HTTP 400, got {status_code}: {response}"
assert isinstance(response["error"]["message"], str) and response["error"]["message"]
@pytest.mark.asyncio

15
tests/test_litellm/caching/test_redis_cache.py
View file

@ -823,15 +823,26 @@ async def test_event_loop_stall_timeout_burst_keeps_breaker_closed():
Every operation already waiting on the loop times out together when the loop resumes,
so a purely consecutive threshold is satisfied instantly even though the Redis on the
other end (here an in-process fake that answers immediately) is healthy.
The fake checks its own client deadline against the clock, the way a client library
does, rather than wrapping the call in asyncio.wait_for: before 3.12 wait_for returns
the inner result when the inner future also completed during the stall, so the burst
never materialises and the test cannot exercise the duration gate.
"""
import time as time_mod
from redis.exceptions import TimeoutError as RedisTimeoutError
from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker
breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0)
async def healthy_redis_call_with_client_timeout():
return await asyncio.wait_for(asyncio.sleep(0.001, result="ok"), timeout=0.05)
deadline = time_mod.monotonic() + 0.05
await asyncio.sleep(0.001)
if time_mod.monotonic() > deadline:
raise RedisTimeoutError("read timed out")
return "ok"
async def stall_the_loop():
await asyncio.sleep(0)
@ -842,7 +853,7 @@ async def test_event_loop_stall_timeout_burst_keeps_breaker_closed():
stall_the_loop(),
return_exceptions=True,
)
timeouts = [r for r in results if isinstance(r, asyncio.TimeoutError)]
timeouts = [r for r in results if isinstance(r, RedisTimeoutError)]
assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst"
assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool"

29
tests/test_litellm/integrations/cloudzero/test_cloudzero.py
View file

@ -1,3 +1,4 @@
import hashlib
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
@ -74,6 +75,8 @@ class TestCloudZeroHourlyExport:
fake_db = MagicMock()
async def query_raw_mock(query: str, *params):
if "sha256(" in query:
return []
start_time_utc = params[0] if len(params) > 0 else None
end_time_utc = params[1] if len(params) > 1 else None
limit = params[2] if len(params) > 2 else None
@ -146,6 +149,9 @@ class TestCloudZeroHourlyExport:
return joined
fake_db.query_raw = AsyncMock(side_effect=query_raw_mock)
fake_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
fake_db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
fake_db.litellm_usertable.find_many = AsyncMock(return_value=[])
fake_client.db = fake_db
mock_prisma_client_getter.return_value = fake_client
@ -160,3 +166,26 @@ class TestCloudZeroHourlyExport:
logger = CloudZeroLogger(api_key="test", connection_id="test")
await logger._hourly_usage_data_export()
class TestLiteLLMDatabaseUsageData:
@pytest.mark.asyncio
async def test_builds_frame_from_rows_recovered_for_double_hashed_keys(self, monkeypatch: pytest.MonkeyPatch):
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
async def query_raw(query: str, *params):
if "sha256(" in query:
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
return [joined_row, dirty_row]
fake_client = MagicMock()
fake_client.db.query_raw = AsyncMock(side_effect=query_raw)
db = LiteLLMDatabase()
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: fake_client)
result = await db.get_usage_data()
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
assert result["team_id"].to_list() == ["team-0", "team-1"]

22
tests/test_litellm/integrations/focus/test_focus_database.py
View file

@ -1,5 +1,6 @@
"""Tests for FocusLiteLLMDatabase query construction."""
import hashlib
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
@ -87,3 +88,24 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
)
assert "ot.organization_alias as organization_alias" in query_text
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
@pytest.mark.asyncio
async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch):
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
async def query_raw(query: str, *params):
if "sha256(" in query:
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
return [joined_row, dirty_row]
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw)))
db = FocusLiteLLMDatabase()
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
result = await db.get_usage_data()
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
assert result["team_id"].to_list() == ["team-0", "team-1"]

65
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py
View file

@ -5,7 +5,10 @@ import pytest
import litellm
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
bedrock_guardrail_cost,
bedrock_guardrail_cost_by_unit,
billed_guardrail_cost_by_unit,
cost_breakdown_with_guardrail,
guardrail_cost_total,
guardrail_information_cost,
)
@ -56,6 +59,68 @@ def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch):
assert bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") == 0.0
def test_bedrock_guardrail_cost_by_unit_prices_every_counter_it_was_given(synthetic_cost_map):
"""LIT-5652: the daily rollup stores one row per counter, so pricing must come
back at that grain, keyed exactly like the usage. An explicit 0.0 in the cost
map is free; a counter the map does not list is unknown (None), never free,
while the scalar the spend path bills still sums only the known prices."""
usage = {"contentPolicyUnits": 2, "topicPolicyUnits": 1, "wordPolicyUnits": 5, "someFutureCounter": 3}
by_unit = bedrock_guardrail_cost_by_unit(usage_units=usage, aws_region_name="us-east-1")
assert by_unit is not None
assert by_unit.keys() == usage.keys()
assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003)
assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015)
assert by_unit["wordPolicyUnits"] == 0.0
assert by_unit["someFutureCounter"] is None
assert guardrail_cost_total(by_unit) == pytest.approx(0.00045)
assert guardrail_cost_total(by_unit) == pytest.approx(
bedrock_guardrail_cost(usage_units=usage, aws_region_name="us-east-1")
)
def test_bedrock_guardrail_cost_by_unit_is_none_without_pricing_so_unpriced_is_not_free(monkeypatch):
"""The scalar keeps returning 0.0 for the spend path; the per-unit view must
say "unknown" instead so the rollup stores NULL rather than a $0 that would
hide the exact silent-spend problem this feature exists to surface."""
monkeypatch.setattr(litellm, "model_cost", {})
assert bedrock_guardrail_cost_by_unit(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") is None
def test_billed_guardrail_cost_by_unit_reads_the_hook_stamp():
entry = {
"guardrail_name": "bedrock",
"guardrail_cost_by_unit": {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0, "someFutureCounter": None},
}
assert billed_guardrail_cost_by_unit(entry) == {
"contentPolicyUnits": 0.15,
"wordPolicyUnits": 0.0,
"someFutureCounter": None,
}
@pytest.mark.parametrize(
"entry",
[
{"guardrail_name": "no-pricing", "guardrail_usage": {"contentPolicyUnits": 1}},
{"guardrail_cost_by_unit": {"text_records": 0.5}, "guardrail_cost_in_spend": False},
{"guardrail_cost_by_unit": {"contentPolicyUnits": -0.5}},
{"guardrail_cost_by_unit": {"contentPolicyUnits": float("nan")}},
{"guardrail_cost_by_unit": {"contentPolicyUnits": float("inf")}},
{"guardrail_cost_by_unit": {"contentPolicyUnits": "bad"}},
{"guardrail_cost_by_unit": "not-a-map"},
{"guardrail_cost_by_unit": {"contentPolicyUnits": 0.1}, "guardrail_cost_in_spend": "maybe"},
"not-an-entry",
],
)
def test_billed_guardrail_cost_by_unit_is_none_when_unpriced_report_only_or_forged(entry):
assert billed_guardrail_cost_by_unit(entry) is None
def test_billed_guardrail_cost_by_unit_treats_none_in_spend_as_billed():
entry = {"guardrail_cost_by_unit": {"contentPolicyUnits": 0.15}, "guardrail_cost_in_spend": None}
assert billed_guardrail_cost_by_unit(entry) == {"contentPolicyUnits": 0.15}
def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(monkeypatch):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")

22
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -4630,6 +4630,28 @@ def test_generic_cost_per_token_grok_46_long_context(_local_model_cost_map):
assert completion_cost == pytest.approx(1_000 * 1.2e-05)
@pytest.mark.parametrize(
("model", "provider", "image_token_rate"),
[
("gpt-realtime-2.1", "openai", 5e-06),
("gpt-realtime-2.1-mini", "openai", 8e-07),
("azure/gpt-realtime-2.1", "azure", 5e-06),
("azure/gpt-realtime-2.1-mini", "azure", 8e-07),
],
)
def test_realtime_image_tokens_priced_per_token(model, provider, image_token_rate, _local_model_cost_map):
"""Realtime image input is billed per 1M image tokens, not per image."""
usage = Usage(
prompt_tokens=1_100,
completion_tokens=0,
total_tokens=1_100,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100, image_tokens=1_000),
)
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider)
text_rate = litellm.model_cost[model]["input_cost_per_token"]
assert prompt_cost == pytest.approx(100 * text_rate + 1_000 * image_token_rate)
@pytest.mark.parametrize(
("response_quality", "requested_quality", "expected_cost"),
[

52
tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py
View file

@ -9,7 +9,6 @@ import os
import pytest
from litellm.litellm_core_utils.fallback_generalizations import (
get_fallback_generalization_rules,
match_capability_generalizations,
@ -248,6 +247,57 @@ def test_azure_ai_claude_1m_context_entries(cost_map: dict):
assert cost_map[model]["max_input_tokens"] == 200000, model
# OpenRouter headline rates from GET https://openrouter.ai/api/v1/models.
# These were the catalog values that disagreed with that API (and, for the
# two spotlight models, the public model pages that their source fields cite).
_OPENROUTER_LIVE_COSTS = {
"openrouter/qwen/qwen3.5-plus-02-15": (2.6e-07, 1.56e-06, None),
"openrouter/openai/gpt-oss-120b": (3.7e-08, 1.7e-07, None),
"openrouter/qwen/qwen3-coder-plus": (6.5e-07, 3.25e-06, None),
"openrouter/qwen/qwen3.5-flash-02-23": (6.5e-08, 2.6e-07, None),
"openrouter/qwen/qwen3.5-27b": (1.95e-07, 1.56e-06, None),
"openrouter/gryphe/mythomax-l2-13b": (6e-08, 6e-08, None),
"openrouter/mancer/weaver": (4e-07, 7.5e-07, None),
"openrouter/xiaomi/mimo-v2.5-pro": (4.35e-07, 8.7e-07, 3.6e-09),
"openrouter/moonshotai/kimi-k2.5": (4.5e-07, 2.25e-06, 7e-08),
"openrouter/z-ai/glm-5": (6e-07, 1.92e-06, None),
}
_OPENROUTER_STALE_COSTS = {
"openrouter/qwen/qwen3.5-plus-02-15": (4e-07, 2.4e-06),
"openrouter/openai/gpt-oss-120b": (1.8e-07, 8e-07),
"openrouter/gryphe/mythomax-l2-13b": (1.875e-06, 1.875e-06),
}
@pytest.mark.parametrize(
"cost_map",
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
ids=["root", "bundled_backup"],
)
def test_openrouter_catalog_costs_match_live_headline_rates(cost_map: dict):
"""openrouter/* spend tracking reads these catalog fields. The values must
stay aligned with OpenRouter's published headline rate, not the stale
figures that over/under-counted by up to 30x. Both maps are checked so
the root file and bundled backup cannot drift apart."""
control = cost_map["openrouter/anthropic/claude-opus-5"]
assert control["input_cost_per_token"] == 5e-06
assert control["output_cost_per_token"] == 2.5e-05
assert control["cache_read_input_token_cost"] == 5e-07
for model, (inp, out, cache) in _OPENROUTER_LIVE_COSTS.items():
entry = cost_map[model]
assert entry["input_cost_per_token"] == inp, model
assert entry["output_cost_per_token"] == out, model
if cache is not None:
assert entry["cache_read_input_token_cost"] == cache, model
for model, (stale_in, stale_out) in _OPENROUTER_STALE_COSTS.items():
entry = cost_map[model]
assert entry["input_cost_per_token"] != stale_in, model
assert entry["output_cost_per_token"] != stale_out, model
def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
"""The load time feeds each pod's reload-due decision; a load that does not stamp it
would make manual reload requests race the proxy's startup"""

108
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py
View file

@ -5,6 +5,7 @@ from datetime import datetime
import pytest
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
@ -32,7 +33,16 @@ class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator):
self.logging_call_count += 1
def _make_logging_obj(test_name: str) -> LiteLLMLoggingObj:
class _FailureRecorder(CustomLogger):
def __init__(self):
super().__init__()
self.failure_kwargs: list = []
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.failure_kwargs.append(kwargs)
def _make_logging_obj(test_name: str, failure_recorder: _FailureRecorder | None = None) -> LiteLLMLoggingObj:
return LiteLLMLoggingObj(
model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": "hi"}],
@ -41,9 +51,19 @@ def _make_logging_obj(test_name: str) -> LiteLLMLoggingObj:
start_time=datetime.now(),
litellm_call_id=test_name,
function_id=test_name,
dynamic_async_failure_callbacks=[failure_recorder] if failure_recorder is not None else None,
)
async def _wait_for_failure_event(recorder: _FailureRecorder) -> dict:
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert len(recorder.failure_kwargs) == 1, "expected exactly one failure event"
return recorder.failure_kwargs[0]
def _make_iterator(test_name: str) -> BaseAnthropicMessagesStreamingIterator:
return BaseAnthropicMessagesStreamingIterator(
litellm_logging_obj=_make_logging_obj(test_name),
@ -524,6 +544,25 @@ async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconn
await asyncio.wait_for(deferred_fired.wait(), timeout=5)
class _DetachedFailureRecorder:
"""Stands in for the closure the proxy arms so a detached-stream failure still reaches its failure hook."""
def __init__(self):
self.exceptions = []
async def __call__(self, exc: Exception) -> None:
self.exceptions.append(exc)
async def _wait_for_detached_failure(recorder: _DetachedFailureRecorder) -> Exception:
for _ in range(200):
if recorder.exceptions:
await asyncio.sleep(0.02)
return recorder.exceptions[0]
await asyncio.sleep(0.01)
raise AssertionError("the detached failure hook never fired")
class _ProviderStreamError(Exception):
"""Stand-in for a provider-specific streaming failure carrying a status code."""
@ -539,7 +578,8 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
before message_stop must propagate the ORIGINAL provider exception to a
still-connected client, so the proxy's failure handling keeps the
provider-specific status. The pump must not swallow it into a generic
api_error event + normal termination.
api_error event + normal termination, and the request is logged as a
failure carrying the partial usage, never as a success.
"""
async def _failing_stream():
@ -547,10 +587,13 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
raise _ProviderStreamError("bedrock stream blew up", status_code=529)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"),
litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
received = []
@ -561,18 +604,25 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
with pytest.raises(_ProviderStreamError) as excinfo:
await _drain()
failure_kwargs = await _wait_for_failure_event(recorder)
assert excinfo.value.status_code == 529
assert received
assert not any(c.startswith(b"event: error\n") for c in received)
assert iterator.logged_chunks == []
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
await asyncio.sleep(0.05)
assert detached_hook.exceptions == [], "the relay re-raised the error, so the proxy failure hook already ran"
@pytest.mark.asyncio
async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect():
async def test_async_sse_wrapper_logs_failure_on_upstream_error_after_disconnect():
"""
When the upstream errors AFTER the client has already disconnected there is
no live client to re-raise to and no failure hook will run, so the pump
salvages partial spend from what it collected instead of dropping the row.
no live client to re-raise to and no proxy failure hook will run, so the
pump logs the failure itself with the partial usage it collected; it must
never bill the broken stream as a success.
"""
tail_gated = asyncio.Event()
@ -582,34 +632,38 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_
await tail_gated.wait()
raise _ProviderStreamError("late failure", status_code=500)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"),
litellm_logging_obj=_make_logging_obj("test_failure_logged_on_late_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
gen = iterator.async_sse_wrapper(_gated_failing_stream())
received = [await gen.__anext__(), await gen.__anext__()]
await gen.aclose() # client disconnects before the upstream error
tail_gated.set() # let the upstream raise now, after disconnect
for _ in range(100):
if iterator.logged_chunks:
break
await asyncio.sleep(0.01)
failure_kwargs = await _wait_for_failure_event(recorder)
assert len(received) == 2
assert iterator.logged_chunks == received
assert iterator.logging_call_count == 0
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
assert isinstance(failure_kwargs["exception"], _ProviderStreamError)
assert await _wait_for_detached_failure(detached_hook) is failure_kwargs["exception"]
assert len(detached_hook.exceptions) == 1
@pytest.mark.asyncio
async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consumed():
async def test_async_sse_wrapper_logs_failure_when_queued_error_is_never_consumed():
"""
When the upstream errors while the client is still connected, the pump
forwards the exception through the queue expecting the relay to re-raise it
into the proxy's failure handling. If the client disconnects before
consuming that queued exception, the handoff never happens and no failure
hook runs, so the pump must notice the unconsumed exception at teardown and
salvage partial spend instead of dropping the row entirely.
forwards the exception through the queue for the relay to re-raise. If the
client disconnects before consuming that queued exception, no proxy failure
hook runs, so the failure logged by the pump itself is the only record of
the request; it must be a failure row, not a salvaged success.
"""
upstream_errored = asyncio.Event()
@ -619,23 +673,27 @@ async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consu
upstream_errored.set()
raise _ProviderStreamError("mid-stream failure", status_code=500)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_salvage_on_unconsumed_queued_error"),
litellm_logging_obj=_make_logging_obj("test_failure_logged_on_unconsumed_queued_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
gen = iterator.async_sse_wrapper(_failing_stream())
received = [await gen.__anext__(), await gen.__anext__()]
await upstream_errored.wait() # exception is now queued behind the consumed chunks
await gen.aclose() # client disconnects without ever consuming the queued exception
for _ in range(100):
if iterator.logged_chunks:
break
await asyncio.sleep(0.01)
failure_kwargs = await _wait_for_failure_event(recorder)
assert iterator.logging_call_count == 1
assert iterator.logged_chunks == received
assert len(received) == 2
assert iterator.logging_call_count == 0
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
assert await _wait_for_detached_failure(detached_hook) is failure_kwargs["exception"]
assert len(detached_hook.exceptions) == 1
@pytest.mark.asyncio

29
tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py
View file

@ -499,3 +499,32 @@ class TestAzureAIServiceTierCostCalculation:
assert flex_prompt < standard_prompt
assert flex_completion < standard_completion
def test_codestral_2501_model_info_and_cost(local_model_cost_map):
model_info = get_model_info(model="Codestral-2501", custom_llm_provider="azure_ai")
usage = Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000)
prompt_cost, completion_cost = cost_per_token(model="Codestral-2501", usage=usage)
assert model_info["mode"] == "chat"
assert model_info["max_input_tokens"] == 256000
assert model_info["max_output_tokens"] == 4096
assert prompt_cost == pytest.approx(0.3)
assert completion_cost == pytest.approx(0.9)
def test_mai_thinking_1_model_info_and_cost(local_model_cost_map):
model_info = get_model_info(model="MAI-Thinking-1", custom_llm_provider="azure_ai")
usage = Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000)
prompt_cost, completion_cost = cost_per_token(model="MAI-Thinking-1", usage=usage)
assert model_info["mode"] == "chat"
assert model_info["max_input_tokens"] == 256000
assert model_info["max_output_tokens"] == 64000
assert model_info["cache_read_input_token_cost"] == pytest.approx(2e-07)
assert model_info["supports_reasoning"] is True
assert model_info["supports_function_calling"] is True
assert prompt_cost == pytest.approx(2.0)
assert completion_cost == pytest.approx(8.0)

25
tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py
View file

@ -176,6 +176,7 @@ def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected):
("FW-MiniMax-M2.5", 0.33, 1.32),
("FW-Inkling", 1.0, 4.05),
("FW-Nemotron-3-Ultra-NVFP4", 0.6, 2.4),
("FW-Nemotron-Lightning-3.5-30B-A3B", 0.06, 0.22),
],
)
def test_azure_ai_fw_cost_per_token(
@ -196,6 +197,30 @@ def test_azure_ai_fw_cost_per_token(
assert completion_cost == pytest.approx(expected_completion)
def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map):
model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B")
assert model_info["litellm_provider"] == "azure_ai"
assert model_info["mode"] == "chat"
assert model_info["input_cost_per_token"] == pytest.approx(6e-08)
assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07)
assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08)
assert model_info["max_input_tokens"] == 262144
assert model_info["supports_function_calling"] is True
assert model_info["supports_reasoning"] is True
assert model_info["supports_tool_choice"] is True
assert model_info["supports_prompt_caching"] is True
assert model_info["supports_vision"] is False
def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map):
from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig
supported_params = AzureAIStudioConfig().get_supported_openai_params("FW-Nemotron-Lightning-3.5-30B-A3B")
assert "tool_choice" in supported_params
def test_azure_ai_fw_kimi_k26_case_insensitive_lookup(use_local_model_cost_map):
upper = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Kimi-K2.6")
lower = use_local_model_cost_map.get_model_info(model="azure_ai/fw-kimi-k2.6")

21
tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
View file

@ -68,6 +68,27 @@ def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws):
assert base_url == vpc_endpoint
@pytest.mark.parametrize(
"lookalike_host",
[
"https://bedrock-mantle.us-east-1.api.aws.internal.example.com",
"https://bedrock-mantle.us-gov-west-1.api.aws-int.example.com",
"https://bedrock-mantle.us-east-1.api.aws:8443",
],
)
def test_lookalike_mantle_host_api_base_is_kept(no_ambient_aws, lookalike_host):
url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
api_base=lookalike_host,
api_key=None,
model="us.openai.gpt-5.6-sol",
endpoint=INVOKE_ENDPOINT,
request_query_params=None,
litellm_params={"api_base": lookalike_host},
)
assert str(url) == f"{lookalike_host}/{INVOKE_ENDPOINT}"
assert base_url == lookalike_host
def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws):
url, _ = BedrockMantlePassthroughConfig().get_complete_url(
api_base=None,

23
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
View file

@ -26,6 +26,12 @@ from litellm.llms.bedrock_mantle.responses.transformation import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
LOOKALIKE_MANTLE_HOSTS = (
"https://bedrock-mantle.us-east-1.api.aws.internal.example.com",
"https://bedrock-mantle.us-gov-west-1.api.aws-int.example.com",
"https://bedrock-mantle.us-east-1.api.aws:8443",
)
class TestBedrockMantleResponsesURL:
def test_url_uses_region_from_env(self, monkeypatch):
@ -1642,6 +1648,23 @@ class TestBedrockMantleResponsesSigV4:
)
assert url == "https://mantle-proxy.internal.example/openai/v1/responses"
@pytest.mark.parametrize("lookalike_host", LOOKALIKE_MANTLE_HOSTS)
def test_lookalike_mantle_host_from_api_base_is_preserved(self, monkeypatch, lookalike_host):
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
cfg = BedrockMantleResponsesAPIConfig()
url = cfg.get_complete_url(
api_base=f"{lookalike_host}/openai/v1",
litellm_params={"aws_region_name": "us-east-2"},
)
assert url == f"{lookalike_host}/openai/v1/responses"
@pytest.mark.parametrize("lookalike_host", LOOKALIKE_MANTLE_HOSTS)
def test_lookalike_mantle_host_from_env_is_preserved(self, monkeypatch, lookalike_host):
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", lookalike_host)
cfg = BedrockMantleResponsesAPIConfig()
url = cfg.get_complete_url(api_base=None, litellm_params={})
assert url == f"{lookalike_host}/openai/v1/responses"
def test_caller_authorization_does_not_override_sigv4(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied Authorization header (e.g.
from extra_headers, surviving the relaxed validate_environment) must not clobber

40
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
View file

@ -399,6 +399,46 @@ class TestBedrockMantleChatAuth:
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
assert "/us-west-2/bedrock/aws4_request" not in headers["Authorization"]
@pytest.mark.parametrize(
("region_params", "env", "expected_region"),
[
({"aws_region_name": "us-west-2"}, {}, "us-west-2"),
({}, {"BEDROCK_MANTLE_REGION": "ap-southeast-2"}, "ap-southeast-2"),
],
)
def test_sigv4_scope_ignores_the_region_segment_of_a_lookalike_host(
self, monkeypatch, region_params, env, expected_region
):
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
for var in (
"BEDROCK_MANTLE_API_KEY",
"AWS_BEARER_TOKEN_BEDROCK",
"BEDROCK_MANTLE_REGION",
"BEDROCK_MANTLE_API_BASE",
"AWS_REGION",
"AWS_REGION_NAME",
):
monkeypatch.delenv(var, raising=False)
for var, value in env.items():
monkeypatch.setenv(var, value)
cfg = BedrockMantleChatConfig(aws_signer=BaseAWSLLM())
headers, _ = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
**region_params,
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.eu-west-1.api.aws.internal.example.com/openai/v1/chat/completions",
api_key=None,
)
assert f"/{expected_region}/bedrock/aws4_request" in headers["Authorization"]
assert "/eu-west-1/bedrock/aws4_request" not in headers["Authorization"]
def test_no_bearer_and_no_credentials_raises_value_error(self, monkeypatch):
from unittest.mock import MagicMock

24
tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py
View file

@ -18,6 +18,10 @@ NEW_MODELS: Final = (
"databricks/databricks-claude-opus-5",
"databricks/databricks-claude-sonnet-5",
"databricks/databricks-claude-fable-5",
"databricks/databricks-claude-fable-5-1",
"databricks/databricks-gpt-5-6-sol",
"databricks/databricks-gpt-5-6-terra",
"databricks/databricks-gpt-5-6-luna",
)
DOLLARS_PER_DBU: Final = Decimal("0.070")
@ -28,6 +32,7 @@ PRICE_FIELDS: Final = (
"cache_read_input_token_cost",
)
PUBLISHED_DBU_PER_MILLION: Final = {
"databricks/databricks-claude-fable-5-1": ("142.858", "714.286", "178.572", "3.572"),
"databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"),
"databricks/databricks-claude-opus-5": ("71.429", "357.143", "89.286", "7.143"),
"databricks/databricks-claude-opus-4-8": ("71.429", "357.143", "89.286", "7.143"),
@ -52,9 +57,17 @@ PUBLISHED_DBU_PER_MILLION: Final = {
"databricks/databricks-gpt-5-2": ("25.000", "200.000", "25.000", "2.500"),
"databricks/databricks-gpt-5-2-codex": ("25.000", "200.000", "25.000", "2.500"),
"databricks/databricks-gpt-5-3-codex": ("25.000", "200.000", "25.000", "2.500"),
"databricks/databricks-gpt-5-6-sol": ("57.143", "285.714", "71.429", "5.714"),
"databricks/databricks-gpt-5-6-terra": ("35.714", "214.286", "44.643", "3.571"),
"databricks/databricks-gpt-5-6-luna": ("14.286", "85.714", "17.857", "1.429"),
"databricks/databricks-gpt-5-5": ("71.429", "428.571", "71.429", "7.143"),
"databricks/databricks-gpt-5-5-pro": ("428.571", "2571.429", "428.571", "428.571"),
"databricks/databricks-gpt-5-4": ("35.714", "214.286", "35.714", "3.571"),
"databricks/databricks-gpt-5-4-mini": ("10.714", "64.286", "10.714", "1.071"),
"databricks/databricks-gpt-5-4-nano": ("2.857", "17.857", "2.857", "0.286"),
"databricks/databricks-gemini-3-6-flash": ("26.786", "133.929", "26.786", "2.679"),
"databricks/databricks-gemini-3-5-flash": ("26.786", "160.714", "26.786", "2.679"),
"databricks/databricks-gemini-3-5-flash-lite": ("5.357", "44.643", "5.357", "0.536"),
"databricks/databricks-gemini-3-1-pro": ("35.714", "214.286", "35.714", "3.571"),
"databricks/databricks-gemini-3-pro": ("35.714", "214.286", "35.714", "3.571"),
"databricks/databricks-gemini-3-flash": ("8.929", "53.571", "8.929", "0.893"),
@ -65,6 +78,13 @@ PUBLISHED_DBU_PER_MILLION: Final = {
"databricks/databricks-deepseek-v4-flash-0731": ("2.000", "4.000", "2.000", "0.400"),
"databricks/databricks-deepseek-v4-pro-0813": ("18.857", "56.571", "18.857", "1.886"),
"databricks/databricks-glm-5-2": ("20.000", "62.857", "20.000", "3.714"),
"databricks/databricks-glm-5-3": ("20.000", "62.857", "20.000", "3.714"),
"databricks/databricks-glm-5-3-flash": ("2.143", "7.143", "2.143", "0.429"),
"databricks/databricks-inkling": ("14.286", "57.857", "14.286", "2.429"),
"databricks/databricks-grok-4-6": ("35.714", "107.143", "35.714", "8.929"),
"databricks/databricks-qwen35-122b-a10b": ("3.143", "31.429", "3.143", "3.143"),
"databricks/databricks-qwen3-next-80b-a3b-instruct": ("2.143", "17.143", "2.143", "2.143"),
"databricks/databricks-qwen3-embedding-0-6b": ("0.286", "0", "0.286", "0.286"),
}
PROMOTIONAL_DISCOUNT: Final = 0.80
PROMOTION_EXPIRES: Final = "2027-01-31"
@ -73,6 +93,10 @@ ENTRIES_STORING_PROMOTIONAL_RATE: Final = (
"databricks/databricks-gemini-2-5-flash",
)
ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION: Final = (
"databricks/databricks-gemini-3-6-flash",
"databricks/databricks-gemini-3-5-flash",
"databricks/databricks-gemini-3-5-flash-lite",
"databricks/databricks-grok-4-6",
"databricks/databricks-gemini-3-1-pro",
"databricks/databricks-gemini-3-pro",
"databricks/databricks-gemini-3-flash",

96
tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py
View file

@ -1,10 +1,13 @@
import math
from datetime import datetime, timezone
import pytest
import litellm
from litellm.llms.fireworks_ai.cost_calculator import cost_per_token
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
from litellm.types.utils import OffPeakPricing, PromptTokensDetailsWrapper, Usage
MODEL = "accounts/fireworks/models/glm-5p2"
INPUT_COST = 1.4e-06
@ -64,3 +67,94 @@ def test_no_cached_tokens_matches_full_input_rate():
assert prompt_cost == pytest.approx(prompt_tokens * INPUT_COST)
assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST)
OFF_PEAK_MODEL = "accounts/fireworks/models/off-peak-test"
OFF_PEAK_WINDOW = "14:00-00:00"
INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)
OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc)
STANDARD_INPUT_COST = 1.5e-07
STANDARD_OUTPUT_COST = 6e-07
STANDARD_CACHE_READ_COST = 1.5e-08
def _register_off_peak_model(off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST) -> None:
litellm.model_cost[f"fireworks_ai/{OFF_PEAK_MODEL}"] = {
"litellm_provider": "fireworks_ai",
"mode": "chat",
"input_cost_per_token": STANDARD_INPUT_COST,
"output_cost_per_token": STANDARD_OUTPUT_COST,
"off_peak_pricing": off_peak_pricing,
**({} if cache_read_cost is None else {"cache_read_input_token_cost": cache_read_cost}),
}
def test_off_peak_window_swaps_in_the_off_peak_rates():
"""
Regression (LIT-6874): a deployment configured with off_peak_pricing kept billing the
standard fireworks_ai rates inside its window, while the same block on a deepseek
deployment billed the off-peak rates.
"""
_register_off_peak_model(
{
"hours_utc": OFF_PEAK_WINDOW,
"input_cost_per_token": 1e-08,
"output_cost_per_token": 2e-08,
"cache_read_input_token_cost": 1e-09,
}
)
usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200)
prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=INSIDE_WINDOW)
assert math.isclose(prompt_cost, (700 * 1e-08) + (300 * 1e-09), rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)
peak_prompt_cost, peak_completion_cost = cost_per_token(
model=OFF_PEAK_MODEL, usage=usage, current_time=OUTSIDE_WINDOW
)
assert math.isclose(peak_prompt_cost, (700 * STANDARD_INPUT_COST) + (300 * STANDARD_CACHE_READ_COST), rel_tol=1e-10)
assert math.isclose(peak_completion_cost, 200 * STANDARD_OUTPUT_COST, rel_tol=1e-10)
def test_off_peak_rates_left_unset_keep_the_standard_rates():
"""A block that only overrides the input rate leaves output and cache reads on the standard rates."""
_register_off_peak_model({"hours_utc": OFF_PEAK_WINDOW, "input_cost_per_token": 1e-08})
usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200)
prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=INSIDE_WINDOW)
assert math.isclose(prompt_cost, (700 * 1e-08) + (300 * STANDARD_CACHE_READ_COST), rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * STANDARD_OUTPUT_COST, rel_tol=1e-10)
def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_a_cache_read_rate():
"""Most fireworks_ai price-map entries carry no cache_read_input_token_cost, so cached tokens
fall back to the input rate, and inside the window that has to be the off-peak one."""
_register_off_peak_model(
{"hours_utc": OFF_PEAK_WINDOW, "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08},
cache_read_cost=None,
)
usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200)
prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=INSIDE_WINDOW)
assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)
peak_prompt_cost, _ = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=OUTSIDE_WINDOW)
assert math.isclose(peak_prompt_cost, 1000 * STANDARD_INPUT_COST, rel_tol=1e-10)
def test_off_peak_defaults_to_the_current_time():
"""The proxy's cost dispatch passes no clock, so an all-day window has to apply on the
default current time."""
_register_off_peak_model({"hours_utc": "00:00-00:00", "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08})
usage = _usage(prompt_tokens=1000, cached_tokens=0, completion_tokens=200)
prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage)
assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)

14
tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py
View file

@ -15,6 +15,7 @@ from litellm.cost_calculator import completion_cost
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
OCR4_COST_PER_PAGE = 0.004
OCR4_ANNOTATION_COST_PER_PAGE = 0.005
REPO_ROOT = Path(__file__).parents[5]
MAIN_COST_MAP = REPO_ROOT / "model_prices_and_context_window.json"
@ -133,3 +134,16 @@ def test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate(local_model_cost_ma
call_type="ocr",
)
assert cost == pytest.approx(AZURE_DOC_AI_COST_PER_PAGE)
def test_azure_ocr4_bills_ocr_and_annotation_pages_at_their_own_rates(local_model_cost_map) -> None:
info = litellm.get_model_info(model="azure_ai/mistral-ocr-4-0", custom_llm_provider="azure_ai")
assert info["ocr_cost_per_page"] == OCR4_COST_PER_PAGE
assert info["annotation_cost_per_page"] == OCR4_ANNOTATION_COST_PER_PAGE
cost = completion_cost(
completion_response=_annotated_ocr_response("mistral-ocr-4-0", 2, 3),
model="azure_ai/mistral-ocr-4-0",
custom_llm_provider="azure_ai",
call_type="ocr",
)
assert cost == pytest.approx(2 * OCR4_COST_PER_PAGE + 3 * OCR4_ANNOTATION_COST_PER_PAGE)

93
tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
View file

@ -8,6 +8,7 @@ search queries, and reasoning tokens.
import json
import math
import os
from datetime import datetime, timezone
from unittest.mock import patch
import pytest
@ -21,6 +22,7 @@ from litellm.llms.perplexity.cost_calculator import (
)
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
OffPeakPricing,
Usage,
PromptTokensDetailsWrapper,
)
@ -523,3 +525,94 @@ class TestPerplexityCostCalculator:
)
assert math.isclose(total_cost, 1000 * 1.4e-06 + 500 * 4.4e-06, rel_tol=1e-9)
OFF_PEAK_MODEL = "sonar-off-peak-test"
OFF_PEAK_WINDOW = "14:00-00:00"
INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)
OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc)
def _register_off_peak_model(self, off_peak_pricing: OffPeakPricing) -> None:
litellm.model_cost[f"perplexity/{self.OFF_PEAK_MODEL}"] = {
"litellm_provider": "perplexity",
"mode": "chat",
"input_cost_per_token": 1e-06,
"output_cost_per_token": 1e-06,
"output_cost_per_reasoning_token": 3e-06,
"citation_cost_per_token": 2e-06,
"search_context_cost_per_query": {"search_context_size_low": 0.005},
"off_peak_pricing": off_peak_pricing,
}
def test_off_peak_window_swaps_in_the_off_peak_rates(self):
"""
Regression (LIT-6874): a deployment configured with off_peak_pricing kept billing the
standard perplexity rates inside its window, while the same block on a deepseek
deployment billed the off-peak rates.
"""
self._register_off_peak_model(
{"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1e-07, "output_cost_per_token": 2e-07}
)
usage = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200)
prompt_cost, completion_cost = perplexity_cost_per_token(
model=self.OFF_PEAK_MODEL, usage=usage, current_time=self.INSIDE_WINDOW
)
assert math.isclose(prompt_cost, 1000 * 1e-07, rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-07, rel_tol=1e-10)
peak_prompt_cost, peak_completion_cost = perplexity_cost_per_token(
model=self.OFF_PEAK_MODEL, usage=usage, current_time=self.OUTSIDE_WINDOW
)
assert math.isclose(peak_prompt_cost, 1000 * 1e-06, rel_tol=1e-10)
assert math.isclose(peak_completion_cost, 200 * 1e-06, rel_tol=1e-10)
def test_off_peak_rates_leave_citation_search_and_reasoning_fees_alone(self):
"""Inside the window only the plain input and output rates change: citation tokens, the
per-request search fee, and a dedicated reasoning rate keep billing as published."""
self._register_off_peak_model(
{"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1e-07, "output_cost_per_token": 2e-07}
)
usage = Usage(
prompt_tokens=1000,
completion_tokens=200,
total_tokens=1200,
prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=1),
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50),
)
usage.citation_tokens = 100
prompt_cost, completion_cost = perplexity_cost_per_token(
model=self.OFF_PEAK_MODEL, usage=usage, current_time=self.INSIDE_WINDOW
)
assert math.isclose(prompt_cost, (1000 * 1e-07) + (100 * 2e-06), rel_tol=1e-10)
assert math.isclose(completion_cost, (150 * 2e-07) + (50 * 3e-06) + 0.005, rel_tol=1e-10)
def test_off_peak_defaults_to_the_current_time(self):
"""The proxy's cost dispatch passes no clock, so an all-day window has to apply on the
default current time."""
self._register_off_peak_model(
{"hours_utc": "00:00-00:00", "input_cost_per_token": 1e-07, "output_cost_per_token": 2e-07}
)
usage = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200)
prompt_cost, completion_cost = perplexity_cost_per_token(model=self.OFF_PEAK_MODEL, usage=usage)
assert math.isclose(prompt_cost, 1000 * 1e-07, rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-07, rel_tol=1e-10)
def test_provider_stated_cost_still_wins_inside_an_off_peak_window(self):
"""A response that carries Perplexity's own metered cost bills that cost whatever the
window says; the caller strips it when the deployment carries custom pricing."""
self._register_off_peak_model(
{"hours_utc": "00:00-00:00", "input_cost_per_token": 1e-07, "output_cost_per_token": 2e-07}
)
usage = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200)
usage.cost = {"total_cost": 0.00501}
prompt_cost, completion_cost = perplexity_cost_per_token(model=self.OFF_PEAK_MODEL, usage=usage)
assert prompt_cost == 0.0
assert completion_cost == 0.00501

64
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py
View file

@ -16,3 +16,67 @@ def test_lookup_mcp_server_auth_in_headers_sanitized_alias():
headers = {"github_mcp": {"Authorization": "Bearer token"}}
result = lookup_mcp_server_auth_in_headers(headers, alias="GitHub-MCP")
assert result == {"Authorization": "Bearer token"}
def test_lookup_mcp_server_auth_in_headers_group_header_is_default_for_members():
headers = {"shared": {"Authorization": "Bearer group-token"}}
assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", server_name="alpha", access_groups=["shared"]) == {
"Authorization": "Bearer group-token"
}
assert lookup_mcp_server_auth_in_headers(headers, alias="beta", server_name="beta", access_groups=["Shared"]) == {
"Authorization": "Bearer group-token"
}
def test_lookup_mcp_server_auth_in_headers_group_header_sanitized_group_name():
headers = {"dev_group": {"Authorization": "Bearer group-token"}}
assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", access_groups=["Dev Group"]) == {
"Authorization": "Bearer group-token"
}
def test_lookup_mcp_server_auth_in_headers_server_header_overrides_group_header():
headers = {
"shared": {"Authorization": "Bearer group-token"},
"beta": {"Authorization": "Bearer beta-token"},
}
assert lookup_mcp_server_auth_in_headers(headers, alias="beta", server_name="beta", access_groups=["shared"]) == {
"Authorization": "Bearer beta-token"
}
def test_lookup_mcp_server_auth_in_headers_group_header_not_forwarded_outside_group():
headers = {"shared": {"Authorization": "Bearer group-token"}}
assert (
lookup_mcp_server_auth_in_headers(headers, alias="gamma", server_name="gamma", access_groups=["other"]) is None
)
assert lookup_mcp_server_auth_in_headers(headers, alias="gamma", server_name="gamma", access_groups=None) is None
def test_lookup_mcp_server_auth_in_headers_alias_colliding_with_group_name_keeps_server_level_match():
headers = {"shared": {"Authorization": "Bearer shared-token"}}
assert lookup_mcp_server_auth_in_headers(headers, alias="shared", access_groups=["other"]) == {
"Authorization": "Bearer shared-token"
}
assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", access_groups=["shared"]) == {
"Authorization": "Bearer shared-token"
}
assert lookup_mcp_server_auth_in_headers(headers, alias="gamma", access_groups=["other"]) is None
def test_lookup_mcp_server_auth_in_headers_conflicting_group_headers_fail_closed():
headers = {
"shared": {"Authorization": "Bearer group-token"},
"other": {"Authorization": "Bearer other-token"},
}
assert lookup_mcp_server_auth_in_headers(headers, alias="delta", access_groups=["shared", "other"]) is None
def test_lookup_mcp_server_auth_in_headers_identical_group_headers_resolve():
headers = {
"shared": {"Authorization": "Bearer group-token"},
"other": {"Authorization": "Bearer group-token"},
}
assert lookup_mcp_server_auth_in_headers(headers, alias="delta", access_groups=["shared", "other"]) == {
"Authorization": "Bearer group-token"
}

37
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -303,6 +303,43 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
assert extra_headers == {"Authorization": "Bearer token"}
def test_prepare_mcp_server_headers_group_header_defaults_for_members_only():
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
def server(alias: str, group: str) -> MCPServer:
return MCPServer(
server_id=f"server-{alias}",
name=alias,
alias=alias,
transport=MCPTransport.http,
access_groups=[group],
)
mcp_server_auth_headers = {
"shared": {"Authorization": "Bearer group-token"},
"beta": {"Authorization": "Bearer beta-token"},
}
def resolve(mcp_server: MCPServer):
server_auth_header, _ = _prepare_mcp_server_headers(
server=mcp_server,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers={"x-litellm-api-key": "Bearer sk-litellm-key"},
)
return server_auth_header
assert resolve(server("alpha", "shared")) == {"Authorization": "Bearer group-token"}
assert resolve(server("beta", "shared")) == {"Authorization": "Bearer beta-token"}
assert resolve(server("gamma", "other")) is None
def test_prepare_mcp_server_headers_passthrough_strips_authorization_without_admission_header():
try:
from litellm.proxy._experimental.mcp_server.server import (

37
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -25,7 +25,8 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.mcp import MCPAuth
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def _rendered_log_message(call):
@ -3268,6 +3269,40 @@ class TestConnectionErrorMessage:
assert "proxy logs" in message.lower()
class TestGetServerAuthHeaderGroupDefault:
"""``x-mcp-<access_group>-authorization`` is the default for group members, the per-server
header still wins, and servers outside the group never see the group credential."""
@staticmethod
def _server(alias: str, group: str) -> MCPServer:
return MCPServer(
server_id=f"server-{alias}",
name=alias,
server_name=alias,
alias=alias,
url="https://example.com/mcp",
transport=MCPTransport.http,
access_groups=[group],
)
def test_group_header_applies_to_members_and_per_server_header_overrides(self):
headers = {
"shared": {"Authorization": "Bearer group-token"},
"beta": {"Authorization": "Bearer beta-token"},
}
assert rest_endpoints._get_server_auth_header(self._server("alpha", "shared"), headers, None) == {
"Authorization": "Bearer group-token"
}
assert rest_endpoints._get_server_auth_header(self._server("beta", "shared"), headers, None) == {
"Authorization": "Bearer beta-token"
}
def test_group_header_falls_back_to_legacy_header_outside_group(self):
headers = {"shared": {"Authorization": "Bearer group-token"}}
assert rest_endpoints._get_server_auth_header(self._server("gamma", "other"), headers, None) is None
assert rest_endpoints._get_server_auth_header(self._server("gamma", "other"), headers, "legacy") == "legacy"
class TestToolResponseMcpInfoEnrichment:
"""The REST tools/list response must expose the user-facing alias and the
server_id alongside the internal server_name so clients (agent builder UIs)

35
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -3312,6 +3312,41 @@ def test_user_daily_activity_routes_reachable_by_non_admin(route, user_role):
)
@pytest.mark.parametrize(
"user_role",
[
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
],
)
def test_team_spend_by_user_reachable_by_non_admin(user_role):
user_obj = LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=user_role,
)
valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role)
request = MagicMock(spec=Request)
request.query_params = {}
def outcome(route: str) -> str:
try:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=user_role,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
except Exception as exc:
return f"denied: {exc}"
return "allowed"
assert outcome("/team/spend/by_user") == "allowed"
assert outcome("/team/spend/by_key").startswith("denied: Only proxy admin")
def test_user_daily_activity_aggregated_not_covered_by_prefix_match():
"""check_route_access is exact-match plus explicit wildcards, so listing the
parent /user/daily/activity does not implicitly cover the /aggregated

44
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -2961,9 +2961,7 @@ async def test_streaming_hook_reraises_guardrail_service_failures():
guardrail = _sse_guardrail()
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.side_effect = HTTPException(
status_code=500, detail="Bedrock guardrail throttle retries exhausted"
)
mock_api.side_effect = HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted")
with pytest.raises(HTTPException) as exc:
await _drain_streaming_hook(guardrail)
@ -5097,13 +5095,46 @@ def test_build_tracing_detail_surfaces_usage_counters_and_cost(monkeypatch):
detail = guardrail._build_tracing_detail(
{
"action": "GUARDRAIL_INTERVENED",
"usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0, "oddball": "not-an-int"},
"usage": {
"topicPolicyUnits": 1,
"contentPolicyUnits": 2,
"wordPolicyUnits": 0,
"someFutureCounter": 3,
"oddball": "not-an-int",
},
},
aws_region_name="us-east-1",
)
assert detail["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0}
assert detail["guardrail_usage"] == {
"topicPolicyUnits": 1,
"contentPolicyUnits": 2,
"wordPolicyUnits": 0,
"someFutureCounter": 3,
}
assert detail["guardrail_cost"] == pytest.approx(0.00045)
by_unit = detail["guardrail_cost_by_unit"]
assert by_unit is not None and by_unit.keys() == detail["guardrail_usage"].keys()
assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015)
assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003)
assert by_unit["wordPolicyUnits"] == 0.0
assert by_unit["someFutureCounter"] is None
assert by_unit["wordPolicyUnits"] == 0.0
def test_build_tracing_detail_omits_cost_by_unit_when_unpriced_but_keeps_scalar_zero(monkeypatch):
"""LIT-5652: without a cost-map entry the spend path still bills 0.0, but the
per-counter stamp must be absent so the rollup records NULL, not $0."""
monkeypatch.setattr(litellm, "model_cost", {})
guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT")
detail = guardrail._build_tracing_detail(
{"action": "NONE", "usage": {"contentPolicyUnits": 5}}, aws_region_name="us-east-1"
)
assert detail["guardrail_usage"] == {"contentPolicyUnits": 5}
assert detail["guardrail_cost"] == 0.0
assert "guardrail_cost_by_unit" not in detail
def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none():
@ -5115,6 +5146,7 @@ def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none():
):
assert "guardrail_usage" not in detail
assert "guardrail_cost" not in detail
assert "guardrail_cost_by_unit" not in detail
@pytest.mark.asyncio
@ -5478,7 +5510,7 @@ async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan():
scan_index = events.index("scan")
chunk_events = [e for e in events if e != "scan"]
assert events.count("scan") == 1
assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index]
assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[:scan_index]
assert ("chunk", "Hello") in events[:scan_index]
assert ("chunk", " world") in events[:scan_index]
assert len(chunk_events) == 3

148
tests/test_litellm/proxy/guardrails/test_usage_endpoints.py
View file

@ -85,7 +85,10 @@ def _units_row(
api_key: str = "",
usage_unit: str = "contentPolicyUnits",
units: int = 1,
cost: float | None = None,
untracked_units: int = 0,
) -> Any:
"""cost=None is a row written before the cost column existed (untracked in full)."""
r = MagicMock()
r.guardrail_id = guardrail_id
r.date = date
@ -93,6 +96,8 @@ def _units_row(
r.api_key = api_key
r.usage_unit = usage_unit
r.units = units
r.cost = cost
r.untracked_units = untracked_units
return r
@ -279,8 +284,8 @@ async def test_detail_breaks_units_down_by_day_team_and_key():
)
assert resp.usage_units == {"contentPolicyUnits": 3, "topicPolicyUnits": 1}
assert [p.model_dump() for p in resp.usage_units_daily] == [
{"date": "2026-04-24", "units": {"topicPolicyUnits": 1}},
{"date": "2026-04-25", "units": {"contentPolicyUnits": 3}},
{"date": "2026-04-24", "units": {"topicPolicyUnits": 1}, "cost": None},
{"date": "2026-04-25", "units": {"contentPolicyUnits": 3}, "cost": None},
]
assert resp.usage_units_by_team == {
"team-a": {"contentPolicyUnits": 2, "topicPolicyUnits": 1},
@ -311,6 +316,120 @@ async def test_overview_degrades_units_to_empty_when_units_table_is_missing():
row = next(r for r in resp.rows if r.id == "yaml-uuid")
assert (row.requestsEvaluated, row.usageUnits) == (4, {})
assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {})
assert (row.cost, resp.totalCost) == (None, None)
assert (row.untrackedUsageUnits, resp.totalUntrackedUsageUnits) == ({}, {})
@pytest.mark.asyncio
async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days():
"""LIT-5652: cost rides the units rollup. Rows written before the cost column
carry NULL and rows whose every unit was unpriced carry 0.0 with
untracked_units == units; both must drop out of the sum rather than read as
$0, and a guardrail with only such rows reports None, not 0.0."""
prisma = _prisma(
find_many=[],
metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
units=[
_units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15),
_units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2000, cost=0.3),
_units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None),
_units_row(
"yaml-pii", date="2026-04-23", usage_unit="topicPolicyUnits", units=9, cost=0.0, untracked_units=9
),
_units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None),
],
)
handler = _config_handler(
_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"),
_yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"),
)
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
by_id = {r.id: r for r in resp.rows}
assert by_id["yaml-uuid"].cost == pytest.approx(0.45)
assert by_id["legacy-uuid"].cost is None
assert resp.totalCost == pytest.approx(0.45)
@pytest.mark.asyncio
async def test_overview_reports_the_units_its_cost_leaves_out_per_row_and_total():
"""A row's cost covers only the units that had a price, so the response must
say exactly which units (per counter) that cost excludes: the row's own
untracked_units, or all of its units when it predates the cost column. A
guardrail whose rows are all priced reports none, one whose rows are all
unpriced reports all of its units, and a mixed row keeps its priced subtotal
while reporting just the unpriced share."""
prisma = _prisma(
find_many=[],
metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
units=[
_units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15, untracked_units=200),
_units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None),
_units_row(
"yaml-pii", date="2026-04-24", usage_unit="topicPolicyUnits", units=40, cost=0.0, untracked_units=40
),
_units_row("yaml-pii", usage_unit="wordPolicyUnits", units=9, cost=0.0),
_units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None),
_units_row("priced-guard", usage_unit="contentPolicyUnits", units=3, cost=0.0003),
],
)
handler = _config_handler(
_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"),
_yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"),
_yaml_guardrail(guardrail_id="priced-uuid", name="priced-guard"),
)
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
by_id = {r.id: r for r in resp.rows}
assert by_id["yaml-uuid"].usageUnits == {"contentPolicyUnits": 6000, "topicPolicyUnits": 40, "wordPolicyUnits": 9}
assert by_id["yaml-uuid"].cost == pytest.approx(0.15)
assert by_id["yaml-uuid"].untrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 40}
assert by_id["legacy-uuid"].untrackedUsageUnits == {"topicPolicyUnits": 7}
assert by_id["priced-uuid"].untrackedUsageUnits == {}
assert resp.totalUntrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 47}
@pytest.mark.asyncio
async def test_detail_breaks_cost_down_by_unit_day_team_and_key():
"""Every cost breakdown keeps the same keys as its units twin so the UI can
render them side by side, with None where that group has no tracked cost."""
prisma = _prisma(
find_unique=None,
units=[
_units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=1000, cost=0.15),
_units_row(
"yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=200, cost=0.03, untracked_units=50
),
_units_row(
"yaml-pii",
date="2026-04-24",
team_id="team-a",
api_key="hash-1",
usage_unit="topicPolicyUnits",
units=10,
cost=None,
),
],
)
handler = _config_handler(_yaml_guardrail())
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_detail(
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
)
assert resp.cost == pytest.approx(0.18)
assert resp.cost_by_unit == {"contentPolicyUnits": pytest.approx(0.18), "topicPolicyUnits": None}
assert [p.model_dump() for p in resp.usage_units_daily] == [
{"date": "2026-04-24", "units": {"topicPolicyUnits": 10}, "cost": None},
{"date": "2026-04-25", "units": {"contentPolicyUnits": 1200}, "cost": pytest.approx(0.18)},
]
assert resp.cost_by_team == {"team-a": pytest.approx(0.15), "": pytest.approx(0.03)}
assert resp.cost_by_key == {"hash-1": pytest.approx(0.15), "hash-2": pytest.approx(0.03)}
assert resp.cost_by_team.keys() == resp.usage_units_by_team.keys()
assert resp.cost_by_key.keys() == resp.usage_units_by_key.keys()
assert resp.untracked_usage_units == {"contentPolicyUnits": 50, "topicPolicyUnits": 10}
@pytest.mark.asyncio
@ -330,6 +449,8 @@ async def test_detail_degrades_units_to_empty_when_units_table_is_missing():
{},
{},
)
assert (resp.cost, resp.cost_by_unit, resp.cost_by_team, resp.cost_by_key) == (None, {}, {}, {})
assert resp.untracked_usage_units == {}
# ---- logs -------------------------------------------------------------------
@ -411,6 +532,29 @@ async def test_detail_rejects_reversed_dates():
assert exc.value.status_code == 400
@pytest.mark.asyncio
async def test_policies_overview_returns_a_full_row_and_totals():
"""Regression: the policies overview shares the guardrail response model, so
every field added there (usage units, cost, untracked units) must be filled
here too or the endpoint 500s on model validation."""
policy = MagicMock(spec=["policy_id", "policy_name"])
policy.policy_id = "pol-1"
policy.policy_name = "block-pii"
metric = _metric("pol-1", requests=10, passed=8, blocked=2)
metric.policy_id = "pol-1"
prisma = _prisma()
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[policy])
prisma.db.litellm_dailypolicymetrics.find_many = AsyncMock(return_value=[metric])
p1, p2 = _patches(prisma, _config_handler())
with p1, p2:
resp = await policies_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
row = next(r for r in resp.rows if r.id == "pol-1")
assert (row.name, row.type, row.requestsEvaluated, row.failRate) == ("block-pii", "Policy", 10, 20.0)
assert (row.usageUnits, row.cost, row.untrackedUsageUnits) == ({}, None, {})
assert (resp.totalRequests, resp.totalBlocked, resp.passRate) == (10, 2, 80.0)
assert (resp.totalUsageUnits, resp.totalCost, resp.totalUntrackedUsageUnits) == ({}, None, {})
@pytest.mark.asyncio
async def test_policies_overview_rejects_range_over_max_days():
prisma = _prisma()

172
tests/test_litellm/proxy/guardrails/test_usage_tracking.py
View file

@ -30,6 +30,8 @@ def _payload(
api_key: str = "hashed-key-1",
usage: dict[str, Any] | None = None,
guardrail_status: str = "success",
cost_by_unit: dict[str, Any] | None = None,
cost_in_spend: bool | None = None,
) -> dict[str, Any]:
entry: dict[str, Any] = {
"guardrail_id": "bedrock-guard",
@ -37,6 +39,10 @@ def _payload(
}
if usage is not None:
entry["guardrail_usage"] = usage
if cost_by_unit is not None:
entry["guardrail_cost_by_unit"] = cost_by_unit
if cost_in_spend is not None:
entry["guardrail_cost_in_spend"] = cost_in_spend
return {
"request_id": request_id,
"startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc),
@ -58,6 +64,19 @@ def _units_upserts(prisma: MagicMock) -> dict[tuple, int]:
return out
def _cost_upserts(prisma: MagicMock) -> dict[str, tuple[float, int]]:
"""usage_unit -> (cost, untracked_units) written on create; the update path must increment by the same."""
calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
out: dict[str, tuple[float, int]] = {}
for c in calls:
create = c.kwargs["data"]["create"]
update = c.kwargs["data"]["update"]
assert update["cost"] == {"increment": create["cost"]}
assert update["untracked_units"] == {"increment": create["untracked_units"]}
out[create["usage_unit"]] = (create["cost"], create["untracked_units"])
return out
@pytest.mark.asyncio
async def test_usage_units_rolled_up_by_guardrail_team_key_and_date():
"""
@ -181,7 +200,9 @@ async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush():
down, [_payload("r1", usage={"topicPolicyUnits": 2})], sleep=sleep, pending=pending
)
assert dict(pending.units) == {("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2}
assert dict(pending.units) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, 0.0, 2)
}
recovered = _prisma()
await process_spend_logs_guardrail_usage(
@ -320,3 +341,152 @@ async def test_payload_without_request_id_is_skipped_like_the_metrics_path():
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
}
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]["requests_evaluated"] == 1
@pytest.mark.asyncio
async def test_cost_rolled_up_per_counter_alongside_units():
"""LIT-5652: the hook's per-counter cost lands on the same daily row as the
units it priced, summed across payloads exactly like the units are, and the
update path increments it so a second flush on the same day keeps adding."""
prisma = _prisma()
logs = [
_payload(
"r1",
usage={"contentPolicyUnits": 1000, "wordPolicyUnits": 50},
cost_by_unit={"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0},
),
_payload(
"r2",
usage={"contentPolicyUnits": 2000, "wordPolicyUnits": 10},
cost_by_unit={"contentPolicyUnits": 0.3, "wordPolicyUnits": 0.0},
),
]
await process_spend_logs_guardrail_usage(prisma, logs)
assert _units_upserts(prisma) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000,
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "wordPolicyUnits"): 60,
}
costs = _cost_upserts(prisma)
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
assert costs["wordPolicyUnits"] == (0.0, 0)
@pytest.mark.asyncio
async def test_counter_the_hook_could_not_price_is_stored_as_untracked_units_not_free():
"""A counter the cost map does not list arrives stamped as None. Its units
must land in untracked_units with no cost, so the row never reads as free,
while the priced counter on the same request keeps its cost."""
prisma = _prisma()
logs = [
_payload(
"r1",
usage={"contentPolicyUnits": 1000, "someFutureCounter": 3},
cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None},
)
]
await process_spend_logs_guardrail_usage(prisma, logs)
assert _units_upserts(prisma) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 1000,
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 3,
}
costs = _cost_upserts(prisma)
assert costs["contentPolicyUnits"] == (pytest.approx(0.15), 0)
assert costs["someFutureCounter"] == (0.0, 3)
@pytest.mark.asyncio
async def test_mixed_priced_and_unpriced_increments_keep_the_subtotal_and_count_the_rest_untracked():
"""Priced and unpriced increments on the same row (a hook without pricing,
a pre-upgrade proxy in a mixed fleet) must keep the priced subtotal and
count exactly the unpriced units as untracked. Nulling the cost would throw
away a known number; keeping it alone would look exact while understating."""
prisma = _prisma()
logs = [
_payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15}),
_payload("r2", usage={"contentPolicyUnits": 700}),
_payload("r3", usage={"contentPolicyUnits": 300}, cost_by_unit={"contentPolicyUnits": None}),
]
await process_spend_logs_guardrail_usage(prisma, logs)
assert _units_upserts(prisma) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 2000,
}
assert _cost_upserts(prisma) == {"contentPolicyUnits": (pytest.approx(0.15), 1000)}
@pytest.mark.asyncio
async def test_report_only_and_forged_costs_are_not_rolled_up_but_units_are():
"""guardrail_cost_in_spend=False (Azure Prompt Shield) keeps its cost out of
spend, so the rollup must not record it either or the dashboard would show
a number the budget never charged. A negative or non-finite per-counter cost
is treated the same way rather than subtracting from the day."""
prisma = _prisma()
logs = [
_payload("r1", usage={"text_records": 3}, cost_by_unit={"text_records": 0.5}, cost_in_spend=False),
_payload("r2", usage={"contentPolicyUnits": 10}, cost_by_unit={"contentPolicyUnits": -0.5}),
_payload("r3", usage={"topicPolicyUnits": 10}, cost_by_unit={"topicPolicyUnits": float("inf")}),
]
await process_spend_logs_guardrail_usage(prisma, logs)
assert _units_upserts(prisma) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "text_records"): 3,
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 10,
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 10,
}
assert _cost_upserts(prisma) == {
"text_records": (0.0, 3),
"contentPolicyUnits": (0.0, 10),
"topicPolicyUnits": (0.0, 10),
}
@pytest.mark.asyncio
async def test_requeued_cost_is_added_to_the_next_flush():
"""Cost and untracked units must survive the connection-error requeue the
same way units do, or a DB blip would silently drop dollars (or the record
that some units had no price) while keeping the units themselves."""
pending = PendingRollups()
down = _prisma()
down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down")
sleep, _ = _fake_sleep()
await process_spend_logs_guardrail_usage(
down,
[
_payload(
"r1",
usage={"contentPolicyUnits": 1000, "someFutureCounter": 3},
cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None},
)
],
sleep=sleep,
pending=pending,
)
recovered = _prisma()
await process_spend_logs_guardrail_usage(
recovered,
[
_payload(
"r2",
usage={"contentPolicyUnits": 2000, "someFutureCounter": 4},
cost_by_unit={"contentPolicyUnits": 0.3, "someFutureCounter": None},
)
],
sleep=sleep,
pending=pending,
)
assert _units_upserts(recovered) == {
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000,
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 7,
}
costs = _cost_upserts(recovered)
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
assert costs["someFutureCounter"] == (0.0, 7)

167
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -15,6 +15,7 @@ from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, Prisma
import litellm
import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module
from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64
from litellm.models.credentials import CredentialItem
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.health_endpoints._health_endpoints import (
@ -2675,6 +2676,172 @@ class TestConfigBaseForHealthCheck:
assert base["litellm_credential_name"] == "OpenAI-prod"
assert base["api_key"] == "sk-configured"
def test_request_naming_another_credential_does_not_inherit_config_credentials(self):
base = self._base(self.CONFIG, {"model": "openai/gpt-4o", "litellm_credential_name": "Another-cred"})
assert "api_key" not in base
assert "api_base" not in base
assert "vertex_credentials" not in base
assert base["rpm"] == 100
def test_blank_credential_name_names_no_credential(self):
base = self._base(self.CONFIG, {"model": "openai/gpt-4o", "litellm_credential_name": ""})
assert base["api_key"] == "sk-configured"
def test_opt_in_does_not_put_config_credentials_over_a_named_credential(self):
base = self._base(
self.CONFIG,
{"model": "openai/gpt-4o", "litellm_credential_name": "Another-cred"},
allow_client_side_credentials=True,
)
assert "api_key" not in base
class TestTestConnectionUsesTheNamedCredential:
CREDENTIAL_KEY = "sk-credential-key"
OTHER_DEPLOYMENT_KEY = "sk-other-deployment-key"
OTHER_DEPLOYMENT_BASE = "https://other-deployment.example/v1"
REQUEST = {
"model": "xai/grok-4",
"custom_llm_provider": "xai",
"litellm_credential_name": "my-xai-cred",
}
COMPLETION = {
"id": "chatcmpl-test",
"object": "chat.completion",
"created": 1700000000,
"model": "grok-4",
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "ok"}}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
@staticmethod
def _credential(**values: str) -> CredentialItem:
return CredentialItem(credential_name="my-xai-cred", credential_info={}, credential_values=values)
@staticmethod
def _wildcard_deployment(**litellm_params: str) -> dict:
return {
"model_name": "xai/*",
"litellm_params": {"model": "xai/*", **litellm_params},
"model_info": {"id": "unrelated-wildcard-deployment"},
}
def _probe(
self,
monkeypatch,
deployment: dict,
request_litellm_params: dict,
deployment_by_id: object | None = None,
request_model_info: dict | None = None,
) -> httpx.Request:
"""Run /health/test_connection and hand back the upstream request it made."""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
app = FastAPI()
app.include_router(_health_endpoints_module.router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
router = MagicMock()
router.get_model_list.return_value = [deployment]
router.get_deployment.return_value = deployment_by_id
with (
patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: the deployment the probe is matched against is a proxy global; it has no injection seam
"litellm.proxy.proxy_server.llm_router", router
),
respx.mock(assert_all_called=True) as respx_mock,
):
respx_mock.post(path__regex=r".*/chat/completions").respond(json=self.COMPLETION)
response = TestClient(app).post(
"/health/test_connection",
json={
"mode": "chat",
"litellm_params": request_litellm_params,
"model_info": request_model_info or {"mode": "chat"},
},
)
probe = respx_mock.calls.last.request
assert response.status_code == 200, response.text
assert response.json()["status"] == "success", response.text
return probe
def test_named_credentials_key_is_sent_not_the_matched_deployments_key(self, monkeypatch):
monkeypatch.setattr(litellm, "credential_list", [self._credential(api_key=self.CREDENTIAL_KEY)])
probe = self._probe(
monkeypatch,
self._wildcard_deployment(api_key=self.OTHER_DEPLOYMENT_KEY),
self.REQUEST,
)
assert probe.headers["authorization"] == f"Bearer {self.CREDENTIAL_KEY}"
def test_named_credentials_api_base_is_used_not_the_matched_deployments(self, monkeypatch):
monkeypatch.setattr(
litellm,
"credential_list",
[self._credential(api_key=self.CREDENTIAL_KEY, api_base="https://credential.example/v1")],
)
probe = self._probe(
monkeypatch,
self._wildcard_deployment(api_base=self.OTHER_DEPLOYMENT_BASE),
self.REQUEST,
)
assert probe.url.host == "credential.example"
def test_named_credential_without_an_api_base_leaves_the_provider_default(self, monkeypatch):
monkeypatch.setattr(litellm, "credential_list", [self._credential(api_key=self.CREDENTIAL_KEY)])
probe = self._probe(
monkeypatch,
self._wildcard_deployment(api_base=self.OTHER_DEPLOYMENT_BASE),
self.REQUEST,
)
assert probe.url.host == "api.x.ai"
def test_configured_model_named_without_a_credential_still_inherits_its_config(self, monkeypatch):
probe = self._probe(
monkeypatch,
self._wildcard_deployment(api_key=self.OTHER_DEPLOYMENT_KEY, api_base=self.OTHER_DEPLOYMENT_BASE),
{"model": "xai/grok-4", "custom_llm_provider": "xai"},
)
assert probe.headers["authorization"] == f"Bearer {self.OTHER_DEPLOYMENT_KEY}"
assert probe.url.host == "other-deployment.example"
def test_deployment_probed_by_id_keeps_the_endpoint_it_is_configured_with(self, monkeypatch):
"""The model detail page always echoes back the credential the deployment already uses."""
from litellm.types.router import Deployment, LiteLLM_Params
monkeypatch.setattr(litellm, "credential_list", [self._credential(api_key=self.CREDENTIAL_KEY)])
probe = self._probe(
monkeypatch,
self._wildcard_deployment(api_key=self.OTHER_DEPLOYMENT_KEY, api_base=self.OTHER_DEPLOYMENT_BASE),
self.REQUEST,
deployment_by_id=Deployment(
model_name="grok-4",
litellm_params=LiteLLM_Params(
model="xai/grok-4",
api_base="https://configured.example/v1",
litellm_credential_name="my-xai-cred",
),
model_info={"id": "configured-deployment"},
),
request_model_info={"id": "configured-deployment", "mode": "chat"},
)
assert probe.url.host == "configured.example"
assert probe.headers["authorization"] == f"Bearer {self.CREDENTIAL_KEY}"
class TestNoRedisWarning:
"""`show_no_redis_warning` drives the Admin UI's default-on "no Redis" banner."""

167
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -3647,6 +3647,173 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing():
assert claimed.reservation_released is True
async def _reserve_tpm_for_owner_call(handler, local_cache, api_key: str, call_id: str) -> int:
await handler.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key=api_key, tpm_limit=10_000),
cache=local_cache,
data={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": call_id,
},
call_type="completion",
)
stash = get_request_stash()
assert stash is not None and stash.reserved_tokens > 0
return stash.reserved_tokens
@pytest.mark.asyncio
async def test_failure_event_settles_tpm_reservation_at_recovered_partial_usage_v3():
"""
A stream that fails mid-way after the model already produced tokens is
logged as a failure carrying the recovered partial usage. Those tokens
were consumed, so the TPM window must settle at them instead of refunding
the whole reservation (which would let repeated timeouts burn output
tokens for free).
"""
_api_key = hash_token("sk-partial-stream-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "partial-call")
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "partial-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
stash = get_request_stash()
assert stash is not None and stash.reservation_released is True
@pytest.mark.asyncio
async def test_failure_event_refunds_reservation_for_input_only_estimate_v3():
"""
A failure with no recovered output carries only the input-token estimate
the proxy lifts onto every failure; that is not consumed usage, so the
reservation is still refunded in full.
"""
_api_key = hash_token("sk-estimated-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "estimate-call")
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "estimate-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0
@pytest.mark.asyncio
async def test_post_call_failure_hook_settles_reservation_at_recovered_partial_usage_v3():
"""
Pass-through streams report a mid-stream failure through the proxy-level
failure hook first, with the recovered usage lifted onto request_data.
That hook must settle at the partial usage too, and the later failure
callback must not double-apply it.
"""
_api_key = hash_token("sk-partial-post-call")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "post-call")
await handler.async_post_call_failure_hook(
request_data={
"model": "gpt-4o-mini",
"litellm_call_id": "post-call",
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
original_exception=Exception("upstream dropped the stream"),
user_api_key_dict=user_api_key_dict,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "post-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
@pytest.mark.asyncio
async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usage_v3():
"""
Project ITPM/OTPM reservations settle the same way: input at the billable
prompt tokens and output at the completion tokens the failed stream
actually produced.
"""
_api_key = hash_token("sk-partial-project-io")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key,
project_id="proj-partial",
project_metadata={
"model_itpm_limit": {"gpt-4o-mini": 10_000},
"model_otpm_limit": {"gpt-4o-mini": 10_000},
},
)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": "project-call",
},
call_type="completion",
)
stash = get_request_stash()
assert stash is not None and stash.itpm_reserved_tokens > 0 and stash.otpm_reserved_tokens > 0
itpm_key = handler.create_rate_limit_keys(
key="model_per_project_itpm", value="proj-partial:gpt-4o-mini", rate_limit_type="tokens"
)
otpm_key = handler.create_rate_limit_keys(
key="model_per_project_otpm", value="proj-partial:gpt-4o-mini", rate_limit_type="tokens"
)
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "project-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=itpm_key) or 0) == 20
assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7
# ----------------------- Per-MCP-server rate limiting (v3) -----------------------

1
tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py
View file

@ -61,6 +61,7 @@ def test_build_logging_obj_seeds_model_call_details():
# Metadata is mapped to the keys the cost-tracking callback reads.
md = details["litellm_params"]["metadata"]
assert md["user_api_key"] == "rust-gateway-test-key"
assert md["user_api_key_hash"] == "rust-gateway-test-key"
assert md["user_api_key_user_id"] == "user-cb-logs-test"
assert md["user_api_key_team_id"] == "team-cb-logs-test"

145
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -454,6 +454,151 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec
assert result["old-key-hash"]["team_id"] == "latest-team"
@pytest.mark.asyncio
async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash():
"""
v1.99 spend logging re-hashed already-hashed api_key values when provenance was
missing. Usage joins DailyUserSpend.api_key to VerificationToken.token, so those
rows looked like key-hash-... with a null alias. Recovery asks Postgres for the
key whose hashed token matches the dirty value and maps it back to its alias.
"""
from litellm.proxy.utils import hash_token
double_hashed = hash_token("a" * 64)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
)
mock_prisma.db.query_raw = AsyncMock(
return_value=[
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
]
)
result = await get_api_key_metadata(
prisma_client=mock_prisma,
api_keys={double_hashed},
)
assert result[double_hashed]["key_alias"] == "batch-worker"
assert result[double_hashed]["team_id"] == "team-1"
assert result[double_hashed]["user_email"] == "alice@example.com"
((digest_sql, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
assert '"LiteLLM_VerificationToken"' in digest_sql
assert digests == [double_hashed]
@pytest.mark.asyncio
async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs():
"""A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan."""
from litellm.proxy.utils import hash_token
double_hashed = hash_token("b" * 64)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
result = await get_api_key_metadata(
prisma_client=mock_prisma,
api_keys={double_hashed},
)
assert double_hashed not in result
issued_sql = [call.args[0] for call in mock_prisma.db.query_raw.call_args_list]
assert len(issued_sql) == 2
assert not any("LiteLLM_SpendLogs" in sql for sql in issued_sql)
token_lookups = (
mock_prisma.db.litellm_verificationtoken.find_many.call_args_list
+ mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list
)
assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups)
def test_key_metadata_includes_recovered_user_email():
from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata
meta = _key_metadata(
{
"dirty-key": {
"key_alias": "batch-worker",
"team_id": "team-1",
"user_email": "alice@example.com",
}
},
"dirty-key",
)
assert meta.key_alias == "batch-worker"
assert meta.user_email == "alice@example.com"
def test_update_breakdown_metrics_includes_user_email():
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
breakdown = BreakdownMetrics()
record = SimpleNamespace(
api_key="dirty-key",
model="gpt-4o-mini",
model_group="grp",
mcp_namespaced_tool_name="srv/tool",
custom_llm_provider="openai",
endpoint="/v1/chat/completions",
spend=1.23,
prompt_tokens=1,
completion_tokens=1,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
compression_saved_tokens=0,
compression_savings_spend=0,
prompt_caching_savings_spend=0,
gateway_injected_caching_savings_spend=0,
autorouter_savings_spend=0,
total_tokens=2,
api_requests=1,
successful_requests=1,
failed_requests=0,
ptu_flat_cost=0.0,
user_id="alice",
)
api_key_metadata = {
"dirty-key": {
"key_alias": "batch-worker",
"team_id": "team-1",
"user_email": "alice@example.com",
}
}
update_breakdown_metrics(
breakdown,
record,
{},
{},
api_key_metadata,
entity_id_field="user_id",
)
expected = ("batch-worker", "alice@example.com")
top = breakdown.api_keys["dirty-key"].metadata
assert (top.key_alias, top.user_email) == expected
assert (
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
assert (
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
assert (
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
@pytest.mark.asyncio
async def test_tag_daily_activity_metadata_totals_not_zero():
"""Test that tag daily activity returns correct metadata totals.

103
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -11912,6 +11912,109 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
@pytest.mark.asyncio
async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token():
"""
LIT-5379: /key/regenerate rewrites the JWT mapping row to the new token (FK
cascade) but left the jwt_key_mapping cache entry pointing at the old hash,
so JWT calls kept resolving the dead token until the cache TTL expired.
Regenerate must evict the entry locally, broadcast the eviction to other
workers, and the very next JWT resolve must return the rotated token.
"""
from litellm.caching.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.auth_method import AuthMethod
from litellm.proxy.auth.resolvers.models import CredentialRef
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
stale_cache_key = "jwt_key_mapping:sub:user1"
existing_key = _make_regenerate_existing_key()
mock_prisma_client = _make_regenerate_mock_prisma()
mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")]
)
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
return_value=MagicMock(token="new-hashed-token")
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key=stale_cache_key, value="abc123")
publish_mock = AsyncMock()
with (
patch( # test-quality-ok: deterministic token; same pattern as sibling regenerate tests
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch( # test-quality-ok: grace-period path not under test; same pattern as sibling regenerate tests
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch( # test-quality-ok: key-object eviction is separate from the mapping eviction under test
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch( # test-quality-ok: background rotation hook is irrelevant to cache eviction
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch( # test-quality-ok: captures the cross-worker broadcast without a redis instance
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
publish_mock,
),
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=existing_key,
hashed_api_key="abc123",
key="abc123",
data=None,
user_api_key_dict=_make_regenerate_user_api_key_dict(),
litellm_changed_by=None,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=MagicMock(),
)
assert await user_api_key_cache.async_get_cache(stale_cache_key) is None
publish_mock.assert_any_await(cache_key=stale_cache_key)
mock_prisma_client.db.litellm_jwtkeymapping.find_many.assert_awaited_once_with(where={"token": "abc123"})
rotated_key = UserAPIKeyAuth(token="new-hashed-token", user_id="user-1")
rotated_principal = IdentityStore._principal_from_key(
rotated_key,
auth_method=AuthMethod.API_KEY,
credential_ref=CredentialRef(token_id="new-hashed-token"),
)
async def fake_resolve(hashed_token):
assert hashed_token == "new-hashed-token", f"JWT resolved stale token {hashed_token!r} after regenerate"
return rotated_principal
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=300
)
with patch( # test-quality-ok: DB-backed resolve; fake asserts it receives the rotated hash
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=fake_resolve,
):
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={"sub": "user1"},
jwt_handler=jwt_handler,
prisma_client=mock_prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "new-hashed-token"
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(monkeypatch):
"""Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields."""

118
tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -726,6 +726,16 @@ async def _run_update_organization_v2(
return mock_prisma_client
def test_v2_update_route_is_public_in_openapi():
"""PATCH /v2/organization/{organization_id} is a public route: hiding it again (include_in_schema=False)
would drop it from openapi.json, /docs, and the generated UI API types."""
from litellm.proxy.proxy_server import get_openapi_schema
v2_path = get_openapi_schema()["paths"].get("/v2/organization/{organization_id}")
assert v2_path is not None
assert "patch" in v2_path
@pytest.mark.asyncio
async def test_v2_update_clears_tpm_limit_and_metadata(monkeypatch):
"""A cleared tpm_limit is written to the budget row as None; a cleared metadata is written as {}."""
@ -814,6 +824,54 @@ async def test_v2_rejects_negative_max_budget(monkeypatch):
assert "max_budget" in str(exc.value.detail)
@pytest.mark.asyncio
@pytest.mark.parametrize("field", ["tpm_limit", "rpm_limit", "max_parallel_requests"])
async def test_v2_rejects_negative_integer_limits(monkeypatch: pytest.MonkeyPatch, field: str):
"""v2 rejects negative tpm/rpm/parallel-request limits with a 422 instead of persisting them to the budget row."""
from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2
prisma_mock = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1")
with pytest.raises(HTTPException) as exc:
await update_organization_v2(
organization_id="org-1",
data=OrganizationUpdateRequestV2.model_validate({field: -1}),
user_api_key_dict=auth,
)
assert exc.value.status_code == 422
assert field in str(exc.value.detail)
prisma_mock.db.tx.assert_not_called()
prisma_mock.db.litellm_budgettable.update.assert_not_awaited()
prisma_mock.db.litellm_organizationtable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_v2_rejects_unparseable_budget_duration(monkeypatch: pytest.MonkeyPatch):
"""v2 rejects a budget_duration the parser can't read with a 422 instead of persisting it alongside a silent
next-midnight fallback reset."""
from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2
prisma_mock = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1")
with pytest.raises(HTTPException) as exc:
await update_organization_v2(
organization_id="org-1",
data=OrganizationUpdateRequestV2.model_validate({"budget_duration": "bogus"}),
user_api_key_dict=auth,
)
assert exc.value.status_code == 422
assert "budget_duration" in str(exc.value.detail)
prisma_mock.db.tx.assert_not_called()
prisma_mock.db.litellm_budgettable.update.assert_not_awaited()
prisma_mock.db.litellm_organizationtable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_v2_rejects_caller_without_org_access(monkeypatch):
"""v2 runs the real _verify_org_access guard: a non-admin without ORG_ADMIN on the org gets 403 and no write."""
@ -963,6 +1021,66 @@ async def test_v2_serializes_model_max_budget_on_budget_write(monkeypatch):
assert json.loads(written) == {"gpt-4o": {"max_budget": 10}}
async def _run_legacy_update_organization(
monkeypatch: pytest.MonkeyPatch, *, body: dict[str, object], existing_budget_id: str
) -> AsyncMock:
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints import organization_endpoints
from litellm.proxy.management_endpoints.organization_endpoints import update_organization
from litellm.proxy.utils import jsonify_object
mock_prisma_client = AsyncMock()
mock_prisma_client.jsonify_object = jsonify_object
existing_org = MagicMock()
existing_org.budget_id = existing_budget_id
existing_org.metadata = {}
mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=existing_org)
mock_prisma_client.db.litellm_organizationtable.update = AsyncMock(return_value=MagicMock())
mock_prisma_client.db.litellm_budgettable.update = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock())
request = MagicMock()
request.json = AsyncMock(return_value=body)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1")
await update_organization(request=request, user_api_key_dict=auth)
return mock_prisma_client
@pytest.mark.asyncio
@pytest.mark.parametrize(
"body",
[
{"organization_id": "org-1", "tpm_limit": None},
{"organization_id": "org-1", "litellm_budget_table": {"tpm_limit": None}},
],
)
async def test_legacy_update_clears_tpm_limit_when_sent_null(monkeypatch, body):
"""PATCH /organization/update with tpm_limit: null writes None to the budget row instead of dropping it."""
prisma = await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1")
budget_write = prisma.db.litellm_budgettable.update.await_args
assert budget_write.kwargs["where"] == {"budget_id": "budget-1"}
assert budget_write.kwargs["data"]["tpm_limit"] is None
assert "rpm_limit" not in budget_write.kwargs["data"]
assert "tpm_limit" not in prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]
@pytest.mark.asyncio
async def test_legacy_update_without_budget_fields_skips_budget_write(monkeypatch):
"""Omitted budget fields are left untouched: renaming the org must not write the budget row."""
prisma = await _run_legacy_update_organization(
monkeypatch,
body={"organization_id": "org-1", "organization_alias": "renamed"},
existing_budget_id="budget-1",
)
prisma.db.litellm_budgettable.update.assert_not_awaited()
assert prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]["organization_alias"] == "renamed"
def test_build_budget_write_data_recomputes_reset_at_on_duration():
"""A sent budget_duration recomputes budget_reset_at so the reset window follows the new duration."""
from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data

172
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -13598,3 +13598,175 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership"
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
def _team_spend_by_user_team(team_id: str, team_alias: str, member: Member, permissions: list[str]) -> MagicMock:
team = MagicMock(spec=LiteLLM_TeamTable)
team.team_id = team_id
team.team_alias = team_alias
team.members_with_roles = [member]
team.team_member_permissions = permissions
team.model_dump.return_value = {
"team_id": team_id,
"team_alias": team_alias,
"members_with_roles": [{"user_id": member.user_id, "role": member.role}],
"team_member_permissions": permissions,
}
return team
def _team_spend_by_user_caller(user_id: str, teams: list[str]) -> LiteLLM_UserTable:
return LiteLLM_UserTable(
user_id=user_id, user_email=f"{user_id}@example.com", teams=teams, user_role="internal_user"
)
def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, requests: int) -> dict:
return {
"team_id": team_id,
"user_id": user_id,
"user_email": f"{user_id}@example.com",
"user_alias": None,
"spend": spend,
"prompt_tokens": 10 * requests,
"completion_tokens": 5 * requests,
"total_tokens": 15 * requests,
"api_requests": requests,
"successful_requests": requests - 1,
"failed_requests": 1,
}
@pytest.mark.asyncio
async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client):
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), [])
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta])
mock_db_client.db.query_raw = AsyncMock(
return_value=[
_team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3),
_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2),
_team_spend_by_user_db_row("team-beta", "alice", 0.1, 1),
]
)
response = await get_team_spend_by_user(
user_api_key_dict=admin,
team_ids="team-alpha,team-beta",
start_date="2026-09-01",
end_date="2026-09-04",
)
sql, *params = mock_db_client.db.query_raw.call_args.args
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"]
assert 'FROM "LiteLLM_SpendLogs" sl' in sql
assert 'sl."startTime" >= $1::timestamp' in sql
assert "sl.\"startTime\" < $2::timestamp + INTERVAL '1 day'" in sql
assert "sl.team_id IN ($3, $4)" in sql
assert 'GROUP BY sl.team_id, sl."user"' in sql
assert 'sl."user" = $' not in sql
assert response.start_date == "2026-09-01"
assert response.end_date == "2026-09-04"
assert [(r.team_id, r.team_alias, r.user_id, r.user_email, r.spend, r.api_requests) for r in response.results] == [
("team-alpha", "Team Alpha", "alice", "alice@example.com", 0.5, 3),
("team-alpha", "Team Alpha", "bob", "bob@example.com", 0.25, 2),
("team-beta", "Team Beta", "alice", "alice@example.com", 0.1, 1),
]
assert (response.results[0].successful_requests, response.results[0].failed_requests) == (2, 1)
assert (response.results[0].prompt_tokens, response.results[0].completion_tokens) == (30, 15)
@pytest.mark.asyncio
async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_client):
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER)
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
mock_db_client.db.query_raw = AsyncMock(return_value=[])
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
)
await get_team_spend_by_user(
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
)
sql, *params = mock_db_client.db.query_raw.call_args.args
assert params == ["2026-09-01", "2026-09-04", "team-alpha"]
assert 'sl."user" = $' not in sql
@pytest.mark.asyncio
async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_client):
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER)
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"])
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
)
response = await get_team_spend_by_user(
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
)
sql, *params = mock_db_client.db.query_raw.call_args.args
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"]
assert "sl.team_id IN ($3)" in sql
assert 'AND sl."user" = $4' in sql
assert [(r.user_id, r.spend) for r in response.results] == [("bob", 0.25)]
@pytest.mark.asyncio
async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_client):
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER)
mock_db_client.db.query_raw = AsyncMock(return_value=[])
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
)
with pytest.raises(HTTPException) as exc_info:
await get_team_spend_by_user(
user_api_key_dict=caller, team_ids="team-beta", start_date="2026-09-01", end_date="2026-09-04"
)
assert exc_info.value.status_code == 404
mock_db_client.db.query_raw.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"team_ids,start_date,end_date,expected_error",
[
(None, "2026-09-01", "2026-09-04", "team_ids"),
("", "2026-09-01", "2026-09-04", "team_ids"),
("team-alpha", None, "2026-09-04", "start_date and end_date"),
("team-alpha", "2026-09-04", "2026-09-01", "on or after"),
("team-alpha", "2020-01-01", "2026-12-31", "at most 400 days"),
("team-alpha", "nope", "2026-09-04", "valid YYYY-MM-DD"),
],
)
async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error):
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
mock_db_client.db.query_raw = AsyncMock(return_value=[])
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as exc_info:
await get_team_spend_by_user(
user_api_key_dict=admin, team_ids=team_ids, start_date=start_date, end_date=end_date
)
assert exc_info.value.status_code == 400
assert expected_error in str(exc_info.value.detail)
mock_db_client.db.query_raw.assert_not_called()

88
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -2441,3 +2441,91 @@ class TestAnthropicPassthroughFastMode:
assert served_standard.usage.speed == "standard"
assert self._cost(served_standard) == pytest.approx(self._cost(standard))
class TestRecordPartialUsageForFailure:
"""A stream that dies mid-way still carries the usage the provider billed in
message_start; the failure row must keep it and its cost instead of logging
a zero-cost failure (or, worse, a success)."""
@staticmethod
def _sse(event, data):
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
@staticmethod
def _make_logging_obj() -> LiteLLMLoggingObj:
return LiteLLMLoggingObj(
model="claude-sonnet-5",
messages=[{"role": "user", "content": "hello"}],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="test-partial-usage-failure",
function_id="test-partial-usage-failure",
)
def _interrupted_chunks(self):
return [
self._sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_abc",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 52, "output_tokens": 1},
},
},
),
self._sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
self._sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}},
),
]
def test_stashes_partial_usage_and_cost_from_interrupted_stream(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-sonnet-5", "stream": True},
all_chunks=self._interrupted_chunks(),
)
usage = logging_obj.model_call_details["combined_usage_object"]
assert usage.prompt_tokens == 52
assert logging_obj.model_call_details["response_cost"] > 0
def test_stashes_partial_usage_at_zero_cost_when_model_is_unpriced(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-unpriced-test-model", "stream": True},
all_chunks=self._interrupted_chunks(),
)
usage = logging_obj.model_call_details["combined_usage_object"]
assert usage.prompt_tokens == 52
assert logging_obj.model_call_details["response_cost"] == 0.0
def test_leaves_logging_obj_untouched_when_nothing_streamed(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-sonnet-5", "stream": True},
all_chunks=[],
)
assert "combined_usage_object" not in logging_obj.model_call_details
assert "response_cost" not in logging_obj.model_call_details

72
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -3988,6 +3988,78 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged(
assert failure_call_kwargs["original_exception"].status_code == 403
class _UpstreamDroppingMidStream(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "hi"}}]}\n\n'
raise httpx.ReadError("upstream dropped the connection mid-stream")
async def _relay_everything(body_iterator) -> list:
return [chunk async for chunk in body_iterator]
@pytest.mark.asyncio
async def test_pass_through_request_mid_stream_upstream_drop_fires_failure_hook():
"""
Regression: a 200 stream whose upstream dies mid-body used to end with no
proxy-level failure hook at all, so the request left no spend row, no
failure metric, and no alert; the pre-stream 4xx/5xx path already fires it.
"""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
def transport_handler(upstream_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, stream=_UpstreamDroppingMidStream(), headers={"content-type": "text/event-stream"})
real_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.PassThroughEndpoint,
params={"timeout": resolve_pass_through_request_timeout(None)},
)
cache_dict = litellm.in_memory_llm_clients_cache.cache_dict
cache_key = next(key for key, cached in cache_dict.items() if cached is real_handler)
cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)))
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.scope = {"path": "/relay-chat"}
mock_request.url = MagicMock()
mock_request.url.path = "/relay-chat"
mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "stream": true}')
mock_request.headers = Headers({"content-type": "application/json"})
mock_request.query_params = QueryParams({})
try:
with patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
):
response = await pass_through_request(
request=mock_request,
target="http://target-api.com/v1/chat/completions",
custom_headers={},
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
stream=True,
)
with pytest.raises(httpx.ReadError):
await _relay_everything(response.body_iterator)
await asyncio.sleep(0)
finally:
cache_dict[cache_key] = real_handler
mock_proxy_logging.post_call_failure_hook.assert_awaited_once()
failure_call_kwargs = mock_proxy_logging.post_call_failure_hook.call_args.kwargs
assert isinstance(failure_call_kwargs["original_exception"], httpx.ReadError)
request_data = failure_call_kwargs["request_data"]
assert request_data["litellm_call_id"]
assert request_data["model"] == "gpt-5.6"
assert isinstance(request_data["litellm_logging_obj"], LiteLLMLoggingObj)
@pytest.mark.asyncio
async def test_pass_through_request_non_streaming_success_unchanged():
"""Success (2xx) passthrough behavior must remain unchanged by the error fix."""

206
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -9,6 +9,8 @@ import httpx
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
@ -632,3 +634,207 @@ async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_arme
mock_enqueue.assert_called_once()
assert logging_obj._deferred_stream_complete_args is None
class _EventRecorder(CustomLogger):
def __init__(self):
super().__init__()
self.failure_kwargs = []
self.success_kwargs = []
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.failure_kwargs.append(kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.success_kwargs.append(kwargs)
def _anthropic_sse(event: str, payload: dict) -> bytes:
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
def _anthropic_stream_that_times_out_mid_stream():
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
async def _aiter_bytes():
yield _anthropic_sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 52, "output_tokens": 1},
},
},
)
yield _anthropic_sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
)
yield _anthropic_sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}},
)
raise httpx.ReadTimeout("Timeout on reading data from socket")
mock.aiter_bytes = _aiter_bytes
return mock
@pytest.mark.asyncio
async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception():
"""A stream that dies after the first chunks is a failed request: the failure
callbacks must fire once with the partial usage and cost, and the success
routing must never run for it."""
recorder = _EventRecorder()
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="test-mid-stream-timeout",
function_id="test-mid-stream-timeout",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
success_routes = []
async def _record_success_route(**kwargs):
success_routes.append(kwargs)
received = []
async def _consume_stream():
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=_anthropic_stream_that_times_out_mid_stream(),
request_body={"model": "claude-sonnet-5", "stream": True},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=_record_success_route,
):
received.append(chunk)
with pytest.raises(httpx.ReadTimeout):
await _consume_stream()
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert len(received) == 3
assert success_routes == []
assert recorder.success_kwargs == []
assert len(recorder.failure_kwargs) == 1
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
assert failure_payload["status"] == "failure"
assert failure_payload["prompt_tokens"] == 52
assert failure_payload["response_cost"] > 0
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)
def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes:
payload = {
"candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}],
"usageMetadata": {
"promptTokenCount": prompt_tokens,
"candidatesTokenCount": completion_tokens,
"totalTokenCount": prompt_tokens + completion_tokens,
},
"modelVersion": "gemini-3.8-flash",
}
return f"data: {json.dumps(payload)}\r\n\r\n".encode()
def _google_stream_that_times_out_mid_stream():
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
async def _aiter_bytes():
yield _google_sse(9, 4, "The sea")
yield _google_sse(9, 12, " is wide and restless")
raise httpx.ReadTimeout("Timeout on reading data from socket")
mock.aiter_bytes = _aiter_bytes
return mock
@pytest.mark.parametrize(
"endpoint_type, url_route",
[
(EndpointType.GEMINI, "/gemini/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse"),
(
EndpointType.VERTEX_AI,
"/vertex_ai/v1/projects/p/locations/us-central1/publishers/google/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
),
],
)
@pytest.mark.asyncio
async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exception(endpoint_type, url_route):
"""Google streams carry cumulative usage on every chunk, so a stream that
dies mid-way must log a failure billed at what was already delivered rather
than a failure at zero usage."""
recorder = _EventRecorder()
logging_obj = LiteLLMLoggingObj(
model="gemini-3.8-flash",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id=f"test-google-mid-stream-timeout-{endpoint_type.value}",
function_id="test-google-mid-stream-timeout",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
logging_obj.update_environment_variables(
model="gemini-3.8-flash",
user="unknown",
optional_params={},
litellm_params={"metadata": {}},
call_type="pass_through_endpoint",
)
success_routes = []
async def _record_success_route(**kwargs):
success_routes.append(kwargs)
async def _consume_stream():
async for _ in PassThroughStreamingHandler.chunk_processor(
response=_google_stream_that_times_out_mid_stream(),
request_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route=url_route,
route_streaming_logging=_record_success_route,
):
pass
with pytest.raises(httpx.ReadTimeout):
await _consume_stream()
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert success_routes == []
assert recorder.success_kwargs == []
assert len(recorder.failure_kwargs) == 1
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
assert failure_payload["status"] == "failure"
assert failure_payload["prompt_tokens"] == 9
assert failure_payload["completion_tokens"] == 12
assert failure_payload["response_cost"] > 12 * 3.75e-06
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)

2
tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
View file

@ -1108,7 +1108,7 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog(
assert payload["1m_context"]["complexity_router_config"]["tiers"] == {
"SIMPLE": ["gpt-5.6-luna"],
"MEDIUM": ["gpt-5.6-terra"],
"COMPLEX": ["claude-opus-5"],
"COMPLEX": ["gpt-5.6-sol"],
"REASONING": ["claude-opus-5"],
}
assert payload["1m_context"]["complexity_router_config"]["tier_model_configs"] == {

220
tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py Normal file
View file

@ -0,0 +1,220 @@
from collections.abc import Sequence
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from prisma.errors import PrismaError
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
recover_double_hashed_key_metadata,
)
from litellm.proxy.utils import hash_token
def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]:
return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id}
def _query_raw_by_table(
active_rows: Sequence[dict[str, str | None]],
deleted_rows: Sequence[dict[str, str | None]],
) -> AsyncMock:
async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]:
if '"LiteLLM_VerificationToken"' in sql:
return list(active_rows)
if '"LiteLLM_DeletedVerificationToken"' in sql:
return list(deleted_rows)
raise AssertionError(f"unexpected query: {sql}")
return AsyncMock(side_effect=query_raw)
@pytest.mark.asyncio
async def test_recover_double_hashed_key_metadata_via_active_token_digest():
double_hashed = hash_token("a" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "batch-worker", "team-1", "alice")],
deleted_rows=[],
)
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result[double_hashed]["key_alias"] == "batch-worker"
assert result[double_hashed]["team_id"] == "team-1"
assert result[double_hashed]["user_id"] == "alice"
((_, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
assert digests == [double_hashed]
@pytest.mark.asyncio
async def test_recover_double_hashed_key_metadata_falls_back_to_deleted_tokens():
double_hashed = hash_token("y" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[],
deleted_rows=[_digest_row(double_hashed, "deleted-key", "team-del", "erin")],
)
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result[double_hashed]["key_alias"] == "deleted-key"
assert result[double_hashed]["team_id"] == "team-del"
assert result[double_hashed]["user_id"] == "erin"
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [[double_hashed], [double_hashed]]
@pytest.mark.asyncio
async def test_recover_only_asks_deleted_tokens_for_digests_active_keys_missed():
found_active = hash_token("1" * 64)
found_deleted = hash_token("2" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(found_active, "active-key", None, None)],
deleted_rows=[_digest_row(found_deleted, "deleted-key", None, None)],
)
result = await recover_double_hashed_key_metadata(mock_prisma, {found_active, found_deleted})
assert result[found_active]["key_alias"] == "active-key"
assert result[found_deleted]["key_alias"] == "deleted-key"
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [
sorted((found_active, found_deleted)),
[found_deleted],
]
@pytest.mark.asyncio
async def test_recover_permanent_miss_costs_two_digest_lookups_and_no_table_walk():
double_hashed = hash_token("b" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(active_rows=[], deleted_rows=[])
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result == {}
assert len(mock_prisma.db.query_raw.call_args_list) == 2
mock_prisma.db.litellm_verificationtoken.find_many.assert_not_called()
mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called()
@pytest.mark.asyncio
async def test_recover_skips_keys_that_are_not_sha256_digests():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
result = await recover_double_hashed_key_metadata(mock_prisma, {"sk-plain-key", "key-hash-short"})
assert result == {}
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_recover_returns_empty_when_digest_lookup_raises_prisma_error():
double_hashed = hash_token("c" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down"))
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result == {}
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows():
double_hashed = hash_token("d" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "recovered-alias", "team-9", "bob")],
deleted_rows=[],
)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")]
)
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": None,
"spend": 12.5,
},
{
"api_key": "already-joined-token",
"api_key_alias": "named-key",
"team_id": "team-ok",
"user_email": "other@example.com",
"spend": 1.0,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled[0]["api_key_alias"] == "recovered-alias"
assert filled[0]["team_id"] == "team-9"
assert filled[0]["user_email"] == "bob@example.com"
assert filled[0]["spend"] == 12.5
assert filled[1]["api_key_alias"] == "named-key"
assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["bob"]}}
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_is_missing():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = ({"api_key": hash_token("e" * 64), "api_key_alias": "named", "user_email": "x@example.com"},)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_keeps_spend_user_email_when_alias_is_missing():
double_hashed = hash_token("f" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "team-key", "team-9", "key-owner")],
deleted_rows=[],
)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="key-owner", user_email="owner@example.com")]
)
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": "spender@example.com",
"spend": 4.0,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled[0]["api_key_alias"] == "team-key"
assert filled[0]["team_id"] == "team-9"
assert filled[0]["user_email"] == "spender@example.com"
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = (
{
"api_key": hash_token("g" * 64),
"api_key_alias": "service-key",
"team_id": "team-svc",
"user_email": None,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()

5
tests/test_litellm/proxy/spend_tracking/test_savings.py
View file

@ -852,7 +852,7 @@ def test_the_served_arm_is_read_from_the_record_not_repriced():
@pytest.mark.parametrize(
"basis, expected_multiplier",
[
pytest.param({"service_tier": "priority"}, 2.0, id="priority tier doubles the baseline"),
pytest.param({"service_tier": "priority"}, 2.5, id="priority tier uplifts the baseline"),
pytest.param({"data_residency": "eu"}, 1.1, id="eu residency uplifts the baseline"),
pytest.param({}, 1.0, id="no basis recorded prices at standard"),
pytest.param(None, 1.0, id="row predating the field prices at standard"),
@ -872,7 +872,8 @@ def test_the_baseline_is_priced_on_the_basis_the_request_was_billed_at(basis, ex
"""
gpt = litellm.get_model_info("gpt-5.5", "openai")
haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic")
assert gpt.get("input_cost_per_token_priority") == 2 * gpt["input_cost_per_token"]
assert gpt.get("input_cost_per_token_priority") == pytest.approx(2.5 * gpt["input_cost_per_token"])
assert gpt.get("output_cost_per_token_priority") == pytest.approx(2.5 * gpt["output_cost_per_token"])
assert gpt.get("regional_processing_uplift_multiplier_eu") == 1.1
assert haiku.get("input_cost_per_token_priority") is None, "served model must not move with the basis"
assert haiku.get("regional_processing_uplift_multiplier_eu") is None

35
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -2712,6 +2712,41 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed():
assert meta["user_api_key"] == hash_token(already_hashed)
def test_get_logging_payload_batch_attribution_keeps_verification_token_hash():
"""
Batch cost rebuilds metadata with the managed object's already-hashed api_key.
That hash must land in SpendLogs.api_key unchanged so Usage/CloudZero can join
LiteLLM_VerificationToken for api_key_alias and user_email. Regression: without
user_api_key_hash provenance, v1.99+ re-hashed the token and broke the join.
"""
token_hash = hash_token("sk-batch-creator-key")
kwargs = {
"model": "gpt-4o",
"call_type": "aretrieve_batch",
"litellm_params": {
"metadata": {
"user_api_key": token_hash,
"user_api_key_hash": token_hash,
"user_api_key_alias": "batch-creator",
"user_api_key_user_id": "alice",
"user_api_key_team_id": "team-1",
}
},
}
payload = get_logging_payload(
kwargs=kwargs,
response_obj={"id": "batch_123", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["api_key"] == token_hash
assert payload["api_key"] != hash_token(token_hash)
parsed_meta = json.loads(payload["metadata"])
assert parsed_meta["user_api_key"] == token_hash
assert parsed_meta["user_api_key_alias"] == "batch-creator"
def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match():
already_hashed = hash_token("sk-some-key")
different_hash = hash_token("sk-other-key")

109
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -7836,3 +7836,112 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()]
assert len(records) == 1
assert (records[0].exc_info is not None) is expect_traceback
class _FailureHookRecorder:
"""Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
def __init__(self, raises: Optional[Exception] = None):
self.calls = []
self._raises = raises
async def post_call_failure_hook(self, **kwargs):
self.calls.append(kwargs)
if self._raises is not None:
raise self._raises
class TestDetachedStreamFailureHook:
"""
Regression for LIT-3798. A streaming /v1/messages request whose client disconnected
before the provider failed mid-stream never reached the proxy's failure hook: the
client-facing generator was gone, and the detached upstream drain only fired the
logging object's callbacks, so no failure spend row was written and the budget
reservation stayed held. base_process_llm_request now arms a closure on the logging
object that the detached drain awaits, and that closure runs post_call_failure_hook
with the request's key and data.
"""
@staticmethod
def _logging_obj():
logging_obj = MagicMock()
logging_obj.litellm_call_id = "call-lit3798"
logging_obj.model_call_details = {}
logging_obj._enqueue_deferred_logging = None
logging_obj._on_deferred_stream_complete = None
logging_obj._on_detached_stream_failure = None
return logging_obj
@staticmethod
def _proxy_logging_obj(recorder: _FailureHookRecorder):
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook
return proxy_logging_obj
@pytest.mark.asyncio
async def test_streaming_messages_arms_the_detached_failure_hook(self, monkeypatch):
import litellm.proxy.common_request_processing as crp
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
async def _stream():
yield b"event: message_start\n\n"
async def fake_route_request(**kwargs):
async def _llm_call():
return _stream()
return _llm_call()
monkeypatch.setattr(crp, "route_request", fake_route_request)
monkeypatch.setattr(litellm, "callbacks", [])
recorder = _FailureHookRecorder()
logging_obj = self._logging_obj()
user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test")
processing_obj = ProxyBaseLLMRequestProcessing(
data={"litellm_logging_obj": logging_obj, "model": "claude-sonnet-4-5"}
)
await processing_obj.base_process_llm_request(
request=MagicMock(spec=Request, headers={}),
fastapi_response=Response(),
user_api_key_dict=user_api_key_dict,
route_type="anthropic_messages",
proxy_logging_obj=self._proxy_logging_obj(recorder),
general_settings={},
proxy_config=MagicMock(spec=ProxyConfig),
select_data_generator=None,
llm_router=None,
skip_pre_call_logic=True,
)
failure = RuntimeError("upstream died after the client left")
await logging_obj._on_detached_stream_failure(failure)
assert recorder.calls == [
{
"user_api_key_dict": user_api_key_dict,
"original_exception": failure,
"request_data": processing_obj.data,
}
]
@pytest.mark.asyncio
async def test_detached_failure_hook_drops_the_replacement_error_it_cannot_deliver(self):
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
recorder = _FailureHookRecorder(raises=HTTPException(status_code=429, detail="budget exceeded"))
logging_obj = self._logging_obj()
processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj})
processing_obj._arm_detached_stream_failure_hook(
logging_obj=logging_obj,
user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=self._proxy_logging_obj(recorder),
)
failure = RuntimeError("upstream died after the client left")
await logging_obj._on_detached_stream_failure(failure)
assert [call["original_exception"] for call in recorder.calls] == [failure]

831
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -15,7 +15,6 @@ from pydantic import ValidationError
import litellm
from litellm import Router
from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
@ -27,6 +26,8 @@ from litellm.router_strategy.complexity_router.complexity_router import (
DimensionScore,
KeywordOverride,
_built_in_prompt,
_ClassifierCircuitBreaker,
_is_classifier_timeout,
_matched_plan_mode_sentinel,
classification_system_prompt,
)
@ -44,6 +45,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
TierGlobalStatistic,
TrainedTierArtifact,
)
from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
from litellm.types.router import (
Deployment,
LiteLLM_Params,
@ -1724,6 +1726,13 @@ class TestLLMClassifierConfig:
assert config.classifier_type == "heuristic"
assert config.classifier_llm_config is None
def test_classifier_circuit_breaker_defaults_on_and_requires_positive_cooldown(self):
config = ClassifierLLMConfig(model="haiku-classifier")
assert config.circuit_breaker_enabled is True
assert config.circuit_breaker_cooldown_seconds == 30.0
with pytest.raises(ValidationError):
ClassifierLLMConfig(model="haiku-classifier", circuit_breaker_cooldown_seconds=0)
@pytest.mark.parametrize("reasoning_effort", ["", "ultra"])
def test_classifier_reasoning_effort_rejects_unsupported_values(self, reasoning_effort):
with pytest.raises(ValidationError):
@ -2000,6 +2009,203 @@ class TestLLMClassifier:
assert outcome.cause == "llm_classifier"
assert outcome.classifier_cost == pytest.approx(1.35e-05)
@pytest.mark.asyncio
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(
self, llm_classifier_config
):
real_router = Router(
model_list=[
{
"model_name": "haiku-classifier",
"litellm_params": {
"model": "openai/mock-classifier",
"api_key": "mock-key",
"mock_timeout": True,
},
},
{
"model_name": "backup-classifier",
"litellm_params": {
"model": "openai/mock-backup-classifier",
"api_key": "mock-key",
"mock_response": '{"tier": "COMPLEX"}',
},
},
],
num_retries=2,
fallbacks=[{"haiku-classifier": ["backup-classifier"]}],
)
config = {
**llm_classifier_config,
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10},
}
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=real_router,
complexity_router_config=config,
)
outcome = await router.aclassify("hi")
next_outcome = await router.aclassify("hi again")
assert outcome.cause == "heuristic_scorer"
assert next_outcome.cause == "heuristic_scorer"
assert "classifier-circuit-open" in next_outcome.signals
assert real_router.total_calls["openai/mock-classifier"] == 1
assert real_router.total_calls["openai/mock-backup-classifier"] == 0
@pytest.mark.asyncio
async def test_aclassify_enforces_total_classifier_deadline(
self, mock_router_instance, llm_classifier_config
):
cancelled = asyncio.Event()
async def slow_classifier(**_kwargs: object) -> None:
try:
await asyncio.sleep(1)
except asyncio.CancelledError:
cancelled.set()
raise
mock_router_instance.acompletion = AsyncMock(side_effect=slow_classifier)
config = {
**llm_classifier_config,
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10},
}
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=config,
)
outcome = await router.aclassify("hi")
assert outcome.cause == "heuristic_scorer"
assert cancelled.is_set()
@pytest.mark.asyncio
async def test_timeout_opens_classifier_circuit_for_other_sessions(
self, mock_router_instance, llm_classifier_config
):
"""One classifier outage is deployment-wide, so a second session must not pay the timeout."""
mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out"))
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=llm_classifier_config,
)
first = await router.aclassify("first ask", request_kwargs={"metadata": {"session_id": "session-a"}})
second = await router.aclassify("second ask", request_kwargs={"metadata": {"session_id": "session-b"}})
assert first.cause == "heuristic_scorer"
assert second.cause == "heuristic_scorer"
assert "classifier-circuit-open" in second.signals
mock_router_instance.acompletion.assert_awaited_once()
def test_classifier_circuit_allows_one_probe_and_closes_on_success(self):
now = 100.0
breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
initial_permit = breaker.acquire_permit()
assert initial_permit is not None
breaker.record_failure(initial_permit, is_timeout=True)
assert breaker.acquire_permit() is None
now = 130.0
probe_permit = breaker.acquire_permit()
assert probe_permit is not None
assert breaker.acquire_permit() is None
breaker.record_success(probe_permit)
assert breaker.acquire_permit() is not None
def test_failed_classifier_probe_restarts_cooldown(self):
now = 100.0
breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
initial_permit = breaker.acquire_permit()
assert initial_permit is not None
breaker.record_failure(initial_permit, is_timeout=True)
now = 130.0
probe_permit = breaker.acquire_permit()
assert probe_permit is not None
breaker.record_failure(probe_permit, is_timeout=False)
assert breaker.acquire_permit() is None
now = 160.0
assert breaker.acquire_permit() is not None
def test_stale_success_cannot_close_circuit_opened_by_overlapping_timeout(self):
breaker = _ClassifierCircuitBreaker(30.0)
timeout_permit = breaker.acquire_permit()
stale_success_permit = breaker.acquire_permit()
assert timeout_permit is not None
assert stale_success_permit is not None
breaker.record_failure(timeout_permit, is_timeout=True)
breaker.record_success(stale_success_permit)
assert breaker.acquire_permit() is None
@pytest.mark.asyncio
async def test_cancelled_classifier_probe_restarts_cooldown(self, mock_router_instance, llm_classifier_config):
now = 100.0
mock_router_instance.acompletion = AsyncMock(
side_effect=[
TimeoutError("classifier timed out"),
asyncio.CancelledError(),
_llm_response('{"tier": "SIMPLE"}'),
]
)
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=llm_classifier_config,
)
router._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
await router.aclassify("open the circuit")
now = 130.0
with pytest.raises(asyncio.CancelledError):
await router.aclassify("cancel the recovery probe")
outcome = await router.aclassify("stay in cooldown")
assert outcome.cause == "heuristic_scorer"
assert "classifier-circuit-open" in outcome.signals
assert mock_router_instance.acompletion.await_count == 2
@pytest.mark.asyncio
async def test_classifier_circuit_can_be_disabled(self, mock_router_instance, llm_classifier_config):
mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out"))
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={
**llm_classifier_config,
"classifier_llm_config": {
**llm_classifier_config["classifier_llm_config"],
"circuit_breaker_enabled": False,
},
},
)
await router.aclassify("first ask")
await router.aclassify("second ask")
assert mock_router_instance.acompletion.await_count == 2
def test_non_timeout_failure_does_not_open_closed_classifier_circuit(self):
breaker = _ClassifierCircuitBreaker(30.0)
permit = breaker.acquire_permit()
assert permit is not None
breaker.record_failure(permit, is_timeout=False)
assert breaker.acquire_permit() is not None
def test_asyncio_timeout_is_a_classifier_timeout_on_python_310(self):
assert _is_classifier_timeout(asyncio.TimeoutError()) is True
@pytest.mark.asyncio
async def test_aclassify_classifier_cost_is_none_when_call_is_unpriced(
self, llm_complexity_router, mock_router_instance
@ -4564,6 +4770,45 @@ class TestSessionAffinity:
# Pinned to the first turn's model, not re-classified down to SIMPLE.
assert second.model == "o1-preview"
@pytest.mark.asyncio
async def test_circuit_open_fallback_does_not_pin_the_session(self, mock_router_instance, session_affinity_config):
"""Regression: the classifier circuit cools down in seconds while a pin lasts for the whole
TTL, so a session whose only turn landed on the cooldown fallback must classify again once
the breaker closes instead of holding that fallback's model."""
now = 100.0
mock_router_instance.cache = DualCache()
mock_router_instance.acompletion = AsyncMock(
side_effect=[TimeoutError("classifier timed out"), _llm_response('{"tier": "REASONING"}')]
)
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={
**session_affinity_config,
"classifier_type": "llm",
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400},
},
)
router._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now)
await router.async_pre_routing_hook(
model="test-model",
request_kwargs=self._request_kwargs("outage-session"),
messages=self.SIMPLE_MESSAGE,
)
cooled_down_kwargs = self._request_kwargs("cooldown-session")
during_cooldown = await router.async_pre_routing_hook(
model="test-model", request_kwargs=cooled_down_kwargs, messages=self.SIMPLE_MESSAGE
)
now = 130.0
after_cooldown = await router.async_pre_routing_hook(
model="test-model", request_kwargs=cooled_down_kwargs, messages=self.SIMPLE_MESSAGE
)
assert during_cooldown.model == "gpt-4o-mini"
assert after_cooldown.model == "o1-preview"
assert mock_router_instance.acompletion.await_count == 2
@pytest.mark.asyncio
async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config):
mock_router_instance.cache = DualCache()
@ -8476,7 +8721,8 @@ class TestClassifierFallbackChoice:
@pytest.mark.asyncio
async def test_a_classifier_failure_does_not_pin_the_session_to_the_default_model(self, mock_router_instance):
"""One transient timeout must not hold a session on default_model for the whole affinity TTL:
that turn was never classified, so there is nothing worth pinning and the next turn retries."""
that turn was never classified, so there is nothing worth pinning. The circuit breaker is
disabled here so the next turn isolates and verifies the affinity contract."""
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
@ -8488,7 +8734,11 @@ class TestClassifierFallbackChoice:
"REASONING": "o1-preview",
},
"classifier_type": "llm",
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400},
"classifier_llm_config": {
"model": "haiku-classifier",
"timeout_ms": 400,
"circuit_breaker_enabled": False,
},
"classifier_fallback": "default_model",
"default_model": "gpt-4o",
"session_affinity": True,
@ -11321,3 +11571,578 @@ class TestModalityRouting:
model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE
)
assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}
class TestTierHealthFailover:
"""A tier whose decided model group is entirely in cooldown falls back to a live peer."""
SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}]
TIERS = {"SIMPLE": ["dead-a", "live-b"], "MEDIUM": "mid", "COMPLEX": "big", "REASONING": "top"}
@staticmethod
def _router(
mock_router_instance,
config,
ids_by_model,
cooling=(),
blocked=(),
excluded=(),
raises_for=None,
health_error=None,
):
"""ids_by_model: model group -> deployment ids the router knows.
The fake mirrors the real async_get_healthy_deployments contract, including how it says
no: BadRequestError for a group with no deployment at all, RouterRateLimitError when every
deployment is filtered out (cooling, admin-paused, or excluded by a request-scoped policy
such as tags, team scoping or access groups), a per-model exception via raises_for (the
RPM verdict), and an unrelated failure via health_error. It records what it was handed so
tests can prove the probe passes a kwargs copy and forwards the prompt arguments.
"""
import litellm as litellm_module
from litellm.types.router import RouterRateLimitError
probed_kwargs = []
probed_prompts = []
async def get_healthy_deployments(
model, request_kwargs, messages=None, input=None, parent_otel_span=None, **kwargs
):
probed_kwargs.append(request_kwargs)
probed_prompts.append((messages, input))
if health_error is not None:
raise health_error
if raises_for and model in raises_for:
raise raises_for[model]
if not ids_by_model.get(model):
raise litellm_module.BadRequestError(
message=f"You passed in model={model}. There are no healthy deployments.",
model=model,
llm_provider="",
)
filtered = (*cooling, *blocked, *excluded)
healthy = [
{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered
]
if not healthy:
raise RouterRateLimitError(
model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[]
)
return healthy
mock_router_instance.async_get_healthy_deployments = get_healthy_deployments
mock_router_instance.probed_kwargs = probed_kwargs
mock_router_instance.probed_prompts = probed_prompts
mock_router_instance.cache = DualCache()
return ComplexityRouter(
model_name="health-test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=config,
)
async def _pinned_hook(self, router, session_id="sess-1", messages=None):
"""Drive the hook twice so the second call replays a pin, which makes the decided
model deterministic instead of a coin flip over the tier pool."""
kwargs = {"metadata": {"session_id": session_id}}
await router.async_pre_routing_hook(model="m", request_kwargs=kwargs, messages=messages or self.SIMPLE_MESSAGE)
return await router.async_pre_routing_hook(
model="m", request_kwargs=kwargs, messages=messages or self.SIMPLE_MESSAGE
)
@pytest.mark.asyncio
async def test_dead_pinned_group_fails_over_to_live_peer_and_reports_the_displacement(self, mock_router_instance):
"""The core regression: a session pinned to a group whose every deployment is cooling
serves from the live peer, and the row says so rather than naming the pinned model."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1", "id-a2"], "live-b": ["id-b1"]},
cooling=("id-a1", "id-a2"),
)
# Seed the pin onto the dead group directly so the replay path is exercised.
key = router._get_session_affinity_cache_key("sess-dead", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-dead"}}, messages=self.SIMPLE_MESSAGE
)
assert result.model == "live-b"
assert result.routing_decision["cause"] == "health_failover"
assert "health_displaced:dead-a" in result.routing_decision["signals"]
assert result.routing_decision["tier"] == "SIMPLE"
@pytest.mark.asyncio
async def test_fresh_classification_never_serves_a_fully_cooled_group(self, mock_router_instance):
"""The pool pick is a uniform draw, so the invariant is asserted over repeated turns:
no turn may land on the dead group while a live peer sits in the same tier."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS)},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
results = [
await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.SIMPLE_MESSAGE)
for _ in range(20)
]
assert {r.model for r in results} == {"live-b"}
assert all(r.routing_decision["cause"] in ("heuristic_scorer", "health_failover") for r in results)
assert any(r.routing_decision["cause"] == "health_failover" for r in results)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"ids_by_model, cooling, health_error, tiers, reason",
[
({"dead-a": ["id-a1"], "live-b": ["id-b1"]}, (), None, None, "nothing_cooling"),
({"dead-a": ["id-a1"], "live-b": ["id-b1"]}, ("id-a1", "id-b1"), None, None, "every_peer_dead"),
(
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
("id-a1",),
RuntimeError("redis down"),
None,
"health_view_unreadable",
),
(
{"only": ["id-1"]},
("id-1",),
None,
{"SIMPLE": "only", "MEDIUM": "mid", "COMPLEX": "big", "REASONING": "top"},
"single_model_tier_has_no_peer",
),
],
)
async def test_gate_fails_open_and_leaves_the_decision_untouched(
self, mock_router_instance, ids_by_model, cooling, health_error, tiers, reason
):
"""Every uncertainty leaves the decided model in place, so the request fails exactly
as it does today rather than being rerouted on a guess."""
router = self._router(
mock_router_instance,
{"tiers": dict(tiers or self.TIERS), "session_affinity": True},
ids_by_model,
cooling=cooling,
health_error=health_error,
)
pinned = "only" if tiers else "dead-a"
key = router._get_session_affinity_cache_key("sess-open", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": pinned, "tier": "SIMPLE"}, ttl=600
)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-open"}}, messages=self.SIMPLE_MESSAGE
)
assert result.model == pinned, reason
assert result.routing_decision["cause"] == "session_affinity_pin", reason
@pytest.mark.asyncio
async def test_a_failed_over_turn_is_never_pinned(self, mock_router_instance):
"""A failover describes the fleet's state, not the session's traffic, so it must not
become the pin: the substitute would outlive the outage that caused it.
Asserted over many sessions because the underlying pool pick is a uniform draw.
"""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
async def pin_after_session(turn: int):
session_id = f"sess-write-{turn}"
await router.async_pre_routing_hook(
model="m",
request_kwargs={"metadata": {"session_id": session_id}},
messages=self.SIMPLE_MESSAGE,
)
return await router.litellm_router_instance.cache.async_get_cache(
key=router._get_session_affinity_cache_key(session_id, {})
)
stored = [await pin_after_session(turn) for turn in range(20)]
assert all(entry in (None, {"model": "live-b", "tier": "SIMPLE"}) for entry in stored)
assert any(entry is None for entry in stored), "a failed-over turn must leave the pin unwritten"
@pytest.mark.asyncio
async def test_an_unpinnable_displaced_cause_stays_unpinnable_after_failover(self, mock_router_instance):
"""A housekeeping turn is deliberately never pinned. Rewriting its cause to health_failover
must not smuggle it past that guard and lock the session onto the cheapest tier."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
session_id = "sess-housekeeping"
result = await router.async_pre_routing_hook(
model="m",
request_kwargs={"metadata": {"session_id": session_id}},
messages=[{"role": "user", "content": TITLE_ASK}],
)
assert result.routing_decision["cause"] in ("housekeeping", "health_failover")
stored = await router.litellm_router_instance.cache.async_get_cache(
key=router._get_session_affinity_cache_key(session_id, {})
)
assert stored is None
@pytest.mark.asyncio
async def test_a_peer_whose_deployments_are_admin_paused_is_not_a_failover_target(self, mock_router_instance):
"""Capacity is the router's own verdict, not just cooldown: a paused peer would be
rejected downstream and the request would fail with a live third peer available."""
router = self._router(
mock_router_instance,
{
"tiers": {
"SIMPLE": ["dead-a", "paused-b", "live-c"],
"MEDIUM": "mid",
"COMPLEX": "big",
"REASONING": "top",
},
"session_affinity": True,
},
{"dead-a": ["id-a1"], "paused-b": ["id-b1"], "live-c": ["id-c1"]},
cooling=("id-a1",),
blocked=("id-b1",),
)
key = router._get_session_affinity_cache_key("sess-paused", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-paused"}}, messages=self.SIMPLE_MESSAGE
)
for _ in range(20)
]
assert {r.model for r in results} == {"live-c"}
@pytest.mark.asyncio
async def test_failover_fails_closed_when_a_routing_plugin_excludes_every_peer(self, mock_router_instance):
"""A plugin's exclusion is policy, so a peer it removed must not be served just because
the plugin's own choice went into cooldown."""
class ExcludeEverythingButDead:
async def run(self, context):
context.candidate_models = [m for m in context.candidate_models if m == "dead-a"]
return context
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "plugins": [ExcludeEverythingButDead()]},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.SIMPLE_MESSAGE)
assert result.model == "dead-a"
assert result.routing_decision["cause"] != "health_failover"
@pytest.mark.asyncio
async def test_failover_moves_the_adaptive_chosen_model_marker(self, mock_router_instance):
"""The adaptive feedback loop scores the marker, so leaving it on the displaced group
would credit a model that never ran."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-adaptive", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
request_kwargs = {"metadata": {"session_id": "sess-adaptive", "adaptive_router_chosen_model": "dead-a"}}
result = await router.async_pre_routing_hook(
model="m", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert result.model == "live-b"
assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "live-b"
@pytest.mark.asyncio
async def test_health_failover_never_undoes_the_modality_gate(self, mock_router_instance):
"""An image turn whose only live peer cannot take images keeps the vision model the
modality gate chose: serving a cooling vision model beats a hard 400."""
vision_by_model = {"dead-vision": True, "live-text": False}
def get_model_list(model_name=None):
if model_name not in vision_by_model:
return []
return [
{
"model_name": model_name,
"litellm_params": {"model": f"openai/unmapped-{model_name}"},
"model_info": {"supports_vision": vision_by_model[model_name]},
}
]
mock_router_instance.get_model_list = get_model_list
router = self._router(
mock_router_instance,
{
"tiers": {
"SIMPLE": ["dead-vision", "live-text"],
"MEDIUM": "mid",
"COMPLEX": "big",
"REASONING": "top",
},
"session_affinity": True,
"modality_routing": True,
},
{"dead-vision": ["id-v1"], "live-text": ["id-t1"]},
cooling=("id-v1",),
)
key = router._get_session_affinity_cache_key("sess-image", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-vision", "tier": "SIMPLE"}, ttl=600
)
image_message = [
{
"role": "user",
"content": [
{"type": "text", "text": "What color is this?"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}},
],
}
]
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-image"}}, messages=image_message
)
assert result.model == "dead-vision"
@pytest.mark.asyncio
async def test_failover_will_not_pick_a_peer_that_cannot_hold_the_prompt(self):
"""The context-window filter is a pre-call check inside the eligibility owner, so this
drives the REAL owner on a real Router and injects only the cooldown. A substitute the
prompt overflows must never be chosen while a peer that holds it exists."""
pool = ["dead-big", "live-small", "live-big"]
router_instance = _windowed_router(
("dead-big", "openai/gpt-4o-mini", 200000),
("live-small", "openai/gpt-3.5-turbo", 16385),
("live-big", "openai/gpt-4o-mini", 200000),
)
router_instance.enable_pre_call_checks = True
dead_ids = {d["model_info"]["id"] for d in router_instance.model_list if d["model_name"] == "dead-big"}
async def active_cooldowns(model_ids, parent_otel_span):
return [(i, {"exception_received": "boom"}) for i in model_ids if i in dead_ids]
router_instance.cooldown_cache.async_get_active_cooldowns = active_cooldowns
router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="health-window-router",
litellm_router_instance=router_instance,
complexity_router_config={
"tiers": {name: list(pool) for name in ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING")},
"session_affinity": True,
"enable_context_window_escalation": True,
},
)
key = router._get_session_affinity_cache_key("sess-window", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-big", "tier": "SIMPLE"}, ttl=600
)
results = [
await router.async_pre_routing_hook(
model="m",
request_kwargs={"metadata": {"session_id": "sess-window"}},
messages=list(_OVERSIZED_TURNS),
)
for _ in range(20)
]
assert "live-small" not in {r.model for r in results}
assert {r.model for r in results} == {"live-big"}
@pytest.mark.asyncio
async def test_a_decision_with_no_tier_is_left_alone(self, mock_router_instance):
"""default_model placements carry no tier, so there is no pool to draw a peer from.
The gate leaves them exactly as they are rather than inventing a tier."""
router = self._router(
mock_router_instance,
{
"tiers": dict(self.TIERS),
"default_model": "fallback-model",
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"classifier_fallback": "default_model",
},
{"fallback-model": ["id-f1"], "dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-f1", "id-a1"),
)
mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier down"))
result = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.SIMPLE_MESSAGE)
assert result.model == "fallback-model"
assert result.routing_decision.get("tier") is None
assert result.routing_decision["cause"] != "health_failover"
@pytest.mark.asyncio
async def test_a_tier_entry_the_router_cannot_serve_fails_over_instead_of_erroring(self, mock_router_instance):
"""A tier naming a model this proxy has no deployment for is unservable, and the
eligibility owner says so, so the peer serves rather than the request 429ing."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"live-b": ["id-b1"]},
)
key = router._get_session_affinity_cache_key("sess-unknown", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-unknown"}}, messages=self.SIMPLE_MESSAGE
)
assert result.model == "live-b"
assert result.routing_decision["cause"] == "health_failover"
@pytest.mark.asyncio
async def test_a_peer_excluded_by_a_request_scoped_policy_is_not_a_failover_target(self, mock_router_instance):
"""Tag, team and access-group filters are request-scoped and live inside the eligibility
owner. A peer they exclude would be rejected downstream, so it must not be chosen."""
router = self._router(
mock_router_instance,
{
"tiers": {
"SIMPLE": ["dead-a", "tagged-out-b", "live-c"],
"MEDIUM": "mid",
"COMPLEX": "big",
"REASONING": "top",
},
"session_affinity": True,
},
{"dead-a": ["id-a1"], "tagged-out-b": ["id-b1"], "live-c": ["id-c1"]},
cooling=("id-a1",),
excluded=("id-b1",),
)
key = router._get_session_affinity_cache_key("sess-tagged", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-tagged"}}, messages=self.SIMPLE_MESSAGE
)
for _ in range(20)
]
assert {r.model for r in results} == {"live-c"}
@pytest.mark.asyncio
async def test_the_eligibility_probe_never_mutates_the_caller_request_kwargs(self, mock_router_instance):
"""The owner pops routing bookkeeping off the dict it is handed, so a probe that passed
the real kwargs would strip them before the request is ever placed."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-kwargs", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
request_kwargs = {
"metadata": {"session_id": "sess-kwargs"},
"_target_order": 1,
"_excluded_deployment_ids": ["id-x"],
}
result = await router.async_pre_routing_hook(
model="m", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert result.model == "live-b"
assert request_kwargs["_target_order"] == 1
assert request_kwargs["_excluded_deployment_ids"] == ["id-x"]
assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs)
@pytest.mark.asyncio
async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(
self, mock_router_instance
):
"""RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer
in that state would be rejected downstream, so it cannot be the substitute."""
from litellm.types.router import RouterRateLimitErrorBasic
router = self._router(
mock_router_instance,
{
"tiers": {
"SIMPLE": ["dead-a", "rpm-full-b", "live-c"],
"MEDIUM": "mid",
"COMPLEX": "big",
"REASONING": "top",
},
"session_affinity": True,
},
{"dead-a": ["id-a1"], "rpm-full-b": ["id-b1"], "live-c": ["id-c1"]},
cooling=("id-a1",),
raises_for={"rpm-full-b": RouterRateLimitErrorBasic(model="rpm-full-b")},
)
key = router._get_session_affinity_cache_key("sess-rpm", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-rpm"}}, messages=self.SIMPLE_MESSAGE
)
for _ in range(20)
]
assert {r.model for r in results} == {"live-c"}
@pytest.mark.asyncio
async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(
self, mock_router_instance
):
"""The Responses API carries its prompt as `input`, never as messages. The owner only
runs its context-window pre-call check when one of them is present, so dropping `input`
would silently skip window filtering on that whole surface."""
router = self._router(
mock_router_instance,
{"tiers": dict(self.TIERS), "session_affinity": True},
{"dead-a": ["id-a1"], "live-b": ["id-b1"]},
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-input", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
result = await router.async_pre_routing_hook(
model="m",
request_kwargs={"metadata": {"session_id": "sess-input"}},
input="summarize this document for me",
)
assert result.model == "live-b"
assert any(
probed_input == "summarize this document for me"
for _, probed_input in router.litellm_router_instance.probed_prompts
), "the eligibility probe must forward `input` to the owner"
@pytest.mark.asyncio
async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(
self, mock_router_instance
):
"""The owner answers an unconfigured group with BadRequestError. Reading that as live
would both skip failover off it and let it be chosen as a substitute."""
router = self._router(
mock_router_instance,
{
"tiers": {
"SIMPLE": ["dead-a", "unconfigured-b", "live-c"],
"MEDIUM": "mid",
"COMPLEX": "big",
"REASONING": "top",
},
"session_affinity": True,
},
{"dead-a": ["id-a1"], "live-c": ["id-c1"]},
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-missing", {})
await router.litellm_router_instance.cache.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-missing"}}, messages=self.SIMPLE_MESSAGE
)
for _ in range(20)
]
assert {r.model for r in results} == {"live-c"}

155
tests/test_litellm/test_baseten_glm_5_3_model_metadata.py Normal file
View file

@ -0,0 +1,155 @@
import json
from pathlib import Path
import pytest
import litellm
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
from litellm.utils import supports_function_calling, supports_prompt_caching
REPO_ROOT = Path(__file__).parents[2]
MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json"
BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json"
MODEL = "baseten/zai-org/GLM-5.3"
INPUT_COST = 1.4e-06
CACHED_INPUT_COST = 1.4e-07
OUTPUT_COST = 4.4e-06
def _load(path):
with open(path) as f:
return json.load(f)
@pytest.fixture
def local_model_cost_map(monkeypatch):
"""Force get_model_info to resolve against the in-repo cost map instead of the
remote one fetched at import time, which still carries the pre-merge registry."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
litellm.get_model_info.cache_clear()
yield
litellm.get_model_info.cache_clear()
def test_baseten_glm_5_3_specs():
info = _load(MAIN_PATH).get(MODEL)
assert info is not None, f"{MODEL} missing from model_prices_and_context_window.json"
assert info["litellm_provider"] == "baseten"
assert info["mode"] == "chat"
assert info["input_cost_per_token"] == INPUT_COST
assert info["output_cost_per_token"] == OUTPUT_COST
assert info["cache_read_input_token_cost"] == CACHED_INPUT_COST
assert info["max_input_tokens"] == 1048576
assert info["max_output_tokens"] == 262144
assert info["max_tokens"] == 262144
assert info["supports_function_calling"] is True
assert info["supports_prompt_caching"] is True
assert info["supports_response_schema"] is True
assert info["supports_tool_choice"] is True
assert info["supports_vision"] is True
assert info["supported_modalities"] == ["text", "image"]
assert info["supported_output_modalities"] == ["text"]
routed_model, provider, _, _ = get_llm_provider(model=MODEL)
assert routed_model == "zai-org/GLM-5.3"
assert provider == "baseten"
def test_baseten_glm_5_3_capabilities_are_visible_to_callers(local_model_cost_map):
"""The entry advertises prompt caching and tool calling, so the helpers every
caller checks before sending a request must say so too."""
assert supports_prompt_caching(model=MODEL) is True
assert supports_function_calling(model=MODEL) is True
info = litellm.get_model_info(model="zai-org/GLM-5.3", custom_llm_provider="baseten")
assert info["max_input_tokens"] == 1048576
assert info["max_output_tokens"] == 262144
def test_cached_prompt_tokens_bill_at_the_cached_rate(local_model_cost_map):
"""A cache hit reports its reused tokens under prompt_tokens_details, and those
tokens cost a tenth of the input rate, not the full rate and not nothing."""
usage = Usage(
prompt_tokens=21010,
completion_tokens=100,
total_tokens=21110,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=20992),
)
prompt_cost, completion_cost = litellm.cost_per_token(
model=MODEL, usage_object=usage, custom_llm_provider="baseten"
)
assert prompt_cost == pytest.approx(18 * INPUT_COST + 20992 * CACHED_INPUT_COST)
assert completion_cost == pytest.approx(100 * OUTPUT_COST)
def test_backup_matches_main():
"""Ensure the bundled (backup) cost map stays in sync with the canonical file.
Both keys are asserted present first: comparing two ``.get`` results alone passes
just as happily when neither file has the entry at all, which is the exact state
this test exists to catch.
"""
main_cost = _load(MAIN_PATH)
backup_cost = _load(BACKUP_PATH)
assert MODEL in main_cost, f"{MODEL} missing from model_prices_and_context_window.json"
assert MODEL in backup_cost, f"{MODEL} missing from model_prices_and_context_window_backup.json"
assert backup_cost[MODEL] == main_cost[MODEL], f"{MODEL} differs between main and backup model cost maps"
def test_entry_advertises_only_what_the_baseten_path_accepts(local_model_cost_map):
"""The entry must not claim a capability whose request parameter BasetenConfig
refuses.
``BasetenConfig.get_supported_openai_params`` returns one hardcoded list for every
Baseten model, and it carries neither ``parallel_tool_calls`` nor
``reasoning_effort``. Baseten's own Model API does take ``reasoning_effort``, but
litellm's Baseten path drops it (``drop_params=True``) or raises
``UnsupportedParamsError`` (``drop_params=False``), so declaring
``supports_parallel_function_calling``, ``supports_reasoning`` or
``reasoning_effort_levels`` here would advertise a level the gateway then refuses to
send. Wiring those params through the Baseten config is separate work; until it
lands, the registry stays honest.
"""
supported = litellm.get_supported_openai_params(model="zai-org/GLM-5.3", custom_llm_provider="baseten")
assert supported is not None
entry = _load(MAIN_PATH)[MODEL]
capability_to_param = {
"supports_function_calling": "tools",
"supports_tool_choice": "tool_choice",
"supports_response_schema": "response_format",
"supports_parallel_function_calling": "parallel_tool_calls",
"supports_reasoning": "reasoning_effort",
}
for capability, param in capability_to_param.items():
if entry.get(capability):
assert param in supported, f"{MODEL} advertises {capability} but baseten drops/rejects {param}"
assert "reasoning_effort_levels" not in entry, (
"reasoning_effort_levels advertises accepted reasoning_effort values, which the Baseten path does not accept"
)
assert "thinking_always_on" not in entry, (
"thinking_always_on is only read by AnthropicModelInfo._is_always_on_thinking_model, "
"which no Baseten route reaches"
)
with pytest.raises(litellm.UnsupportedParamsError):
litellm.utils.get_optional_params(
model="zai-org/GLM-5.3",
custom_llm_provider="baseten",
parallel_tool_calls=True,
reasoning_effort="high",
drop_params=False,
)

5
tests/test_litellm/test_cost_calculator.py
View file

@ -175,6 +175,11 @@ def test_wandb_model_api_pricing_entries(_local_model_cost_map):
expected_pricing = {
"wandb/moonshotai/Kimi-K2.5": (6e-07, 3e-06),
"wandb/MiniMaxAI/MiniMax-M2.5": (3e-07, 1.2e-06),
"wandb/Qwen/Qwen3-235B-A22B-Instruct-2507": (1e-07, 1e-07),
"wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": (1e-07, 1e-07),
"wandb/deepseek-ai/DeepSeek-R1-0528": (1.35e-06, 5.4e-06),
"wandb/deepseek-ai/DeepSeek-V3-0324": (1.14e-06, 2.75e-06),
"wandb/meta-llama/Llama-4-Scout-17B-16E-Instruct": (1.7e-07, 6.6e-07),
}
for model_name, (input_cost, output_cost) in expected_pricing.items():

Some files were not shown because too many files have changed in this diff Show more