Merge branch 'main' into litellm_tts_true_streaming

This commit is contained in:
mateo-berri 2026-11-04 12:04:45 -08:00
commit 5fa3c58bac
98 changed files with 2528 additions and 7319 deletions

View file

@ -280,6 +280,7 @@ jobs:
tests/unit/llms/e2b
tests/unit/llms/edenai
tests/unit/llms/elevenlabs
tests/unit/llms/empiriolabs
tests/unit/llms/exa_ai
tests/unit/llms/fal_ai
tests/unit/llms/fastcrw

View file

@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | |
| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | |
| [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | |
| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | |
| [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | |
| [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | |
| [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | |

View file

@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
unpriced: Final = _without_read_price(billed.response)
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
end_time: Final = datetime.now()
logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time)
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time

View file

@ -1796,6 +1796,9 @@ if TYPE_CHECKING:
from .llms.bespoke.decisions.transformation import (
BespokeDecisionsConfig as BespokeDecisionsConfig,
)
from .llms.empiriolabs.decisions.transformation import (
EmpirioLabsDecisionsConfig as EmpirioLabsDecisionsConfig,
)
from .llms.eu_jev.decisions.transformation import (
EuJevDecisionsConfig as EuJevDecisionsConfig,
)

View file

@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = (
"OllamaDecisionsConfig",
"LayaDecisionsConfig",
"BespokeDecisionsConfig",
"EmpirioLabsDecisionsConfig",
"EuJevDecisionsConfig",
"OpenAIDecisionsConfig",
"AzureDecisionsConfig",
@ -759,6 +760,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"),
"LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"),
"BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"),
"EmpirioLabsDecisionsConfig": (".llms.empiriolabs.decisions.transformation", "EmpirioLabsDecisionsConfig"),
"EuJevDecisionsConfig": (".llms.eu_jev.decisions.transformation", "EuJevDecisionsConfig"),
"OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
"AzureDecisionsConfig": (".llms.azure.decisions.transformation", "AzureDecisionsConfig"),

View file

@ -1,7 +1,7 @@
"""Public accounting contracts. Importing these values does not enable accounting."""
from litellm.accounting.measurements import measurements_for
from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter
from litellm.accounting.protocols import CostAggregator
from litellm.accounting.types import (
AdmissionDecision,
AggregateQuery,
@ -33,9 +33,7 @@ __all__ = [
"AdmissionDecision",
"AggregateQuery",
"AggregateRow",
"BudgetEnforcer",
"CostAggregator",
"CostCompiler",
"CostEvent",
"DeliveryState",
"DeliveryTarget",
@ -56,7 +54,6 @@ __all__ = [
"PricingSnapshot",
"Rate",
"Unpriced",
"UsageMeter",
"measurements_for",
"validate_aggregate_rows",
]

View file

@ -0,0 +1,304 @@
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.constants import (
BILLED_RETRY_ATTEMPT_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MODEL_INSIGHTS_DEFAULT_TASK,
MODEL_INSIGHTS_TASK_TAG_PREFIX,
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids
from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks
from litellm.proxy.db.model_usage_rollup import date_from_start_time
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
_FIELDS: Final = TypeAdapter(dict[str, object])
_STRINGS: Final = TypeAdapter(tuple[str, ...])
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
_TASK_TYPES: Final = frozenset(load_model_insight_tasks())
@dataclass(frozen=True, slots=True)
class Contribution:
table: str
key: tuple[tuple[str, str], ...]
increments: tuple[tuple[str, int | float], ...]
values: tuple[tuple[str, str | None], ...] = ()
def _fields(value: object) -> Mapping[str, object]:
if isinstance(value, str):
return _FIELDS.validate_json(value)
return _FIELDS.validate_python(value) if isinstance(value, Mapping) else MappingProxyType({})
def _text(value: object) -> str:
return "" if value is None else str(value)
def _number(value: object) -> int | float:
return value if isinstance(value, (int, float)) else 0
def _tags(row: Mapping[str, object]) -> tuple[str, ...]:
value: Final = row.get("request_tags")
if value is None:
return ()
return _STRINGS.validate_json(value) if isinstance(value, str) else _STRINGS.validate_python(value)
def _date(row: Mapping[str, object]) -> str | None:
start: Final = row.get("startTime")
return date_from_start_time(start) if isinstance(start, (datetime, str)) else None
def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contribution:
return Contribution(table, key, (("spend", spend),))
def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
spend: Final = _number(row.get("spend"))
entities: Final = _fields(metadata.get("accounting_entities"))
identities: Final = (
("LiteLLM_VerificationToken", "token", entities.get("api_key")),
("LiteLLM_UserTable", "user_id", entities.get("user")),
("LiteLLM_UserTable", "user_id", entities.get("proxy_budget_name")),
("LiteLLM_TeamTable", "team_id", entities.get("team_id")),
("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")),
("LiteLLM_EndUserTable", "user_id", entities.get("end_user")),
("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")),
("LiteLLM_ProjectTable", "project_id", entities.get("project_id")),
)
groups: Final = _STRINGS.validate_python(metadata.get("accounting_model_access_groups", ()))
return (
*(
_spend(table, ((column, identity),), spend)
for table, column, identity in identities
if isinstance(identity, str)
),
*(_spend("LiteLLM_TagTable", (("tag_name", tag),), spend) for tag in _tags(row) if tag),
*(_spend("LiteLLM_ModelAccessGroupBudgetTable", (("access_group_name", group),), spend) for group in groups),
*_membership_contributions(entities, spend),
)
def _membership_contributions(row: Mapping[str, object], spend: float) -> tuple[Contribution, ...]:
user: Final = row.get("user")
if not isinstance(user, str):
return ()
owners: Final = (
("LiteLLM_TeamMembership", "team_id", row.get("team_id")),
("LiteLLM_OrganizationMembership", "organization_id", row.get("organization_id")),
)
return tuple(
_spend(table, ((column, identity), ("user_id", user)), spend)
for table, column, identity in owners
if isinstance(identity, str)
)
def _daily_counters(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[tuple[str, int | float], ...]:
usage: Final = _fields(metadata.get("usage_object"))
savings: Final = _fields(metadata.get("accounting_savings"))
internal: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY))
caller: Final = not internal and metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is not True
success: Final = metadata.get("status") != "failure"
duration: Final = row.get("request_duration_ms")
timed: Final = not internal and success and isinstance(duration, int) and duration >= 0
return (
("spend", _number(row.get("spend"))),
("prompt_tokens", _number(row.get("prompt_tokens"))),
("completion_tokens", _number(row.get("completion_tokens"))),
("api_requests", int(caller)),
("successful_requests", int(caller and success)),
("failed_requests", int(caller and not success)),
("cache_read_input_tokens", extract_cache_read_tokens(usage)),
("cache_creation_input_tokens", extract_cache_creation_tokens(usage)),
("compression_saved_tokens", extract_compression_saved_tokens(metadata)),
("compression_savings_spend", _number(savings.get("compression"))),
("prompt_caching_savings_spend", _number(savings.get("prompt_caching"))),
("gateway_injected_caching_savings_spend", _number(savings.get("gateway_injected_caching"))),
("autorouter_savings_spend", 0.0 if internal else _number(savings.get("autorouter"))),
("total_response_time_ms", _number(duration) if timed else 0),
("timed_requests", int(timed)),
)
def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
date: Final = _date(row)
if date is None or "api_key" not in row:
return ()
if "mcp_namespaced_tool_name" not in row and (
"model" not in row or "custom_llm_provider" not in row or "model_group" not in row
):
return ()
counters: Final = _daily_counters(row, metadata)
endpoint: Final = ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type")))
shared_key: Final = (
("date", date),
("api_key", _text(row.get("api_key"))),
("model", _text(row.get("model"))),
("custom_llm_provider", _text(row.get("custom_llm_provider"))),
("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))),
("endpoint", _text(endpoint)),
)
# The legacy daily organization writer uses the charged org argument, not
# SpendLogs' normalized empty-string identity for a request with no org.
daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")}
return tuple(_daily_rows(daily_row, shared_key, counters))
def _daily_rows(
row: Mapping[str, object],
shared_key: tuple[tuple[str, str], ...],
counters: tuple[tuple[str, int | float], ...],
) -> Iterator[Contribution]:
for entity, table in DAILY_SPEND_TABLES.items():
for identity in daily_spend_entity_ids(row, entity):
yield Contribution(
table.name,
((table.entity_id_column, _text(identity)), *shared_key),
counters,
(
("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))),
*((("request_id", _text(row.get("request_id"))),) if table.carries_request_id else ()),
),
)
def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
windows: Final = _ROWS.validate_python(metadata.get("accounting_windows", ()))
return tuple(
Contribution(
"LiteLLM_BudgetWindowSpend",
tuple((column, _text(window[column])) for column in ("entity_type", "entity_id", "window_duration")),
(("spend", _number(window.get("spend", row.get("spend")))),),
(("window_start", _text(window["window_start"])), ("started_at", _text(row.get("startTime")))),
)
for window in windows
)
def project(row: Mapping[str, object]) -> tuple[Contribution, ...]:
metadata: Final = _fields(row.get("metadata"))
return (
*_entity_contributions(row, metadata),
*_daily_contributions(row, metadata),
*_window_contributions(row, metadata),
*_model_contributions(row, metadata),
*_guardrail_contributions(row, metadata),
)
def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
date: Final = _date(row)
origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)
if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN):
return ()
model: Final = _text(row.get("model")) or "unknown"
caller: Final = not origin
success: Final = row.get("status") == "success"
return (
Contribution(
"LiteLLM_DailyModelUsage",
(
("date", date),
("model_group", _text(row.get("model_group")) or model),
("model", model),
("custom_llm_provider", _text(row.get("custom_llm_provider")) or "unknown"),
("task_type", _model_task(row)),
),
(
("spend", _number(row.get("spend"))),
("prompt_tokens", _number(row.get("prompt_tokens"))),
("completion_tokens", _number(row.get("completion_tokens"))),
("request_count", int(caller)),
("successful_requests", int(caller and success)),
("failed_requests", int(caller and not success)),
),
),
)
def _utc_date(row: Mapping[str, object]) -> str | None:
start: Final = row.get("startTime")
try:
parsed: Final = datetime.fromisoformat(start.replace("Z", "+00:00")) if isinstance(start, str) else start
except ValueError:
return None
if not isinstance(parsed, datetime):
return None
return (
(parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed)
.astimezone(timezone.utc)
.date()
.isoformat()
)
def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
date: Final = _utc_date(row)
info: Final = metadata.get("guardrail_information") or metadata.get("standard_logging_guardrail_information")
if not row.get("request_id") or date is None or not isinstance(info, list):
return ()
return tuple(chain.from_iterable(_guardrail_entry(row, date, entry) for entry in _ROWS.validate_python(info)))
def _guardrail_entry(row: Mapping[str, object], date: str, entry: Mapping[str, object]) -> tuple[Contribution, ...]:
identity: Final = _text(entry.get("guardrail_id") or entry.get("guardrail_name"))
usage: Final = entry.get("guardrail_usage")
if not identity or not isinstance(usage, Mapping):
return ()
costs: Final = _guardrail_unit_costs(entry)
return tuple(
_guardrail_unit(row, (identity, date, unit), units, None if costs is None else costs.get(unit))
for unit, units in _fields(usage).items()
if isinstance(units, int) and not isinstance(units, bool) and units > 0
)
def _guardrail_unit(
row: Mapping[str, object], identity: tuple[str, str, str], units: int, cost: float | None
) -> Contribution:
guardrail_id, date, unit = identity
return Contribution(
"LiteLLM_DailyGuardrailUsageUnits",
(
("guardrail_id", guardrail_id),
("date", date),
("team_id", _text(row.get("team_id"))),
("api_key", _text(row.get("api_key"))),
("usage_unit", unit),
),
(("units", units), ("cost", cost or 0.0), ("untracked_units", units if cost is None else 0)),
)
def _model_task(row: Mapping[str, object]) -> str:
return next(
(
task
for tag in _tags(row)
if tag.startswith(MODEL_INSIGHTS_TASK_TAG_PREFIX)
and (task := tag.removeprefix(MODEL_INSIGHTS_TASK_TAG_PREFIX)) in _TASK_TYPES
),
MODEL_INSIGHTS_DEFAULT_TASK,
)
def _guardrail_unit_costs(entry: Mapping[str, object]) -> Mapping[str, float | None] | None:
try:
cost: Final = GuardrailCostByUnitEntry.model_validate(entry)
except ValidationError:
return None
return None if cost.guardrail_cost_in_spend is False else cost.guardrail_cost_by_unit

View file

@ -1,58 +1,14 @@
"""Contracts only. Implementations own persistence, pricing, projections and live budgets.
The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``.
Usage envelopes retain the public event and request-scoped pricing context: quantities with units,
model/provider, requested and billed tier/region when known, and per-request/deployment overrides.
These inputs, including the effective cost-map rates, must survive request teardown.
Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot.
T5/T6 own envelope serialization; this public API does not introduce a second persisted schema.
"""
from __future__ import annotations
from collections.abc import Mapping
from typing import Protocol
from litellm.accounting.types import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
Dimension,
InboundUsageEvent,
Measurement,
Metric,
OutboundUsageEvent,
Unpriced,
)
class UsageMeter(Protocol):
async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None:
"""Acknowledge commit, no-op identical IDs and reject conflicting content.
Persist the event and its request-scoped pricing context before acknowledging.
HTTP acknowledgement is a separate policy. Unknown usage remains None.
"""
...
class CostCompiler(Protocol):
async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced:
"""Reuse the committed charge or finalize one charge and its delivery intent.
Resolve captured request-scoped context and rates by event identity.
Persist the resolved inputs in the required snapshot. Replay
reuses the committed result and never consumes free allowance again.
Missing pricing is Unpriced, never an inferred zero.
"""
...
from litellm.accounting.types import AggregateQuery, AggregateRow
class CostAggregator(Protocol):
async def apply(self, measurements: tuple[Measurement, ...]) -> None:
"""Apply one source event's complete contribution set per target/version.
Each target commits all contributions and its receipt atomically. A receipt
for the first metric must not suppress the rest. Usage need not wait for cost.
"""
async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None:
"""Apply stored request contributions atomically with their delivery receipts."""
...
async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]:
@ -63,23 +19,3 @@ class CostAggregator(Protocol):
combinations. Different groups may overlap, including non-additive tags.
"""
...
class BudgetEnforcer(Protocol):
async def reserve(
self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...]
) -> AdmissionDecision:
"""Apply existing admission policies and estimates; duplicate IDs add no hold."""
...
async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None:
"""Settle the committed charge with atomic deduplication and hold mutation.
Preserve scope/window attribution and retry unfinished settlement after
restart. SQL projection receipts do not provide live-counter atomicity.
"""
...
async def release(self, reservation_id: str) -> None:
"""Idempotently release unused exposure at the request boundary, preserving charges."""
...

View file

@ -0,0 +1,97 @@
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection
if TYPE_CHECKING:
from litellm.router import Router
_FIELDS: Final = TypeAdapter(dict[str, object])
def _mapping(value: object) -> Mapping[str, object] | None:
return _FIELDS.validate_python(value) if isinstance(value, Mapping) else None
_METADATA_FIELDS: Final = frozenset(
{
"usage_object",
"status",
"guardrail_information",
"standard_logging_guardrail_information",
"compression_savings",
"internal_call_origin",
"billed_retry_attempt",
}
)
_ACCOUNTING_FIELDS: Final = frozenset(
{
"request_id",
"call_type",
"api_key",
"spend",
"prompt_tokens",
"completion_tokens",
"startTime",
"model",
"model_group",
"custom_llm_provider",
"user",
"team_id",
"organization_id",
"end_user",
"agent_id",
"request_tags",
"status",
"request_duration_ms",
"mcp_namespaced_tool_name",
}
)
def build_accounting_payload(
payload: SpendLogsPayload,
*,
entities: Mapping[str, str | None],
model_access_groups: Sequence[str],
windows: Sequence[Mapping[str, object]],
router: "Callable[[], Router | None]",
) -> Mapping[str, object]:
metadata: Final = _FIELDS.validate_json(payload["metadata"])
savings: Final = compute_savings_spend(
model=payload["model"],
custom_llm_provider=payload["custom_llm_provider"],
compression_saved_tokens=extract_compression_saved_tokens(metadata),
gateway_injected_cache=marks_gateway_injection(metadata, payload["model_id"]),
routing_decision=_mapping(metadata.get("routing_decision")),
model_id=payload["model_id"],
llm_router=router,
usage_object=_mapping(metadata.get("usage_object")),
cost_breakdown=_mapping(metadata.get("cost_breakdown")),
recorded_autorouter_savings=metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")),
billed_at=payload["endTime"],
)
accounting_metadata: Final = {
**{key: value for key, value in metadata.items() if key in _METADATA_FIELDS},
"accounting_entities": entities,
"accounting_savings": savings._asdict(),
"accounting_model_access_groups": model_access_groups,
"accounting_windows": windows,
}
return {
**{
key: value.isoformat() if isinstance(value, datetime) else value
for key, value in payload.items()
if key in _ACCOUNTING_FIELDS
},
"metadata": _FIELDS.validate_python(safe_json_structure(accounting_metadata, non_finite_as_null=True)),
}

View file

@ -50,6 +50,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
get_batch_cost_rates,
get_billable_input_tokens,
get_cost_per_unit,
get_fedramp_uplift_multiplier,
get_modality_token_cost_key,
get_regional_uplift_multiplier,
get_service_tier_cost_key,
@ -3311,7 +3312,9 @@ def batch_cost_calculator(
text_tokens: Final = usage.completion_tokens - image_tokens
total_completion_cost = text_tokens * text_rate + image_tokens * image_rate
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency)
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_fedramp_uplift_multiplier(
model_info, data_residency
)
if uplift != 1.0:
total_prompt_cost *= uplift
total_completion_cost *= uplift

View file

@ -19,7 +19,7 @@ from datetime import datetime as dt_object
from enum import Enum
from functools import lru_cache
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast
from httpx import Response, Timeout
from openai import AsyncStream as OpenAIAsyncStream
@ -283,8 +283,6 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector
from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation
from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
EnterpriseCallbackControls,
@ -662,54 +660,7 @@ class UsageObservation(Enum):
ESTIMATED = "estimated"
@dataclass(frozen=True, slots=True)
class UsageCaptureObservation:
usage_json: str | None
provenance: UsageObservation
response_cost: float | None = None
@classmethod
def from_usage(
cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None
) -> "UsageCaptureObservation | UsageCaptureError":
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
try:
serialized: Final = (
None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json()
)
except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling
return UsageCaptureError(str(error))
else:
return cls(serialized, provenance, response_cost)
class Logging(LiteLLMLoggingBaseClass):
usage_capture_callback: ClassVar[
Callable[
["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | None"
]
| None
] = None
def capture_usage_before_logging(
self,
response: object,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
*,
freeze_only: bool = False,
) -> None:
capture: Final = type(self).usage_capture_callback
if capture is not None and (not self.stream or self._is_assembled_stream_success(response)):
capture(
self,
response,
start_time or self.start_time,
end_time or datetime.datetime.now(datetime.timezone.utc),
freeze_only,
)
global \
supabaseClient, \
promptLayerLogger, \
@ -949,7 +900,6 @@ class Logging(LiteLLMLoggingBaseClass):
self._stash_owed_on_success: bool = False
self._stash_billed_by_continuation: bool = False
self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = ()
self._failed_attempt_ids_owed_on_success: tuple[str, ...] = ()
self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = ()
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
@ -2582,11 +2532,6 @@ class Logging(LiteLLMLoggingBaseClass):
)
return
self.capture_usage_before_logging(
TypeAdapter(object).validate_python(result),
TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(start_time),
TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(end_time),
)
await self.async_success_handler(
result,
start_time=start_time,
@ -2698,20 +2643,12 @@ class Logging(LiteLLMLoggingBaseClass):
self.finalized_cost = finalized
self.model_call_details["response_cost"] = finalized.amount
def record_usage_capture(self, captured: "CapturedUsage") -> None:
self.model_call_details["_accounting_capture"] = captured
def record_usage_pricing(self, pricing: object) -> None:
self.model_call_details["_accounting_pricing_snapshot"] = pricing
def record_partial_usage_for_failure(
self, usage: Usage, response_cost: float, *, observation: UsageObservation = UsageObservation.KNOWN
) -> 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
if type(self).usage_capture_callback is not None:
self.model_call_details["_accounting_usage_observed"] = observation
def record_failure_after_earlier_rounds(
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
@ -2795,12 +2732,6 @@ class Logging(LiteLLMLoggingBaseClass):
owed_on_success: Final = self._stash_owed_on_success
self._stash_billed_by_continuation = False
self._stash_owed_on_success = False
captured_attempt_ids: Final = (
self.capture_failed_attempt_usage(usage)
if isinstance(usage, Usage) and type(self).usage_capture_callback is not None
else ()
)
self.model_call_details.pop("_accounting_capture", None)
self.model_call_details.pop("_accounting_pricing_snapshot", None)
self.model_call_details.pop("_accounting_usage_observed", None)
self.model_call_details.pop("combined_usage_object", None)
@ -2816,60 +2747,10 @@ class Logging(LiteLLMLoggingBaseClass):
if not owed_on_success or billed_by_continuation:
return
self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt)
self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids)
def mark_stash_billed_by_continuation(self) -> None:
self._stash_billed_by_continuation = True
def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]:
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
capture: Final = type(self).usage_capture_callback
if capture is None:
return ()
attempt: Final = self.attempt_scoped_copy()
attempt.prepared_cost = None
attempt.finalized_cost = None
identity: Final = (
f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}"
)
params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {})
metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params)
observed: Final[object] = self.model_call_details.get("_accounting_usage_observed")
cost: Final[object] = self.model_call_details.get("response_cost")
attempt.model_call_details = {
**attempt.model_call_details,
"litellm_call_id": identity,
"parent_request_id": self.litellm_call_id,
"_accounting_capture": None,
"_accounting_prior_captures": (),
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
usage,
observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN,
float(cost) if isinstance(cost, (int, float)) else None,
),
"litellm_params": {
**params,
"metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"},
"litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"},
},
}
captured: Final = capture(
attempt,
{"id": identity, "usage": usage.model_dump(), "status": "failed"},
self.start_time,
datetime.datetime.now(datetime.timezone.utc),
False,
)
if not isinstance(captured, CapturedUsage):
return ()
self.model_call_details["_accounting_prior_captures"] = (
*self.model_call_details.get("_accounting_prior_captures", ()),
captured,
)
return tuple(write.event.event_id for write in captured.writes)
def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None:
"""Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw."""
self.model_call_details["continued_partial_usage"] = usage
@ -2909,8 +2790,6 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details.pop("failed_attempts_usage", None)
self.failed_attempts_billed_on_success = failed_attempts
self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts
if self._failed_attempt_ids_owed_on_success:
self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success
self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects(
usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)]
)
@ -3073,45 +2952,6 @@ class Logging(LiteLLMLoggingBaseClass):
return self._handle_a2a_response_logging(result=result)
return result
def usage_capture_response(self, result: object) -> object:
if self.stream and not self._is_assembled_stream_success(result):
return None
copied: Final = result.model_copy(deep=True) if isinstance(result, BaseModel) else result
normalized: Final = self.normalize_logging_result(self.normalize_endpoint_logging_result(copied))
if self.model_call_details.get("call_type") == CallTypes.arealtime.value and isinstance(result, list):
frames: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(result)
measured: Final = any(
isinstance(_LITELLM_PARAMS.validate_python(frame.get("response") or {}).get("usage"), Mapping)
for frame in frames
if frame.get("type") == "response.done"
)
if not measured:
metered: Final = self.model_call_details.get(REALTIME_METERED_USAGE_KEY)
return {
"id": self.litellm_call_id,
"usage": {"metered_usage": TypeAdapter(object).dump_python(metered, mode="json")}
if metered is not None
else None,
}
assembled: Final = self._get_assembled_streaming_response(
normalized, self.start_time, self.start_time, True, []
)
return normalized if assembled is None else assembled
def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]:
from litellm.utils import count_characters
return TypeAdapter(dict[str, JsonValue]).validate_python(
{
"custom_pricing": self.custom_pricing_for(result),
"custom_cost_per_token": self.custom_cost_per_token,
"custom_cost_per_second": self.custom_cost_per_second,
"prompt_characters": count_characters(self._prompt_for_cost_calculation())
if self.model_call_details.get("call_type") in (CallTypes.speech.value, CallTypes.aspeech.value)
else None,
}
)
def normalize_logging_result(self, result: object) -> object:
"""
Some endpoints return a different type of result than what is expected by the logging system.
@ -3826,7 +3666,6 @@ class Logging(LiteLLMLoggingBaseClass):
self._as_served_chat_response(complete_streaming_response)
)
self.model_call_details["complete_streaming_response"] = logged_streaming_response
self.capture_usage_before_logging(logged_streaming_response, start_time, end_time)
## STANDARDIZED LOGGING PAYLOAD
if not for_async_request:
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
@ -4352,8 +4191,6 @@ class Logging(LiteLLMLoggingBaseClass):
cost_for_built_in_tools_cost_usd_dollar=0.0,
)
self.capture_usage_before_logging(result, start_time, end_time)
self.truncated_messages_for_logging = await truncate_base64_in_messages_async(
StandardLoggingPayloadSetup.append_system_prompt_messages(
kwargs=self.model_call_details, messages=self.model_call_details.get("messages")

View file

@ -1441,6 +1441,29 @@ def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str |
_get_regional_uplift_multiplier = get_regional_uplift_multiplier
def get_fedramp_uplift_multiplier(model_info: ModelInfo, data_residency: str | None) -> float:
"""
Resolve the per-model uplift multiplier for the OpenAI FedRAMP endpoint.
OpenAI applies a flat percentage uplift (e.g. +10%) on all token costs for
requests served from gov.api.openai.com. The multiplier is stored on the
model entry as ``fedramp_uplift_multiplier`` (e.g. 1.10).
Returns 1.0 (no uplift) unless ``data_residency`` is ``"fedramp"`` and the
model has a multiplier configured.
"""
if data_residency is None or data_residency.lower() != DataResidency.FEDRAMP.value:
return 1.0
multiplier: Final = model_info.get("fedramp_uplift_multiplier")
if multiplier is None:
return 1.0
try:
return float(multiplier)
except (TypeError, ValueError):
verbose_logger.exception("Invalid fedramp_uplift_multiplier for model; defaulting to 1.0")
return 1.0
def get_vertex_regional_endpoint_uplift(model_info: ModelInfo, vertex_location: str | None) -> float:
"""
Resolve the per-model uplift multiplier for Vertex AI non-global (regional and
@ -1550,8 +1573,8 @@ def generic_cost_per_token(
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
used to apply the per-model regional-processing uplift multiplier.
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us") or
"fedramp", used to apply the per-model regional-processing or FedRAMP uplift multiplier.
- vertex_location: optional Vertex AI location the request was served from
(e.g. "us-east5", "global"), used to apply the per-model
regional-endpoint uplift multiplier when non-global.
@ -1728,10 +1751,9 @@ def generic_cost_per_token(
)
completion_cost += float(video_tokens) * _output_cost_per_video_token
## REGIONAL DATA-RESIDENCY UPLIFT
# Applied as a flat multiplier across all token costs for the request
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).
uplift: Final = get_regional_uplift_multiplier(resolved_model_info, data_residency)
uplift: Final = get_regional_uplift_multiplier(resolved_model_info, data_residency) * get_fedramp_uplift_multiplier(
resolved_model_info, data_residency
)
if uplift != 1.0:
prompt_cost *= uplift
completion_cost *= uplift
@ -1874,6 +1896,7 @@ def _cost_map_billed_rates(
)
multiplier: Final = (
get_regional_uplift_multiplier(model_info, data_residency)
* get_fedramp_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
* get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
* get_provider_specific_speed_multiplier(model_info=model_info, usage=usage)
@ -2059,8 +2082,10 @@ def calculate_prompt_caching_savings(
{ttl: rate - prompt_base_cost for ttl, rate in base_costs.cache_creation_rates_by_ttl.items()}
),
)
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(
model_info, vertex_location
uplift: Final = (
get_regional_uplift_multiplier(model_info, data_residency)
* get_fedramp_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
)
return (read_discount - write_premium) * uplift

View file

@ -836,8 +836,6 @@ class RealTimeStreaming:
if self.session_tools or self.tool_calls:
self.logging_obj.model_call_details["realtime_tools"] = self.session_tools
self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls
if LiteLLMLogging.usage_capture_callback is not None:
self.logging_obj.capture_usage_before_logging(self.messages)
# Route through the bounded logging worker (per-coroutine timeout +
# concurrency cap) instead of a bare create_task, so a slow callback
# can't leave suspended tasks pinning each call's response in memory.

View file

@ -272,15 +272,31 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[
return {
**fields,
**({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}),
**({"reasoning": REDACTED_BY_LITELLM} if fields.get("reasoning") is not None else {}),
**({"reasoning_text": REDACTED_BY_LITELLM} if fields.get("reasoning_text") is not None else {}),
**{
key: None
for key in (
"thinking_blocks",
"reasoningContentBlocks",
"reasoningContent",
"annotations",
"citations",
"citation",
"citationsContent",
"compaction_blocks",
"compaction_start",
"compaction_delta",
"reasoning_details",
"web_search_calls",
"web_search_results",
"tool_results",
"code_interpreter_results",
"search_results",
"rerank_results",
"executed_tools",
"custom_outputs",
"tool_search_items",
"token_ids",
)
if key in fields

View file

@ -0,0 +1,10 @@
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
class EmpirioLabsDecisionsConfig(BaseDecisionsConfig):
path = "/v1/decisions"
api_key_env = ("EMPIRIOLABS_API_KEY",)
api_base_env = ("EMPIRIOLABS_API_BASE",)
def get_default_api_base(self) -> str | None:
return "https://api.empiriolabs.ai"

View file

@ -31,9 +31,9 @@ def cost_per_token(
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
inferred from api_base. Applies the model's regional-processing
uplift multiplier when set.
- data_residency: optional OpenAI data-residency mode (e.g. "eu", "us", "fedramp"),
inferred from api_base, including gov.api.openai.com. Applies the
model's regional-processing uplift multiplier when set.
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd

View file

@ -7,9 +7,12 @@ enabled and rejects requests sent to the wrong host, so the api_base hostname
is the authoritative signal of which region a request was processed in.
"""
import os
from typing import Final
from urllib.parse import urlparse
import litellm
# Mapping of OpenAI regional hostnames to the corresponding data-residency
# value used by the cost calculator. See
# https://developers.openai.com/api/docs/pricing for the regional-processing
@ -17,6 +20,7 @@ from urllib.parse import urlparse
_OPENAI_REGIONAL_HOSTS: Final[dict[str, str]] = {
"eu.api.openai.com": "eu",
"us.api.openai.com": "us",
"gov.api.openai.com": "fedramp",
}
@ -25,13 +29,22 @@ def infer_openai_data_residency(custom_llm_provider: str | None, api_base: str |
Derive the OpenAI data-residency region from an api_base URL.
Returns ``"eu"`` for the EU regional host, ``"us"`` for the US regional
host, and ``None`` for the default global host, any non-OpenAI provider,
or any non-OpenAI URL.
host, ``"fedramp"`` for ``gov.api.openai.com``, and ``None`` for the
default global host, any non-OpenAI provider, or any non-OpenAI URL.
When ``api_base`` is unset, falls back to the same defaults the OpenAI
handler uses (``litellm.api_base``, ``OPENAI_BASE_URL``, ``OPENAI_API_BASE``)
so env-configured regional endpoints are billed with their uplift.
"""
if custom_llm_provider != "openai" or not api_base:
if custom_llm_provider != "openai":
return None
resolved_api_base: Final = (
api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE")
)
if not resolved_api_base:
return None
try:
host: Final = urlparse(api_base).hostname
host: Final = urlparse(resolved_api_base).hostname
except (TypeError, ValueError):
return None
if not host:

View file

@ -6111,6 +6111,11 @@ def completion(
messages if readable_messages is None else list(readable_messages)
)
supplied_client_api_base: Final = (
str(client.base_url)
if custom_llm_provider == "openai" and isinstance(client, openai.OpenAI | openai.AsyncOpenAI)
else None
)
# For logging - save the values of the litellm-specific params passed in
requested_litellm_params: Final = get_litellm_params(
acompletion=acompletion,
@ -6119,7 +6124,7 @@ def completion(
logger_fn=logger_fn,
verbose=verbose,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_base=supplied_client_api_base or api_base,
litellm_call_id=kwargs.get("litellm_call_id", None),
model_alias_map=litellm.model_alias_map,
completion_call_id=id,

File diff suppressed because it is too large Load diff

View file

@ -5277,7 +5277,7 @@
}
},
"empiriolabs": {
"display_name": "EmpirioLabs (`empiriolabs`)",
"display_name": "EmpirioLabs AI (`empiriolabs`)",
"url": "https://docs.litellm.ai/docs/providers/empiriolabs",
"endpoints": {
"chat_completions": false,

File diff suppressed because one or more lines are too long

View file

@ -38,7 +38,6 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_no_callback_env_reference,
)
from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from litellm.types.agents import (
AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export
)
@ -2481,22 +2480,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
Documents all the fields supported by `general_settings` in config.yaml
"""
accounting_capture: UsageCaptureSettings | None = Field(
None,
description="Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes.",
)
accounting_sync_ack: bool = Field(
False,
description="Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. "
"Other requests remain supported with asynchronous usage capture.",
)
accounting_sync_ack_timeout_seconds: float = Field(
20.0, gt=0, description="Maximum response wait for usage commit, including the batch flush interval"
)
accounting_sync_ack_failure_policy: Literal["fail_closed", "log"] = Field(
"fail_closed", description="Return 503 or log the persistence failure after the response wait expires"
)
request_tag_mappings: Sequence[RequestTagMapping] | None = Field(
None,
description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.",

View file

@ -140,7 +140,6 @@ from litellm.proxy.route_llm_request import (
route_request,
)
from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost
from litellm.proxy.spend_tracking.usage_capture import capture_proxy_response
from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id
from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
ProxyLogging,
@ -3397,9 +3396,6 @@ class ProxyBaseLLMRequestProcessing:
)
),
)
await capture_proxy_response(
logging_obj, response, background=_REQUEST_FIELDS.validate_python(self.data).get("background") is True
)
except (Exception, asyncio.CancelledError) as exc:
_exception_raised = True
if logging_obj.prepared_cost is not None:

View file

@ -1,352 +0,0 @@
from __future__ import annotations
import hashlib
import json
from collections.abc import AsyncGenerator, Callable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
if TYPE_CHECKING:
from prisma import Prisma
from litellm.proxy.utils import PrismaClient
AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent
_JSON: Final = TypeAdapter(dict[str, JsonValue])
_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime)
_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...])
_TIMEOUT: Final = timedelta(seconds=10)
class AccountingConflictError(ValueError):
pass
@dataclass(frozen=True, slots=True)
class UsageWrite:
event: InboundUsageEvent | OutboundUsageEvent
envelope: str
targets: tuple[DeliveryTarget, ...] = ()
_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"]
@dataclass(frozen=True, slots=True)
class _AcceptedEvent:
event_id: str
kind: EventKind
source_event_id: str | None
occurred_at: datetime
payload: str
payload_hash: str
targets: tuple[str, ...]
schema_version: int
version: int
def _json(value: object) -> str:
return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
def _prepare(
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...],
envelope: str,
schema_version: int,
version: int,
) -> _AcceptedEvent:
if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1:
raise ValueError("Accounting versions must be positive integers")
validated: Final = type(event).model_validate_json(event.model_dump_json())
payload: Final = _JSON.validate_json(envelope)
if payload.get("event") != validated.model_dump(mode="json"):
raise ValueError("Accounting envelope must contain the matching public event")
destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets)))
serialized: Final = _json(payload)
identity: Final = _json((schema_version, version, destinations, payload))
kind: Final = (
EventKind.COST
if isinstance(validated, CostEvent)
else EventKind.INBOUND_USAGE
if isinstance(validated, InboundUsageEvent)
else EventKind.OUTBOUND_USAGE
)
return _AcceptedEvent(
event_id=validated.event_id,
kind=kind,
source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None,
occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at,
payload=serialized,
payload_hash=hashlib.sha256(identity.encode()).hexdigest(),
targets=destinations,
schema_version=schema_version,
version=version,
)
_INSERT_EVENT: Final = """
INSERT INTO "LiteLLM_AccountingEvent"
(event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload)
VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb)
ON CONFLICT (event_id) DO NOTHING RETURNING event_id
"""
_INSERT_DELIVERIES: Final = """
INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb)
"""
_LOCK_SOURCE: Final = """
SELECT event_id FROM "LiteLLM_AccountingEvent"
WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE
"""
_EXISTING_COST: Final = """
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1
"""
_LOCK_EVENT: Final = """
SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE
"""
_INSERT_USAGE_BATCH: Final = """
INSERT INTO "LiteLLM_AccountingEvent"
(event_id, kind, schema_version, payload_hash, occurred_at, payload)
SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload
FROM jsonb_to_recordset($1::jsonb) AS records(
event_id text, kind text, payload_hash text, occurred_at text, payload jsonb)
ORDER BY event_id
ON CONFLICT (event_id) DO NOTHING RETURNING event_id
"""
_LOCK_USAGE_BATCH: Final = """
SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent"
WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
ORDER BY event_id FOR UPDATE
"""
_LOAD_USAGE_BATCH: Final = """
SELECT event.payload, ARRAY(
SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery
WHERE delivery.event_id = event.event_id ORDER BY target::text
) AS targets
FROM "LiteLLM_AccountingEvent" event
WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
AND event.kind = 'outbound_usage' AND event.schema_version = 1
"""
_INSERT_USAGE_DELIVERIES: Final = """
INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1
FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb)
CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination
"""
_COMPLETED_FAMILY: Final = """
root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage')
AND root.recorded_at < $1::timestamp
AND (root.kind = 'inbound_usage' OR EXISTS (
SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost'
))
AND NOT EXISTS (
SELECT 1 FROM "LiteLLM_AccountingEvent" AS member
WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id))
AND (member.recorded_at >= $1::timestamp
OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery
WHERE delivery.event_id = member.event_id AND delivery.state != 'applied'))
)
"""
_RETENTION_CANDIDATES: Final = (
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE '
+ _COMPLETED_FAMILY
+ " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED"
)
_ELIGIBLE_GROUP: Final = (
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY
)
_DELETE_DELIVERIES: Final = """
DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN (
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
)
"""
_DELETE_EVENTS: Final = """
DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
"""
_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"]
async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict:
if record.source_event_id is not None:
source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id))
if not source:
return "missing_source"
costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id)))
if any(cost["event_id"] != record.event_id for cost in costs):
return "competing_cost"
inserted: Final = _STRINGS.validate_python(
await db.query_raw(
_INSERT_EVENT,
record.event_id,
record.kind.value,
record.source_event_id,
record.schema_version,
record.payload_hash,
record.occurred_at.astimezone(timezone.utc).isoformat(),
record.payload,
)
)
if not inserted:
existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id)))
if not existing or existing[0]["payload_hash"] != record.payload_hash:
return "identity_conflict"
return "duplicate"
await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version)
return "inserted"
async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int:
if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)):
return 0
await db.execute_raw(_DELETE_DELIVERIES, event_id)
return await db.execute_raw(_DELETE_EVENTS, event_id)
class AccountingLedger:
def __init__(self, transaction: Callable[[], _Transaction]) -> None:
self._transaction: Final = transaction
@classmethod
def for_client(cls, client: PrismaClient) -> AccountingLedger:
from prisma.client import TransactionManager
managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True))
@asynccontextmanager
async def transaction() -> AsyncGenerator[Prisma]:
primary: Final = writer_wrapper(client.db)
manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT))
async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
yield db
return cls(transaction)
async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None:
records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes)
if not records:
return
unique: Final = {record.event_id: record for record in records}
if any(record.payload_hash != unique[record.event_id].payload_hash for record in records):
raise AccountingConflictError("Batch contains conflicting copies of a usage identity")
batch: Final = _json(
tuple(
{
"event_id": record.event_id,
"kind": record.kind.value,
"payload_hash": record.payload_hash,
"occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(),
"payload": _JSON.validate_json(record.payload),
}
for record in unique.values()
)
)
async with self._transaction() as db:
inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch))
accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique))))
if len(accepted) != len(unique) or any(
row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted
):
raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set")
deliveries: Final = _json(
tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted)
)
await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries)
async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]:
if not event_ids:
return ()
async with self._transaction() as db:
rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(
await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids))
)
return tuple(
UsageWrite(
OutboundUsageEvent.model_validate_json(_json(_JSON.validate_python(row["payload"])["event"])),
_json(row["payload"]),
TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]),
)
for row in rows
)
async def insert_event(
self,
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...],
*,
envelope: str,
schema_version: int = 1,
version: int = 1,
) -> Literal["inserted", "duplicate"]:
"""Accept T5/T6's serialized envelope, containing its public event under `event`.
The immutable acceptance hash includes the target set and delivery version.
Changing either for an accepted event conflicts. T7 must provide a separate
add-target or new-version workflow that preserves receipts and prevents
duplicate contributions; redelivery cannot silently create or reset work.
The Prisma adapter opens one interactive transaction per event, including
two SET LOCAL statements and a source-row lock for costs. I/O is async,
but awaiting this method inline adds commit latency per event, potentially
several times per request. T5 owns bounded asynchronous capture and an
explicit batching strategy; looping over this method does not batch commits.
Success acknowledges commit, independently of HTTP response acknowledgement.
"""
record: Final = _prepare(event, targets, envelope, schema_version, version)
async with self._transaction() as db:
result: Final = await _accept(db, record)
match result:
case "inserted" | "duplicate":
return result
case "missing_source":
raise AccountingConflictError("Cost source must be a durably accepted outbound usage event")
case "competing_cost":
raise AccountingConflictError("Usage event already has an authoritative cost")
case "identity_conflict":
raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set")
case _:
assert_never(result)
async def purge_completed(
self,
*,
redelivery_before: datetime | None = None,
recovery_before: datetime | None = None,
limit: int = 100,
) -> int:
"""Delete at most `limit` completed families once both retention owners permit it.
T8 must supply a coherent recovery watermark including outstanding holds.
Omitting either watermark disables deletion. Each usage row and its cost
remain together so a retained charge never loses its source or vice versa.
Applied receipts are deleted with the family, ending its retained dedup
protection. The caller must guarantee events recorded before
redelivery_before cannot be redelivered through supported paths.
T8's recovery_before must also protect checkpoints,
outstanding holds and unfinished settlements; age alone is insufficient.
"""
if type(limit) is not int or not 1 <= limit <= 1000:
raise ValueError("Accounting retention batch size must be between 1 and 1000")
if redelivery_before is None or recovery_before is None:
return 0
cutoff: Final = min(
_INSTANT.validate_python(redelivery_before).astimezone(timezone.utc),
_INSTANT.validate_python(recovery_before).astimezone(timezone.utc),
).isoformat()
async with self._transaction() as db:
roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit)))
deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots])
return sum(deleted)

View file

@ -88,14 +88,14 @@ def _internal_call_origin(metadata: str) -> object:
return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)
def _date_from_start_time(start_time: datetime | str) -> str | None:
def date_from_start_time(start_time: datetime | str) -> str | None:
if isinstance(start_time, datetime):
return start_time.date().isoformat()
return start_time[:10] if len(start_time) >= 10 else None
def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None:
date: Final = _date_from_start_time(payload["startTime"])
date: Final = date_from_start_time(payload["startTime"])
origin: Final = _internal_call_origin(payload["metadata"])
if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN):
return None

View file

@ -24,7 +24,6 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
UsageCaptureObservation,
UsageObservation,
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
@ -67,7 +66,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 #
sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export
should_store_prompts_and_responses_in_spend_logs,
)
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs
from litellm.proxy.utils import (
REJECTED_BEFORE_REQUEST_SETUP,
REJECTED_BEFORE_REQUEST_SETUP_KEY,
@ -375,26 +373,13 @@ def _attempt_spend_row_kwargs(
"combined_usage_object": attempt.usage,
"response_cost": attempt.response_cost,
"_accounting_usage_observed": UsageObservation.ESTIMATED,
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost
),
"_accounting_prior_captures": (),
}
row.pop("billed_failed_attempts", None)
row.pop("_accounting_capture", None)
call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id")
if isinstance(call_id, str):
row["litellm_call_id"] = f"{call_id}{suffix}"
row["parent_request_id"] = call_id
metadata["parent_request_id"] = call_id
captured_attempts: Final = request_data.get("_accounting_billed_attempts")
if isinstance(captured_attempts, tuple):
for captured in TypeAdapter(tuple[object, ...]).validate_python(captured_attempts):
if (
isinstance(captured, CapturedUsage)
and captured.payload["litellm_call_id"] == row["litellm_call_id"]
):
row["_accounting_capture"] = captured
standard_logging_object: Final = request_data.get("standard_logging_object")
if isinstance(standard_logging_object, dict):
payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict
@ -480,9 +465,6 @@ class _ProxyDBLogger(CustomLogger):
"Failed to invalidate budget reservation counters after failure release failed"
)
capture_metadata: Final = await self._try_capture_failure_usage(
request_data, original_exception, user_api_key_dict
)
if completed_cost and tracking_disabled:
return
@ -519,10 +501,8 @@ class _ProxyDBLogger(CustomLogger):
_error_information.pop("traceback", None)
_metadata["error_information"] = _error_information
_metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture
{**_metadata, **capture_metadata}
if capture_metadata is not None
else await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled(
_metadata = ( # rebind-ok: enrich identity before building failure spend rows
await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled(
metadata=_metadata, original_exception=original_exception
)
)
@ -629,106 +609,6 @@ class _ProxyDBLogger(CustomLogger):
),
)
async def _failure_capture_metadata(
self, request_data: Mapping[str, object], exception: Exception, auth: UserAPIKeyAuth
) -> Mapping[str, object] | None:
from litellm.proxy.proxy_server import usage_capture
if usage_capture.meter is None:
return None
captured_identity: Final = (
_batch_accounting_for(request_data.get("litellm_call_id")) is not None
and request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is not REJECTED_BEFORE_REQUEST_SETUP
)
current_identity: Final = {
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth),
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth),
}
metadata: Final = (
{
**dict.fromkeys(current_identity),
**{
key: value
for key, value in _REQUEST_DATA.validate_python(
get_litellm_metadata_from_kwargs(dict(request_data))
).items()
if key in current_identity or key == "tags"
},
}
if captured_identity
else current_identity
)
return await self._enrich_failure_metadata_unless_db_stalled(
metadata=metadata, original_exception=exception, resolve_missing_key_identity=not captured_identity
)
async def _try_capture_failure_usage(
self,
request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer
original_exception: Exception,
auth: UserAPIKeyAuth,
) -> Mapping[str, object] | None:
try:
metadata: Final = await self._failure_capture_metadata(request_data, original_exception, auth)
self._capture_failure_usage(request_data, original_exception, auth, metadata)
except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing or reservation release
verbose_proxy_logger.exception("Usage capture failed on the request failure path")
return None
return metadata
def _capture_failure_usage(
self,
request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer
original_exception: Exception,
auth: UserAPIKeyAuth,
metadata: Mapping[str, object] | None = None,
) -> None:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.db.exception_handler import exception_chain
from litellm.proxy.proxy_server import usage_capture
from litellm.proxy.spend_tracking.usage_capture import UsageAcknowledgementError
if any(isinstance(error, UsageAcknowledgementError) for error in exception_chain(original_exception)):
return
usage_meter: Final = usage_capture.meter
if usage_meter is None:
return
logging_obj: Final = request_data.get("litellm_logging_obj")
details: Final[Mapping[str, object]] = (
_str_keyed_fields(logging_obj.model_call_details) if isinstance(logging_obj, Logging) else {}
)
params: Final = _str_keyed_fields(request_data.get("litellm_params") or details.get("litellm_params"))
failure_data: Final = {
**details,
**request_data,
"litellm_params": {
**params,
"metadata": {
**_str_keyed_fields(params.get("metadata")),
**(
metadata
if metadata is not None
else {
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth),
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth),
}
),
"status": "failure",
},
},
}
start: Final = logging_obj.start_time if isinstance(logging_obj, Logging) else request_data.get("start_time")
if not isinstance(start, datetime):
return
end: Final = datetime.now(timezone.utc)
captured: Final = tuple(
capture_callback(usage_meter, row, original_exception, start, end)
for _, row in _failure_spend_rows(failure_data, 0.0, None, self._deployment_lookup)
)
request_data["_accounting_capture"] = captured[0] # rebind-ok: share one receipt with the legacy writer
request_data["_accounting_billed_attempts"] = captured[1:] # rebind-ok: share attempt receipts
async def _write_failure_spend_rows(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -914,7 +794,6 @@ class _ProxyDBLogger(CustomLogger):
resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES,
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
kwargs.update(await captured_callback_kwargs(kwargs, completion_response, start_time, end_time))
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier):
accounting.write_started = True

View file

@ -877,8 +877,6 @@ from litellm.proxy.spend_tracking.spend_event_producer import (
SpendEventProducer,
build_spend_event_producer,
)
from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from litellm.proxy.telemetry.endpoints import router as telemetry_router
from litellm.proxy.telemetry.middleware import TelemetryMiddleware
from litellm.proxy.telemetry.runtime import TelemetryRuntime
@ -1393,7 +1391,6 @@ async def _flush_spend_around_logging_drain_on_shutdown(
finally:
await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock))
await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock))
await usage_capture.stop(GracefulShutdownManager.remaining(clock=clock))
await flush_spend()
@ -3310,39 +3307,9 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False):
spend_event_producer: SpendEventProducer | None = None
usage_capture: Final = UsageCaptureRuntime()
def _configure_usage_capture() -> None:
settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings)
capture_settings: Final = settings.get("accounting_capture")
if settings.get("accounting_sync_ack") and capture_settings is None:
raise ValueError("accounting_sync_ack requires accounting_capture")
if capture_settings is None:
if usage_capture.meter is not None:
raise ValueError("Disabling accounting_capture requires a worker restart")
return
if (
prisma_client is None
or CollectorSettings().enabled
or os.getenv("SPEND_LOGS_URL")
or litellm.apply_search_free_tier
):
raise ValueError(
"Usage capture requires local Prisma writes and disabled collector, remote writer and search free tier"
)
from litellm.proxy.db.accounting_ledger import AccountingLedger
ledger: Final = AccountingLedger.for_client(prisma_client)
writer: Final = UsageBatchWriter(ledger.insert_usage_batch, ledger.load_usage_batch)
usage_capture.start(
writer.write, UsageCaptureSettings.model_validate(capture_settings), load=ledger.load_usage_batch
)
def cost_tracking():
global prisma_client, spend_event_producer
_configure_usage_capture()
if prisma_client is not None:
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger

View file

@ -1553,6 +1553,34 @@
],
"default_model_placeholder": "gpt-3.5-turbo"
},
{
"provider": "EMPIRIOLABS",
"provider_display_name": "EmpirioLabs AI",
"litellm_provider": "empiriolabs",
"credential_fields": [
{
"key": "api_base",
"label": "API Base",
"placeholder": "https://api.empiriolabs.ai/v1",
"tooltip": null,
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "api_key",
"label": "API Key",
"placeholder": null,
"tooltip": null,
"required": true,
"field_type": "password",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "empiriolabs/aplomb-1"
},
{
"provider": "EMPOWER",
"provider_display_name": "Empower",

View file

@ -3,7 +3,7 @@ import json
import os
import re
import secrets
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from datetime import datetime as dt
from functools import reduce
@ -58,7 +58,6 @@ from litellm.litellm_core_utils.internal_call_metadata import (
)
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
UsageObservation,
coerce_model_access_groups,
is_valid_sha256_hash,
request_model_access_groups_from_litellm_params,
@ -592,7 +591,6 @@ def get_logging_payload(
start_time: datetime,
end_time: datetime,
llm_router: "Router | None" = None,
capture_usage: Callable[[SpendLogsPayload, str | None], None] | None = None,
) -> SpendLogsPayload:
if kwargs is None:
kwargs = {}
@ -982,14 +980,6 @@ def get_logging_payload(
payload.get("spend"),
)
if capture_usage is not None:
observed_usage: Final = (
StandardLoggingPayloadSetup.get_usage_as_dict({"usage": usage})
if usage and kwargs.get("_accounting_usage_observed") is not UsageObservation.UNKNOWN
else None
)
capture_usage(payload, safe_dumps(observed_usage) if observed_usage is not None else None)
# Explicitly clear large intermediate objects to reduce memory pressure
del response_obj_dict, usage, clean_metadata, additional_usage_values

File diff suppressed because it is too large Load diff

View file

@ -1,330 +0,0 @@
from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import Awaitable, Callable, Mapping
from concurrent.futures import Future
from dataclasses import dataclass
from itertools import accumulate, chain, count, takewhile
from threading import Lock
from typing import Final
from pydantic import BaseModel, ConfigDict, Field
from litellm._logging import verbose_proxy_logger
from litellm.accounting import (
DeliveryTarget,
Dimension,
DimensionName,
InboundUsageEvent,
OutboundUsageEvent,
)
from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite
_DAILY_TARGETS: Final = (
(DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND),
(DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND),
(DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND),
(DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND),
(DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND),
(DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND),
(DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND),
(DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE),
(DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND),
)
def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]:
names: Final = frozenset(dimension.name for dimension in dimensions)
return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names))
class UsageCaptureSettings(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
batch_size: int = Field(default=1000, ge=2, le=1000)
flush_interval_seconds: float = Field(default=5.0, gt=0, le=10)
max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024)
max_pending_records: int = Field(default=10000, ge=2)
class UsageCaptureError(RuntimeError):
pass
@dataclass(frozen=True, slots=True)
class _Pending:
writes: tuple[UsageWrite, ...]
size: int
committed: Future[None]
finished: asyncio.Event
prerequisites: tuple[asyncio.Future[None], ...]
async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None:
try:
await asyncio.wait_for(wakeup.wait(), timeout)
except asyncio.TimeoutError:
return
def _report_failure(receipt: Future[None]) -> None:
if not receipt.cancelled() and receipt.exception() is not None:
verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception())
def _retryable_persistence_failure(error: Exception) -> bool:
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain
outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error)
return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any(
isinstance(cause, Exception)
and (
PrismaDBExceptionHandler.is_deadlock_error(cause)
or PrismaDBExceptionHandler.is_lock_timeout_error(cause)
or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001"
)
for cause in exception_chain(error)
)
def _isolatable_persistence_failure(error: Exception) -> bool:
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain
return isinstance(error, AccountingConflictError) or any(
isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22")
for cause in exception_chain(error)
)
class DurableUsageMeter:
def __init__(
self,
write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]],
settings: UsageCaptureSettings,
*,
wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch,
retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep,
) -> None:
self._write_batch: Final = write_batch
self.settings: Final = settings
self._wait_for_batch: Final = wait_for_batch
self._retry_delay: Final = retry_delay
self._loop: Final = asyncio.get_running_loop()
self._admission: Final = Lock()
self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue
self._wakeup: Final = asyncio.Event()
self._capacity_changed: Final = asyncio.Event()
self._reservations: Mapping[object, tuple[int, int]] = {}
self._reserved_bytes = 0
self._task: asyncio.Task[None] | None = None
self._closed = False
self.pending_bytes = 0
self.pending_records = 0
def submit(
self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = ()
) -> asyncio.Future[None]:
size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes)
with self._admission:
if (
self._closed
or not writes
or len(writes) > self.settings.batch_size
or self.pending_bytes + size > self.settings.max_pending_bytes
or self.pending_records + len(writes) > self.settings.max_pending_records
):
raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted")
receipt: Final = Future[None]()
receipt.set_running_or_notify_cancel()
receipt.add_done_callback(_report_failure)
self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites))
self.pending_bytes += size
self.pending_records += len(writes)
self._loop.call_soon_threadsafe(self._ensure_worker)
wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop)
wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception())
return wrapped
def reserve_producer(self, size: int, headroom: int) -> object:
token: Final = object()
self.resize_producer(token, size, headroom)
return token
def resize_producer(self, token: object, size: int, headroom: int) -> None:
with self._admission:
previous: Final = self._reservations.get(token, (0, 0))[0]
delta: Final = size - previous
others: Final = (value[1] for key, value in self._reservations.items() if key is not token)
required: Final = max((headroom, *others))
if (
self._closed
or self.pending_bytes + delta > self.settings.max_pending_bytes
or self._reserved_bytes + delta + required > self.settings.max_pending_bytes
):
raise UsageCaptureError("Batch capture producer byte capacity exhausted")
self._reservations = {**self._reservations, token: (size, headroom)}
self._reserved_bytes += delta
self.pending_bytes += delta
self._capacity_changed.set()
def release_producer(self, token: object) -> None:
with self._admission:
size, _ = self._reservations[token]
self._reservations = {key: value for key, value in self._reservations.items() if key is not token}
self._reserved_bytes -= size
self.pending_bytes -= size
self._capacity_changed.set()
def _ensure_worker(self) -> None:
if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites:
self._wakeup.set()
if self._task is None or self._task.done():
self._task = asyncio.create_task(self._run(), name="accounting usage capture")
async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None:
envelope: Final = '{"event":' + event.model_dump_json() + "}"
await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),)))
async def record_writes(self, writes: tuple[UsageWrite, ...]) -> None:
size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes)
capacity: Final = self.settings.max_pending_bytes - self._reserved_bytes
if len(writes) > min(self.settings.batch_size, self.settings.max_pending_records) or (
size > capacity and len(writes) > 1
):
middle: Final = len(writes) // 2
await self.record_writes(writes[:middle])
await self.record_writes(writes[middle:])
return
if size > capacity:
raise UsageCaptureError("One usage observation exceeds capture capacity")
while (receipt := self._submit_available(writes)) is None:
self._wakeup.set()
await self._capacity_changed.wait()
self._wakeup.set()
await asyncio.shield(receipt)
def _submit_available(self, writes: tuple[UsageWrite, ...]) -> asyncio.Future[None] | None:
self._capacity_changed.clear()
try:
receipt: Final = self.submit(writes)
except UsageCaptureError:
if self._closed or not self.pending_records or not writes:
raise
return None
return receipt
async def flush(self) -> None:
self._ensure_worker()
self._wakeup.set()
if self._task is not None:
await asyncio.shield(self._task)
async def close(self) -> None:
with self._admission:
self._closed = True
await self.flush()
async def stop(self, timeout: float) -> None:
with self._admission:
self._closed = True
try:
await asyncio.wait_for(self.close(), timeout)
except asyncio.TimeoutError:
if self._task is not None:
self._task.cancel()
await asyncio.gather(self._task, return_exceptions=True)
while self._queue:
self._finish(
self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement")
)
verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete")
def _take_batch(self) -> tuple[_Pending, ...]:
with self._admission:
if self._queue[0].prerequisites:
return (self._queue.popleft(),)
independent: Final = takewhile(lambda item: not item.prerequisites, self._queue)
sizes: Final = accumulate(len(item.writes) for item in independent)
count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes))
return tuple(self._queue.popleft() for _ in range(count))
async def _persist(self, writes: tuple[UsageWrite, ...]) -> None:
for attempt in count():
try:
await self._write_batch(writes)
except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes
if not _retryable_persistence_failure(error):
raise
verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry")
await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6)))
else:
return
async def _commit(self, batch: tuple[_Pending, ...]) -> None:
writes: Final = tuple(chain.from_iterable(item.writes for item in batch))
try:
if batch[0].prerequisites:
outcomes: Final = await asyncio.gather(
*(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True
)
for outcome in outcomes:
if isinstance(outcome, BaseException):
raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome
await self._persist(writes)
except asyncio.CancelledError:
self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted"))
raise
except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers
if len(batch) > 1 and _isolatable_persistence_failure(error):
try:
for item in batch:
await self._commit((item,))
except asyncio.CancelledError:
self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted"))
raise
return
self._finish(
batch,
UsageCaptureError("Usage identity conflicts with accepted content")
if isinstance(error, AccountingConflictError)
else error,
)
else:
self._finish(batch, None)
def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None:
for item in batch:
if item.finished.is_set():
continue
item.finished.set()
with self._admission:
self.pending_bytes -= item.size
self.pending_records -= len(item.writes)
self._capacity_changed.set()
if item.committed.done():
continue
if error is None:
item.committed.set_result(None)
else:
item.committed.set_exception(error)
async def _run(self) -> None:
try:
while self._queue:
await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds)
self._wakeup.clear()
await self._commit(self._take_batch())
if (
self._closed
or self.pending_records >= self.settings.batch_size
or self._queue
and self._queue[0].prerequisites
):
self._wakeup.set()
except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally
while self._queue:
self._finish(
self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement")
)
raise

View file

@ -168,7 +168,6 @@ from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int
from litellm.litellm_core_utils.litellm_logging import (
BilledRoundsDispatch,
Logging,
UsageObservation,
deployment_pricing_model_info,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
@ -1805,39 +1804,6 @@ def _failure_fields_to_lift(
if _prompt_estimate is not None
else ()
)
from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation
_recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed")
_current_usage: Final = (
_current_failure_usage_to_lift(
_model_call_details,
request_data,
_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None,
_billed,
_deployment_id,
_base_model,
_estimate_call_type,
)
if Logging.usage_capture_callback is not None
and (_billed or _model_call_details.get("combined_usage_object") is not None)
else None
)
_current_observation: Final = (
_recorded_observation
if isinstance(_recorded_observation, UsageObservation)
else UsageObservation.UNKNOWN
if _current_usage is None
else UsageObservation.KNOWN
if _model_call_details.get("combined_usage_object") is not None
else UsageObservation.ESTIMATED
)
_capture_observation: Final = UsageCaptureObservation.from_usage(
_current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None,
_current_observation,
float(_current_usage[1])
if _current_usage is not None and isinstance(_current_usage[1], (int, float))
else None,
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
("first_api_call_start_time", _first_handoff),
@ -1846,16 +1812,12 @@ def _failure_fields_to_lift(
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
("standard_logging_object", _model_call_details.get("standard_logging_object")),
("billed_failed_attempts", _earlier_attempts or None),
("_accounting_capture", _model_call_details.get("_accounting_capture")),
("_accounting_pricing_snapshot", _model_call_details.get("_accounting_pricing_snapshot")),
("_accounting_finalized_cost", _logging_obj.finalized_cost if isinstance(_logging_obj, Logging) else None),
(
"_accounting_allocation",
_logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None,
),
("_accounting_usage_observed", _current_observation if Logging.usage_capture_callback is not None else None),
("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None),
("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")),
)
return MappingProxyType({key: value for key, value in _entries if value is not None})

View file

@ -3866,11 +3866,6 @@ class PythonCore(DeprecatedRouterMembers):
priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor
key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost"
}
if LiteLLMLogging.usage_capture_callback is not None:
from litellm._internal_context import internal_sub_call
with internal_sub_call(cut_logging.litellm_call_id):
cut_logging.capture_usage_before_logging(priced_partial)
return ContinuedPartial(
usage=BaseTokenUsageProcessor.combine_usage_objects(
usage_objects=[partial_usage, inherited_usage]
@ -5240,8 +5235,19 @@ class PythonCore(DeprecatedRouterMembers):
# _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping
# the static shared deployment's id instead would let one tenant's bad credentials
# cool down the deployment every other tenant sharing this config relies on.
if getattr(exception, "failed_deployment_id", None):
return
effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({})
self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info}))
configured_deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info).get(
"original_model_id"
)
if isinstance(configured_deployment_id, str) and configured_deployment_id:
exception.failed_configured_deployment_id = configured_deployment_id # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: the exception carries the failure to _maybe_run_weighted_failover
def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None:
if picked:
self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs)
@staticmethod
def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None:
@ -6118,9 +6124,10 @@ class PythonCore(DeprecatedRouterMembers):
def _image_generation(self, prompt: str, model: str, **kwargs):
model_name: Final = ""
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
deployment: Final = self.get_available_deployment(
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6156,6 +6163,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def aimage_generation(self, prompt: str, model: str, **kwargs):
@ -6181,11 +6191,12 @@ class PythonCore(DeprecatedRouterMembers):
async def _aimage_generation(self, prompt: str, model: str, **kwargs):
model_name = model
start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
try:
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6246,6 +6257,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def atranscription(self, file: FileTypes, model: str, **kwargs):
@ -6292,10 +6306,11 @@ class PythonCore(DeprecatedRouterMembers):
async def _atranscription(self, file: FileTypes, model: str, **kwargs):
model_name: Final = model
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _atranscription()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6329,6 +6344,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
@ -6389,10 +6407,11 @@ class PythonCore(DeprecatedRouterMembers):
async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
model_name: Final = model
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _aspeech()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6433,6 +6452,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def arerank(self, model: str, **kwargs):
@ -6458,9 +6480,10 @@ class PythonCore(DeprecatedRouterMembers):
async def _arerank(self, model: str, **kwargs):
model_name = None
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _rerank()- model: %s", model)
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
@ -6492,6 +6515,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
def text_completion(
@ -6573,10 +6599,11 @@ class PythonCore(DeprecatedRouterMembers):
raise e
async def _atext_completion(self, model: str, prompt: str, **kwargs):
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _atext_completion()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": prompt}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6612,6 +6639,9 @@ class PythonCore(DeprecatedRouterMembers):
if model is not None:
self.fail_calls[model] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def aadapter_completion(
@ -6643,10 +6673,11 @@ class PythonCore(DeprecatedRouterMembers):
raise e
async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs):
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=[{"role": "user", "content": "default text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
@ -6682,6 +6713,9 @@ class PythonCore(DeprecatedRouterMembers):
if model is not None:
self.fail_calls[model] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object:
@ -7670,6 +7704,7 @@ class PythonCore(DeprecatedRouterMembers):
"""
handler_name: Final = original_function.__name__
metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call")
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
self._update_kwargs_before_fallbacks(
@ -7678,7 +7713,7 @@ class PythonCore(DeprecatedRouterMembers):
metadata_variable_name=metadata_variable_name,
)
with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name):
deployment: Final = self.get_available_deployment(
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
@ -7724,6 +7759,9 @@ class PythonCore(DeprecatedRouterMembers):
verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
if model is not None:
self.fail_calls[model] += 1
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
def embedding(
@ -7745,10 +7783,11 @@ class PythonCore(DeprecatedRouterMembers):
def _embedding(self, input: str | list[object], model: str, **kwargs):
model_name = None
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside embedding()- model: %s", model)
with embedding_deployment_selection():
deployment: Final = self.get_available_deployment(
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
@ -7805,6 +7844,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
async def aembedding(
@ -7834,11 +7876,12 @@ class PythonCore(DeprecatedRouterMembers):
async def _aembedding(self, input: str | list[object], model: str, **kwargs):
model_name = None
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
try:
verbose_router_logger.debug("Inside _aembedding()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
with embedding_deployment_selection():
deployment: Final = await self._core_async_get_available_deployment(
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
@ -7885,6 +7928,9 @@ class PythonCore(DeprecatedRouterMembers):
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
self._stamp_failed_attempt(
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
)
raise e
#### FILES API ####
@ -9123,8 +9169,9 @@ class PythonCore(DeprecatedRouterMembers):
kwargs[metadata_variable_name] = meta
if not isinstance(meta, dict):
return None
configured_failed_id: Final[str | None] = getattr(exception, "failed_configured_deployment_id", None)
prev_excluded: Final = set(meta.get("_failover_excluded_ids") or [])
excluded: Final = prev_excluded | {failed_id}
excluded: Final = prev_excluded | {configured_failed_id or failed_id}
all_ids: Final = {
(d.get("model_info") or {}).get("id")

View file

@ -431,6 +431,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
regional_processing_uplift_multiplier_us: (
float | None
) # OpenAI US data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%)
fedramp_uplift_multiplier: ReadOnly[
float | None
] # OpenAI FedRAMP endpoint (gov.api.openai.com) uplift multiplier applied to all token costs (e.g. 1.10 = +10%)
regional_endpoint_uplift_multiplier: ReadOnly[
float | None
] # Vertex AI non-global (regional) endpoint uplift multiplier applied to all token costs (e.g. 1.10 = +10%)
@ -4258,6 +4261,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
annotation_cost_per_page_batches: float | None = None
regional_processing_uplift_multiplier_eu: float | None = None
regional_processing_uplift_multiplier_us: float | None = None
fedramp_uplift_multiplier: float | None = None
regional_endpoint_uplift_multiplier: float | None = None
@classmethod
@ -4646,6 +4650,7 @@ class LlmProviders(str, Enum):
STRANDS_DECIDER = "strands_decider"
LAYA = "laya"
BESPOKE = "bespoke"
EMPIRIOLABS = "empiriolabs"
EU_JEV = "eu_jev"
MISTRAL = "mistral"
MILVUS = "milvus"
@ -5260,11 +5265,13 @@ class DataResidency(Enum):
Inferred from the OpenAI api_base host (eu.api.openai.com -> EU,
us.api.openai.com -> US). Used to apply the regional-processing
cost uplift (see ``regional_processing_uplift_multiplier_<region>``
on ModelInfo).
on ModelInfo). gov.api.openai.com maps to FEDRAMP, which applies
``fedramp_uplift_multiplier`` instead.
"""
US = "us"
EU = "eu"
FEDRAMP = "fedramp"
LLMResponseTypes = (

View file

@ -1647,7 +1647,6 @@ def dispatch_success_logging(
is_litellm_internal_call: bool,
) -> None:
if is_litellm_internal_call:
logging_obj.capture_usage_before_logging(result, start_time, end_time)
return
_schedule_async_success_logging(
@ -1771,7 +1770,6 @@ def _sync_speech_end_logger(
_mark_speech_stream_closed_early(logging_obj)
case StreamCompleted():
pass
logging_obj.capture_usage_before_logging(result, start_time, end_time)
if is_litellm_internal_call:
return
executor.submit(ctx.run, logging_obj.success_handler, result, start_time, end_time)
@ -1800,14 +1798,9 @@ def _schedule_async_success_logging(
if not logging_obj.claim_async_success_log():
return
if logging_obj.defer_async_logging and not is_completion_with_fallbacks:
logging_obj.capture_usage_before_logging(result, start_time, end_time, freeze_only=True)
def _enqueue_async_logging() -> None:
from litellm.responses.streaming_iterator import spawn_logging_task
if not is_completion_with_fallbacks:
logging_obj.capture_usage_before_logging(result, start_time, end_time)
if logging_obj.defer_async_logging:
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result, start_time=start_time, end_time=end_time
@ -2421,21 +2414,17 @@ def client(original_function):
executor=executor,
),
)
else:
logging_obj.capture_usage_before_logging(
TypeAdapter(object).validate_python(result), start_time, end_time
elif not is_internal_call.get():
# Copy the current context to propagate it to the background thread
# This is essential for OpenTelemetry span context propagation
ctx: Final = contextvars.copy_context()
executor.submit(
ctx.run,
logging_obj.success_handler,
result,
start_time,
end_time,
)
if not is_internal_call.get():
# Copy the current context to propagate it to the background thread
# This is essential for OpenTelemetry span context propagation
ctx: Final = contextvars.copy_context()
executor.submit(
ctx.run,
logging_obj.success_handler,
result,
start_time,
end_time,
)
# RETURN RESULT
return result
except Exception as e:
@ -7676,6 +7665,7 @@ def get_model_info_helper(
regional_processing_uplift_multiplier_us=_model_info.get(
"regional_processing_uplift_multiplier_us", None
),
fedramp_uplift_multiplier=_model_info.get("fedramp_uplift_multiplier", None),
regional_endpoint_uplift_multiplier=_model_info.get("regional_endpoint_uplift_multiplier", None),
output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None),
output_cost_per_character=_model_info.get("output_cost_per_character", None),
@ -10474,6 +10464,7 @@ class ProviderConfigManager:
if provider == LlmProviders.BESPOKE:
return litellm.BespokeDecisionsConfig()
if provider in (
LlmProviders.EMPIRIOLABS,
LlmProviders.OPENAI,
LlmProviders.AZURE,
LlmProviders.AZURE_AI,
@ -10481,6 +10472,7 @@ class ProviderConfigManager:
LlmProviders.PARALLEL_AI,
):
configs: Final = {
LlmProviders.EMPIRIOLABS: litellm.EmpirioLabsDecisionsConfig,
LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig,
LlmProviders.AZURE: litellm.AzureDecisionsConfig,
LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig,

File diff suppressed because it is too large Load diff

View file

@ -5277,7 +5277,7 @@
}
},
"empiriolabs": {
"display_name": "EmpirioLabs (`empiriolabs`)",
"display_name": "EmpirioLabs AI (`empiriolabs`)",
"url": "https://docs.litellm.ai/docs/providers/empiriolabs",
"endpoints": {
"chat_completions": false,

File diff suppressed because one or more lines are too long

View file

@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig):
)
def _register_deployment(proxy, deployment: CompatDeployment) -> str:
"""Register one deployment and return its proxy-assigned model_id
once it is servable on the data plane. The aliases are shared by every
cell and, under xdist, by every worker, so no call to them belongs to
one test and none is cached: the matrix exists to reach real providers."""
return proxy.create_model(
deployment.model_name,
deployment.litellm_params,
def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]):
"""The aliases are shared by every cell and, under xdist, by every worker,
so no call to them belongs to one test and none is cached: the matrix
exists to reach real providers."""
from models import ModelInfoBody, ModelNewBody
return proxy.register_models(
tuple(
ModelNewBody(
model_name=deployment.model_name,
litellm_params=deployment.litellm_params,
model_info=ModelInfoBody(),
)
for deployment in deployments
),
provider_live=True,
)
@ -632,18 +639,14 @@ def _compat_models_registered() -> Any:
yield
return
from proxy_client import ModelNotRegistered, ModelRegistered
from requests import RequestException
proxy = _build_control_plane_client(proxy_config)
registered_ids: list[str] = []
failures: list[tuple[str, str]] = []
registrations = _register_deployments(proxy, load_all_deployments())
registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered))
failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered))
try:
for deployment in load_all_deployments():
try:
model_id = _register_deployment(proxy, deployment)
registered_ids.append(model_id)
except (AssertionError, RequestException) as exc:
failures.append((deployment.model_name, str(exc)))
if failures:
summary = "\n".join(
f" - {name}: {reason}" for name, reason in failures

View file

@ -35,6 +35,7 @@ from e2e_config import (
PROMPT_CACHING_OPT_IN_ENV,
PROVIDER_EDGE_HOST_OPT_IN_ENV,
PROXY_BASE_URL,
PROXY_WIDE_CONFIG_OPT_IN_ENV,
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
SAGEMAKER_NOVA_OPT_IN_ENV,
SECRET_MANAGER_OPT_IN_ENV,
@ -54,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from provider_edge import replay_leftover_error
from proxy_client import ProxyClient, build_proxy_client
from stack_lock import stack_lock
_E2E_TEST_RAN = pytest.StashKey[bool]()
_CALL_PASSED = pytest.StashKey[bool]()
@ -77,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"secret_manager": SECRET_MANAGER_OPT_IN_ENV,
"sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
"sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
"proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV,
}
)
@ -162,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set",
)
config.addinivalue_line(
"markers",
"quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; "
"every other test waits for it to finish",
)
config.addinivalue_line(
"markers",
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
@ -200,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None:
"sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
"deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
)
config.addinivalue_line(
"markers",
"proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite "
"is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:
@ -269,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None:
ahead of xdist's own collection-finish report, and the controller schedules no
test until every worker has reported, so this is the idle footprint of a stack
that just passed its readiness gate. The fixture hands the capture to the
idle-budget test in router/test_reliability_memory_e2e.py."""
idle-budget test in perf/test_reliability_memory_e2e.py."""
if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items):
return
session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS)
@ -299,12 +300,6 @@ def _proxy_fail_reason() -> str | None:
return None
@pytest.hookimpl(wrapper=True)
def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]:
with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None):
return (yield)
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.

View file

@ -44,8 +44,6 @@
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
- {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"}
- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"}
- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"}
- {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
- {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}

View file

@ -210,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))
@ -234,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES",
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1"))
MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200"))
MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768"))
MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64"))
def ws_base_url() -> str:

View file

@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t
/config/update. A chat completion that attaches the server and asks to search logs must report,
in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all
of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's
OPENAI_API_KEY for the embedding deployment.
OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is
marked `proxy_wide_config` and runs on a stack the main suite is not driving.
"""
from __future__ import annotations
@ -23,7 +24,7 @@ from mcp_client import McpClient
from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool
from pydantic import BaseModel
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config]
TOP_K: Final = 3

View file

@ -1262,27 +1262,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]):
pass
class ConfigPatchResponse(BaseModel):
status: str
values: dict[str, float | dict[str, float]]
class CostDiscountConfig(RootModel[dict[str, float]]):
pass
class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]):
pass
class CostDiscountConfigResponse(BaseModel):
values: dict[str, float]
class CostMarginConfigResponse(BaseModel):
values: dict[str, float | dict[str, float]]
class BedrockGuardrailParams(BaseModel):
guardrail: Literal["bedrock"] = "bedrock"
mode: Literal["pre_call"] = "pre_call"

View file

@ -1,223 +0,0 @@
from __future__ import annotations
import math
import time
from collections.abc import Callable
from dataclasses import replace
from typing import Final
import psycopg
from cases.small_bodies_5rps import CASE
from e2e_config import MASTER_KEY
from e2e_http import send_text_external
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from setup.performance_test_case import PerformanceTestCase
from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes
from setup.wire import WireRecord
MAX_TRAFFIC_SECONDS: Final = 6 * 3600
_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...])
_ACTIVITY: Final = TypeAdapter(tuple[list[int], int])
_COUNT: Final = TypeAdapter(tuple[int])
_SERVICES: Final = ("postgres", "batch_write_to_db")
def upgrade_case() -> PerformanceTestCase:
return replace(
CASE,
name="database_upgrade_5rps",
traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)),
limits=tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
),
)
class MetricsHeaders(BaseModel):
model_config = ConfigDict(frozen=True)
authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization")
def postgres_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
names: Final = frozenset(family.name for family in families)
required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES)
samples: Final = tuple(sample for family in families for sample in family.samples)
has_database_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in samples
)
assert required <= names or has_database_activity, f"Missing database service counters: {required - names}"
failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required})
assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter"
return sum(failures)
def assert_database_metrics(urls: tuple[str, ...]) -> None:
assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint"
for url in urls:
assert_database_metrics_url(url)
def assert_database_metrics_url(url: str) -> None:
result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10)
assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}"
assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}"
def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]:
ids: Final = tuple(record.call_id for record in records if record.call_id)
if not ids:
return frozenset()
rows: Final = _ROWS.validate_python(
connection.execute(
'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),)
).fetchall()
)
return matched_response_ids(records, rows)
def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]:
call_ids: Final = frozenset(call_id for _, call_id in rows if call_id)
return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids)
def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]:
return _ACTIVITY.validate_python(
connection.execute("""SELECT
ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()),
(SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock')
""").fetchone()
)
def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None:
assert state.phase == "complete", f"Upgrade did not finish: {state.phase}"
assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid"
(missing,) = _COUNT.validate_python(
connection.execute(
"""
SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS (
SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name
AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""",
(list(state.required_migrations),),
).fetchone()
)
assert missing == 0, f"{missing} target migrations were not applied"
(unfinished,) = _COUNT.validate_python(
connection.execute("""
SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL
""").fetchone()
)
assert unfinished == 0, f"{unfinished} unfinished migrations"
assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, (
"Baseline rows changed or disappeared"
)
class UpgradeMonitor:
def __init__(
self,
*,
state: Callable[[], UpgradeState],
check_metrics: Callable[[], None],
read_activity: Callable[[], tuple[list[int], int]],
read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]],
start_indexes: Callable[[], None],
clock: Callable[[], float] = time.monotonic,
wall_clock: Callable[[], float] = time.time,
) -> None:
self.state: Final = state
self.check_metrics: Final = check_metrics
self.read_activity: Final = read_activity
self.read_persisted: Final = read_persisted
self.start_indexes: Final = start_indexes
self.clock: Final = clock
self.wall_clock: Final = wall_clock
self.started: Final = clock()
self.status_limit: Final = next(
limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count" and limit.max is not None
)
self.pending: tuple[WireRecord, ...] = ()
self.processed = 0
self.successes = 0
self.persisted_count = 0
self.overlap_writes = 0
self.max_lock_waiters = 0
self.previous_builds: frozenset[int] = frozenset()
self.build_seen_at: float | None = None
self.released = False
self.completed_at: float | None = None
def poll(self, records: tuple[WireRecord, ...]) -> bool:
self.check_metrics()
current: Final = self.state()
assert current.phase != "failed", "Target migration/index worker failed"
assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout"
successful: Final = tuple(record for record in records if record.status == 200)
self.processed += len(records)
self.successes += len(successful)
assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses"
assert all(record.ok and record.text for record in successful), "Failed or empty inference response"
assert all(record.response_id and record.call_id for record in successful), (
"Response missing spend-log identifiers"
)
builds, lock_waiters = self.read_activity()
if not self.previous_builds.intersection(builds):
self.build_seen_at = self.wall_clock() if builds else None
self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters)
self.pending += successful
found: Final = self.read_persisted(self.pending)
self.persisted_count += len(found)
still_building, _ = self.read_activity()
since: Final = self.build_seen_at
if since is not None and frozenset(builds).intersection(still_building):
self.overlap_writes += sum(
1 for record in self.pending if record.response_id in found and record.started_at >= since
)
self.previous_builds = frozenset(builds)
self.pending = tuple(record for record in self.pending if record.response_id not in found)
if self.successes and not self.released:
self.start_indexes()
self.released = True
if current.phase == "complete" and self.completed_at is None:
self.completed_at = self.clock()
return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300
def assert_complete(self, state: UpgradeState) -> None:
assert self.successes >= 300, f"Only {self.successes} requests completed"
assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log"
assert self.persisted_count == self.successes, "Spend-log count does not match successful requests"
assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work"
def upgrade_properties(state: UpgradeState, monitor: UpgradeMonitor) -> tuple[tuple[str, float], ...]:
"""The upgrade's own numbers under the `perf_` prefix the dashboard ingests, next to the perf rows."""
return (
("perf_upgrade_schema_seconds", state.schema_seconds),
("perf_upgrade_index_seconds", state.index_seconds),
("perf_upgrade_overlap_writes", float(monitor.overlap_writes)),
("perf_upgrade_max_lock_waiters", float(monitor.max_lock_waiters)),
("perf_upgrade_pending_indexes", float(len(state.pending_indexes))),
("perf_upgrade_unmatched_requests", float(len(monitor.pending))),
)
def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: read_state(connection),
check_metrics=lambda: assert_database_metrics(urls),
read_activity=lambda: activity(connection),
read_persisted=lambda records: persisted(connection, records),
start_indexes=lambda: release_indexes(connection),
)

View file

@ -1,102 +0,0 @@
from __future__ import annotations
import os
import sys
import time
from pathlib import Path
from typing import Final
import pytest
from idp import Keycloak
from lifecycle import ResourceManager
from memory_readings import read_rss_everywhere
from proxy_client import ProxyClient
from setup.datastore import read_db_counters, read_redis_counters
from setup.observe import (
Observation,
evaluate,
junit_properties,
read_mock_requests,
render,
reset_mock,
set_mock_latency,
)
from setup.plan import build_plan
from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, mock_url
from setup.run import run_plan
from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties
from setup.upgrade_worker import connect, read_state
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only")
def test_database_upgrade(
proxy: ProxyClient, idp: Keycloak, resources: ResourceManager, request: pytest.FixtureRequest
) -> None:
case: Final = upgrade_case()
urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS)
for url in urls:
_ = reset_mock(url)
_ = set_mock_latency(url, case.mock_latency)
setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources)
metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
with connect() as connection:
assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting"
rss_before: Final = read_rss_everywhere(proxy)
db_before: Final = read_db_counters()
redis_before: Final = read_redis_counters()
monitor: Final = monitor_for(connection, metrics)
with gateway_cpu_sampler() as cpu:
outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30)
assert outcome.locust_exit == 0, outcome.locust_output
remaining: Final = outcome.records[monitor.processed :]
_ = monitor.poll(remaining)
deadline: Final = time.monotonic() + 120
while monitor.pending and time.monotonic() < deadline:
time.sleep(1)
_ = monitor.poll(())
state: Final = read_state(connection)
_ = sys.stdout.write(
f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} "
f"unmatched_requests={len(monitor.pending)}\n"
)
rss_after: Final = read_rss_everywhere(proxy)
observed: Final = Observation(
records=outcome.records,
mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)),
guardrail_calls=(),
rss_before=rss_before,
rss_after=rss_after,
wall_seconds=outcome.wall_seconds,
spend_logged_share=1,
db_before=db_before,
db_after=read_db_counters(),
redis_before=redis_before,
redis_after=read_redis_counters(),
gateway_cpu_seconds=cpu.seconds,
)
rows, exceeded = evaluate(case, observed)
item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped
assert isinstance(item, pytest.Item)
item.user_properties[:] = [
entry for entry in item.user_properties if not entry[0].startswith("perf_")
] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set
item.user_properties.extend(junit_properties(rows))
item.user_properties.extend(upgrade_properties(state, monitor))
_ = sys.stdout.write(render(case, observed, rows) + "\n")
_ = sys.stdout.write(
f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} "
f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} "
f"pending_indexes={state.pending_indexes}\n"
)
artifact: Final = os.environ.get("E2E_ARTIFACT_DIR")
if artifact:
_ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8")
monitor.assert_complete(state)
assert_finished(connection, state)
assert {(r.replica, r.hostname) for r in rss_before.readings} == {
(r.replica, r.hostname) for r in rss_after.readings
}, "Gateway pods replaced during the upgrade"
assert not exceeded, "\n".join(exceeded)

View file

@ -48,7 +48,7 @@ from setup.ramp import step_properties
from setup.run import run_plan
from setup.telemetry import read_gateway_telemetry, telemetry_sampler
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
pytestmark = pytest.mark.perf
_SPEND_SAMPLE: Final = 40
_WARMUP_SECONDS: Final = 60

View file

@ -1,17 +1,19 @@
"""Live e2e: a few hundred requests that fail before any provider answers must not
grow the proxy's resident memory past a loose ceiling once the proxy is warm.
grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads
the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack.
The regression this guards shipped in v1.100.0: every retry breadcrumb copied the
whole request and the copies nested into one router-global list, so a proxy under
retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that
shape: a model group whose deployments refuse at the socket (an unreachable base
URL) with cooldown_time 0 so the router keeps retrying them, per-request retries,
and a fallback group that refuses the same way, each request carrying a long chat
URL, with a literal key since the perf stack carries no provider keys) with
cooldown_time 0 so the router keeps retrying them, per-request retries, and a
fallback group that refuses the same way, each request carrying a long chat
transcript so every whole-request copy costs hundreds of containers instead of a
handful. Under the stack's cooldown policy a
deployment that fails a handful of times in a row is benched (a bad-credential 401
included), the router answers "No deployments available" without retrying, and the
retry loop that leaks stops running; cooldown_time 0 keeps it running.
handful. Under the stack's cooldown policy a deployment that fails a handful of
times in a row is benched (a bad-credential 401 included), the router answers "No
deployments available" without retrying, and the retry loop that leaks stops running;
cooldown_time 0 keeps it running.
Four identical batches run back to back. The first two are the warmup that grows
the proxy's caches and allocator arenas, the last two are the measured window, and
@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time,
a worker that answered only one checkpoint would otherwise drop out of the
comparison, which is where a leaking worker could hide.
RSS alone is a coarse gauge: on the release stack (spend logs storing prompts,
json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS
by only about 15 MB per 300 failing requests, while every failing request's stored
request snapshot carried a copy of the request per failed attempt, over 100 KB on
the first call and a couple of MB once the copies nested, against tens of KB with
the fix. So the first check sends one failing request before the phases, reads its
spend log back through /spend/logs, and holds the stored request body to a fixed
size budget: the deterministic catch for a breadcrumb that copies the whole
request. It runs before the phases because the leaking writer drops its own rows
under the phases' traffic (a queue budget hit, a recursion limit on the nested
copies), which would turn the size check into a missing-row check.
The stored-request half of that regression (every failing request's spend log carrying a
copy of the request per failed attempt) is a deterministic size check, so it lives in
tests/integration/spend/test_failing_request_stored_snapshot.py.
A second, cheaper check holds the idle footprint: every worker's RSS as the harness
A cheaper check holds the idle footprint: every worker's RSS as the harness
read it at collection time, before this pytest process sent any traffic (see
conftest.pytest_collection_finish), must sit under a fixed budget. On the
release gate that is a fresh stack right after its readiness gate, one worker per
@ -74,7 +68,6 @@ from typing import Final
import pytest
from complexity_router_client import ComplexityRouterClient
from e2e_config import (
MEMORY_CONCURRENCY,
MEMORY_IDLE_RSS_BUDGET_MB,
@ -83,7 +76,6 @@ from e2e_config import (
MEMORY_RSS_CEILING_MB,
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS,
MEMORY_RSS_SETTLE_SAMPLES,
MEMORY_STORED_REQUEST_BUDGET_KB,
MEMORY_TRANSCRIPT_TURNS,
PROXY_REPLICA_URLS,
unique_marker,
@ -91,11 +83,12 @@ from e2e_config import (
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere
from models import ChatMessage, RouterSettingsOverride, SpendLogRow
from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride
from proxy_client import ProxyClient
from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.perf]
REFUSING_MODEL: Final = "openai/gpt-5.5"
DEPLOYMENTS_PER_GROUP: Final = 2
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
@ -110,7 +103,6 @@ class FailedCall:
status_code: int
seconds: float
body_head: str
call_id: str | None
@dataclass(frozen=True, slots=True)
@ -123,8 +115,17 @@ class WorkerGrowth:
return self.after.ram_usage_mb - self.warm.ram_usage_mb
def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
return proxy.create_model(
name,
LiteLLMParamsBody(
model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0
),
)
def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None:
for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
resources.defer(lambda model_id=model_id: proxy.delete_model(model_id))
@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS)
def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall:
started: Final = time.perf_counter()
resp: Final = chat_override(
proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT
resp: Final = proxy.transport.send(
"/chat/completions",
headers=proxy.transport.bearer(key),
json=ReliabilityChatBody(
model=model,
messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")],
max_tokens=512,
router_settings_override=override,
cache={"no-cache": True},
),
)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300])
def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]:
@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall
)
def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float:
assert call.call_id, (
f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}"
)
rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id)
assert rows, (
f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack "
"writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored "
"request outgrew the writer's queue budget"
)
snapshot: Final = rows[0].proxy_server_request
assert snapshot, (
f"spend log {call.call_id} stored no request body, so the stack is not running with "
"general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously"
)
return len(json.dumps(snapshot).encode()) / 1024
class TestReliabilityMemory:
@pytest.mark.covers("reliability.perf.idle_memory.under_slo")
@meta(
@ -309,13 +300,13 @@ class TestReliabilityMemory:
Subject(
domain=Domain.ROUTING,
providers=(Provider.OPENAI,),
models=(REAL_MODEL,),
models=(REFUSING_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_failing_requests_do_not_grow_rss_or_stored_request(
def test_failing_requests_do_not_grow_rss(
self,
client: ComplexityRouterClient,
proxy: ProxyClient,
resources: ResourceManager,
scoped_key: str,
record_property: Callable[[str, object], None],
@ -323,35 +314,25 @@ class TestReliabilityMemory:
marker: Final = unique_marker()
primary: Final = f"reliability-memory-{marker}"
fallback: Final = f"reliability-memory-fb-{marker}"
_register_refusing_group(client.proxy, resources, primary)
_register_refusing_group(client.proxy, resources, fallback)
_register_refusing_group(proxy, resources, primary)
_register_refusing_group(proxy, resources, fallback)
override: Final = RouterSettingsOverride(
num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}]
)
probe: Final = _fail_once(client.proxy, scoped_key, primary, override)
_assert_every_call_failed_through_fallback((probe,), fallback)
stored_kb: Final = _stored_request_kb(client.proxy, probe)
assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, (
f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the "
f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with "
f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole "
f"request into the stored snapshot the way the v1.100.0 ones did"
)
warm: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,
fallback,
record_property,
_settled_rss_per_worker(client.proxy),
_settled_rss_per_worker(proxy),
1,
WARMUP_BATCHES,
)
after: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,

View file

@ -33,7 +33,6 @@ from transport import HttpTransport
pytestmark = [
pytest.mark.perf,
pytest.mark.quiet_stack,
pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"),
]

View file

@ -11,7 +11,7 @@ from __future__ import annotations
import os
import time
import warnings
from collections.abc import Callable, Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field, replace
from datetime import datetime
from functools import reduce
@ -155,6 +155,21 @@ class NotServableOn:
last_result: Result[ModelsListResponse] | None
@dataclass(frozen=True, slots=True)
class ModelRegistered:
model_name: str
model_id: str
@dataclass(frozen=True, slots=True)
class ModelNotRegistered:
model_name: str
reason: str
ModelRegistration = ModelRegistered | ModelNotRegistered
def await_servable(
list_models: Callable[[float], Result[ModelsListResponse]],
*,
@ -461,6 +476,7 @@ class ProxyClient:
model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS
model_servable_interval: float = MODEL_SERVABLE_INTERVAL
model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
settle: Callable[[float], None] = settle_propagation
def with_caller(self, caller: Caller) -> ProxyClient:
return replace(self, caller=caller)
@ -694,31 +710,59 @@ class ProxyClient:
balancer address is configured (every request opens a fresh connection, so
the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is
what makes the model safe to use anywhere."""
model_id = unwrap(
self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
).model_id
model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id
written_at = time.monotonic()
try:
self._await_model_servable(body.model_name, listed_for)
except BaseException:
self.delete_model(model_id)
raise
settle_propagation(written_at)
self.settle(written_at)
return model_id
@step("Add a batch of deployments, then wait once for all of them to propagate")
def register_models(
self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False
) -> tuple[ModelRegistration, ...]:
posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies)
written_at: Final = time.monotonic()
registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted)
self.settle(written_at)
return registrations
def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]:
return self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration:
if not isinstance(posted, Success):
return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}")
match self._servable_outcome(model_name, None):
case Servable():
return ModelRegistered(model_name=model_name, model_id=posted.data.model_id)
case NotServableOn() as missing:
self.delete_model(posted.data.model_id)
return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing))
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
"""Block until every replica lists `model_name`, or fail at model_servable_timeout."""
match self._servable_outcome(model_name, listed_for):
case Servable():
return
case NotServableOn() as missing:
raise AssertionError(self._not_servable_message(model_name, missing))
def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn:
headers: Final = self.management_headers(listed_for)
outcome: Final = await_servable_everywhere(
return await_servable_everywhere(
{url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
model_name=model_name,
timeout=self.model_servable_timeout,
@ -728,19 +772,15 @@ class ProxyClient:
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Servable():
return
case NotServableOn(replica=replica, last_result=last_result):
raise AssertionError(
servable_timeout_message(
model_name=model_name,
replica=replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=last_result,
)
)
def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str:
return servable_timeout_message(
model_name=model_name,
replica=missing.replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=missing.last_result,
)
@staticmethod
def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller:

View file

@ -13,7 +13,6 @@ markers =
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set
quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish
provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set
owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
@ -21,3 +20,4 @@ markers =
secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it
proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set

View file

@ -17,11 +17,6 @@ from models import (
ChatBody,
ChatMessage,
ChatResponse,
ConfigPatchResponse,
CostDiscountConfig,
CostDiscountConfigResponse,
CostMarginConfig,
CostMarginConfigResponse,
GuardrailCreateBody,
GuardrailCreateResponse,
KeyDeleteBody,
@ -31,14 +26,11 @@ from models import (
)
from spend_e2e_client import SpendClient
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = pytest.mark.e2e
BACKEND_MODEL: Final = "openai/gpt-6.1-sol"
INPUT_RATE: Final = 0.00005
OUTPUT_RATE: Final = 0.0001
DISCOUNT: Final = 0.25
MARGIN_PERCENT: Final = 0.1
MARGIN_FIXED: Final = 0.0005
def _register_bedrock_guardrail(
@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea
)
def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
json=CostDiscountConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
json=CostMarginConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _get_discount(client: SpendClient) -> dict[str, float]:
return unwrap(
client.proxy.transport.get(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostDiscountConfigResponse,
)
).values
def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
return unwrap(
client.proxy.transport.get(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostMarginConfigResponse,
)
).values
def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str:
model: Final = f"{prefix}-{unique_marker()}"
model_id: Final = client.proxy.create_model(
@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str:
return key
@pytest.fixture
def restored_pricing_config(client: SpendClient) -> Iterator[None]:
discount: Final = _get_discount(client)
margin: Final = _get_margin(client)
try:
_set_discount(client, {})
reset_at: Final = _set_margin(client, {})
settle_propagation(reset_at)
yield
finally:
try:
_set_discount(client, discount)
finally:
restored_at: Final = _set_margin(client, margin)
settle_propagation(restored_at)
class TestPricingConfigSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.discount_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_discount_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT})
settle_propagation(discount_set_at)
model: Final = _register_model(client, strict_resources, "discount-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT)
assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT)
assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT))
@pytest.mark.covers(
"quota_management.spend_tracking.margin_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_margin_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
margin_set_at: Final = _set_margin(
client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}
)
settle_propagation(margin_set_at)
model: Final = _register_model(client, strict_resources, "margin-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED
assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT)
assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED)
assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin)
assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin)
class TestGuardrailCostSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.guardrail_cost.logs_cost",
exercised_on=["chat_completions"],
@ -311,7 +164,6 @@ class TestPricingConfigSpend:
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"]

View file

@ -1,45 +0,0 @@
"""Cross-process reader/writer lock over the proxy stack every xdist worker shares.
Every collected test holds it shared, marker or not, since the Claude Code cells and
other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive,
and the `gate` file makes a waiting exclusive holder win over readers that arrive
after it."""
from __future__ import annotations
import fcntl
import hashlib
import tempfile
from collections.abc import Generator
from contextlib import ExitStack, contextmanager
from pathlib import Path
from typing import Final
from e2e_config import PROXY_BASE_URL
STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12]
LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}"
GATE_FILE: Final = LOCK_DIR / "gate"
STACK_FILE: Final = LOCK_DIR / "stack"
@contextmanager
def _flock(path: Path, operation: int) -> Generator[None]:
with path.open("a") as handle:
fcntl.flock(handle, operation)
try:
yield
finally:
fcntl.flock(handle, fcntl.LOCK_UN)
@contextmanager
def stack_lock(exclusive: bool) -> Generator[None]:
LOCK_DIR.mkdir(parents=True, exist_ok=True)
if exclusive:
with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX):
yield
return
with ExitStack() as held:
with _flock(GATE_FILE, fcntl.LOCK_SH):
held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH))
yield

View file

@ -1,6 +1,6 @@
# e2e harness tests
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter, and `codex/` covers how the Codex driver reads the CLI's JSONL events and the compactions in a thread's rollout file. Put a new harness test under the folder that mirrors the suite folder whose module it covers

View file

@ -8,7 +8,29 @@ from typing import Final
import pytest
from setup.run import run_until
from setup.wire import WireRecord
from test_upgrade import record
def record(seq: int = 0, *, started_at: float = 100) -> WireRecord:
return WireRecord(
seq=seq,
marker=f"request-{seq}",
caller="key",
endpoint="chat_completions",
stream=False,
blocked=False,
started_at=started_at,
status=200,
header_s=1,
first_byte_s=1,
first_token_s=None,
total_s=1,
text="answer",
ok=True,
mock_delay_s=1,
call_id=f"call-{seq}",
error=None,
response_id=f"response-{seq}",
)
class StopAfterRecord:

View file

@ -1,375 +0,0 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
import pytest
from cases.small_bodies_5rps import CASE
from memory_readings import RssCapture, RssReading
from setup.observe import Observation, evaluate
from setup.upgrade import (
MAX_TRAFFIC_SECONDS,
UpgradeMonitor,
matched_response_ids,
postgres_failures,
upgrade_case,
upgrade_properties,
)
from setup.upgrade_worker import UpgradeState
from setup.wire import PlanProgress, WireRecord, response_id
def record(seq: int = 0, *, started_at: float = 100) -> WireRecord:
return WireRecord(
seq=seq,
marker=f"request-{seq}",
caller="key",
endpoint="chat_completions",
stream=False,
blocked=False,
started_at=started_at,
status=200,
header_s=1,
first_byte_s=1,
first_token_s=None,
total_s=1,
text="answer",
ok=True,
mock_delay_s=1,
call_id=f"call-{seq}",
error=None,
response_id=f"response-{seq}",
)
def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None:
records: Final = (record(0), record(1), record(2))
rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call"))
assert matched_response_ids(records, rows) == frozenset({"response-0"})
def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState:
return UpgradeState.model_validate(
{
"phase": phase,
"started_at": 0,
"required_indexes": ["index"],
"required_migrations": ["migration"],
"pending_indexes": ["index"] if pending else [],
"baseline_rows": {"seed-row": "hash"},
}
)
class Probe:
def __init__(self) -> None:
self.now = 0.0
self.current = state()
self.builds: list[int] = []
self.found: frozenset[str] = frozenset()
self.releases = 0
self.metric_checks = 0
self.errors = False
def check_metrics(self) -> None:
self.metric_checks += 1
assert not self.errors, "Failed PostgreSQL operation"
def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]:
return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id)
def release(self) -> None:
self.releases += 1
def monitor(self) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: self.current,
check_metrics=self.check_metrics,
read_activity=lambda: (self.builds, 2),
read_persisted=self.persisted,
start_indexes=self.release,
clock=lambda: self.now,
wall_clock=lambda: self.now,
)
def test_stop_drains_issued_requests_without_issuing_the_rest() -> None:
progress: Final = PlanProgress(100)
assert progress.take(stop=False)
assert progress.take(stop=False)
progress.completed = 1
assert not progress.take(stop=True)
assert not progress.drained
assert not progress.take(stop=False)
progress.completed = 2
assert progress.drained
assert progress.issued == 2
def test_finite_plan_still_waits_for_its_last_response() -> None:
progress: Final = PlanProgress(1)
assert progress.take(stop=False)
assert not progress.take(stop=False)
assert not progress.drained
progress.completed = 1
assert progress.drained
@pytest.mark.parametrize(
"payload,expected",
[
({"id": "chat-id"}, "chat-id"),
({"type": "response.created", "response": {"id": "response-id"}}, "response-id"),
({"type": "message_start", "message": {"id": "message-id"}}, "message-id"),
({"type": "response.output_item.added", "item": {"id": "item-id"}}, None),
({"type": "content_block_delta", "delta": {"text": "hello"}}, None),
],
)
def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None:
assert response_id(payload) == expected
def metrics(*, failure: str = "0", batch: str = "0") -> str:
return f"""# TYPE litellm_postgres_failed_requests_total counter
litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure}
# TYPE litellm_batch_write_to_db_failed_requests_total counter
litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch}
"""
def test_failed_sql_and_async_batch_writes_are_both_counted() -> None:
assert postgres_failures(metrics(failure="2", batch="3")) == 5
assert postgres_failures(metrics()) == 0
def test_no_error_label_series_is_a_valid_zero() -> None:
assert (
postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter
# TYPE litellm_batch_write_to_db_failed_requests_total counter
""")
== 0
)
@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"])
def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None:
total: Final = """# TYPE litellm_postgres_total_requests_total counter
litellm_postgres_total_requests_total{postgres="postgres"} 1
"""
failure: Final = (
f"# TYPE litellm_{service}_failed_requests_total counter\n"
f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n'
if service is not None
else ""
)
assert postgres_failures(total + failure) == (1 if service is not None else 0)
@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"])
def test_missing_failure_families_require_observed_database_activity(total: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(
"# TYPE litellm_postgres_total_requests_total counter\n"
f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n'
)
@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")])
def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(body)
def test_index_completion_stops_after_the_post_upgrade_window() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll((record(),))
assert probe.releases == 1
probe.builds = [17]
probe.now = 101
assert not monitor.poll(())
probe.found = frozenset({"response-0", "response-1"})
probe.now = 103
assert not monitor.poll((record(1, started_at=102),))
assert monitor.overlap_writes == 1
assert monitor.persisted_count == 2
assert not monitor.pending
assert monitor.max_lock_waiters == 2
probe.current = state("complete")
assert not monitor.poll(())
probe.now = 162
assert not monitor.poll(())
probe.now = 163
assert not monitor.poll(())
probe.found = frozenset(f"response-{i}" for i in range(300))
assert monitor.poll(tuple(record(i) for i in range(2, 300)))
assert probe.releases == 1
assert probe.metric_checks == 7
def test_http_200_without_spend_row_cannot_complete() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
_ = monitor.poll(records)
with pytest.raises(AssertionError, match="no matching spend log"):
monitor.assert_complete(state("complete", pending=False))
def test_no_pending_index_work_still_requires_300_persisted_requests() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
probe.found = frozenset(f"response-{i}" for i in range(300))
_ = monitor.poll(records)
monitor.assert_complete(state("complete", pending=False))
assert monitor.successes == monitor.persisted_count == 300
with pytest.raises(AssertionError, match="during index work"):
monitor.assert_complete(state("complete", pending=True))
def test_postgres_failure_fails_even_when_inference_succeeded() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.errors = True
with pytest.raises(AssertionError, match="PostgreSQL"):
_ = monitor.poll((record(),))
def test_failed_migration_and_operational_timeout_are_not_success() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.current = state("failed")
with pytest.raises(AssertionError, match="worker failed"):
_ = monitor.poll(())
probe.current = state()
probe.now = MAX_TRAFFIC_SECONDS
with pytest.raises(AssertionError, match="operational timeout"):
_ = monitor.poll(())
def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None:
case: Final = upgrade_case()
assert case.traffic.rps == 5
assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection(
limit.metric for limit in case.limits
)
assert case.limits == tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
)
def failed_record(seq: int) -> WireRecord:
return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None})
def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
assert probe.releases == 0
assert monitor.processed == 3
assert monitor.successes == 0
assert not monitor.pending
probe.found = frozenset(f"response-{i}" for i in range(3, 303))
assert not monitor.poll(tuple(record(i) for i in range(3, 303)))
monitor.assert_complete(state("complete", pending=False))
assert probe.releases == 1
assert monitor.processed == 303
assert monitor.successes == monitor.persisted_count == 300
def test_fourth_non_200_fails_across_poll_batches() -> None:
monitor: Final = Probe().monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
with pytest.raises(AssertionError, match="Too many unexpected non-200"):
_ = monitor.poll((failed_record(3),))
@pytest.mark.parametrize(
"failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})]
)
def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None:
with pytest.raises(AssertionError, match="Failed or empty inference"):
_ = Probe().monitor().poll((failed,))
def _ceiling(metric: str, percentile: int | None = None) -> float:
return next(
limit.max
for limit in upgrade_case().limits
if limit.metric == metric and limit.percentile == percentile and limit.max is not None
)
@pytest.mark.parametrize(
"over,expected",
[
("", ()),
("unexpected_status_count", ("unexpected_status_count",)),
("gateway_added_seconds", ("gateway_added_seconds percentile 50",)),
("time_to_first_token_seconds", ("time_to_first_token_seconds percentile 99",)),
("memory_growth_megabytes", ("memory_growth_megabytes",)),
],
)
def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples(
over: str, expected: tuple[str, ...]
) -> None:
"""Each sample sits exactly on the case's ceiling, so only the one nudged past it goes over."""
failures: Final = int(_ceiling("unexpected_status_count")) + (over == "unexpected_status_count")
overhead: Final = _ceiling("gateway_added_seconds", 50) * (1.01 if over == "gateway_added_seconds" else 1)
ttft: Final = _ceiling("time_to_first_token_seconds", 99) * (1.01 if over == "time_to_first_token_seconds" else 1)
growth: Final = _ceiling("memory_growth_megabytes") + (over == "memory_growth_megabytes")
nonstream: Final = tuple(record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100))
stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft})
observed: Final = Observation(
records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))),
mock_requests=(),
guardrail_calls=(),
rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()),
rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()),
wall_seconds=60,
)
rows, exceeded = evaluate(upgrade_case(), observed)
assert tuple(row.label for row in rows if row.over or row.missing) == expected
assert len(exceeded) == len(expected)
def test_upgrade_properties_carry_the_upgrade_timings_and_what_the_monitor_saw() -> None:
state: Final = UpgradeState(
phase="complete",
started_at=0,
schema_seconds=12.5,
index_seconds=240.25,
required_indexes=("idx_a", "idx_b"),
required_migrations=("20260101_a",),
pending_indexes=("idx_b",),
baseline_rows={},
)
monitor: Final = UpgradeMonitor(
state=lambda: state,
check_metrics=lambda: None,
read_activity=lambda: ([], 0),
read_persisted=lambda records: frozenset(),
start_indexes=lambda: None,
)
monitor.overlap_writes = 37
monitor.max_lock_waiters = 3
monitor.pending = (record(0), record(1))
assert dict(upgrade_properties(state, monitor)) == {
"perf_upgrade_schema_seconds": 12.5,
"perf_upgrade_index_seconds": 240.25,
"perf_upgrade_overlap_writes": 37,
"perf_upgrade_max_lock_waiters": 3,
"perf_upgrade_pending_indexes": 1,
"perf_upgrade_unmatched_requests": 2,
}

View file

@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything.
from __future__ import annotations
import json
import time
from builtins import ExceptionGroup
from collections.abc import Callable, Generator, Iterable, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import chain, repeat
from queue import SimpleQueue
@ -40,7 +41,10 @@ from models import (
LiteLLMParamsBody,
McpServerCreateBody,
McpServerUpdateBody,
ModelDeleteBody,
ModelInfoBody,
ModelListEntry,
ModelNewBody,
ModelsListResponse,
OrgNewBody,
OrgUpdateBody,
@ -59,6 +63,8 @@ from proxy_client import (
ConvergeOutcome,
CredentialKind,
EverywhereConverged,
ModelNotRegistered,
ModelRegistered,
ModelsPoller,
NeverConvergedOn,
NotConverged,
@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation:
)
assert json.loads(bodies.get_nowait()) == expected
assert bodies.empty()
@contextmanager
def model_registry(
listed: frozenset[str], refused: frozenset[str]
) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]:
writes: Final[SimpleQueue[float]] = SimpleQueue()
deleted: Final[SimpleQueue[str]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
def log_message(self, format: str, *args: object) -> None:
pass
def reply(self, status: int, body: object) -> None:
self.send_response(status)
self.end_headers()
self.wfile.write(json.dumps(body).encode())
def do_GET(self) -> None:
self.reply(200, {"data": [{"id": name} for name in sorted(listed)]})
def do_POST(self) -> None:
raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
if self.path == "/model/delete":
deleted.put(ModelDeleteBody.model_validate_json(raw).id)
self.reply(200, {})
return
writes.put(time.monotonic())
name: Final = ModelNewBody.model_validate_json(raw).model_name
if name in refused:
self.reply(500, {"error": "refused"})
return
self.reply(200, {"model_id": f"id-{name}"})
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
thread.start()
url: Final = f"http://127.0.0.1:{server.server_port}"
proxy: Final = build_proxy_client(
base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap"
)
try:
yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
class TestRegisterModels:
def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None:
names: Final = ("first", "refused", "ghost", "last")
settled: Final[SimpleQueue[float]] = SimpleQueue()
with (
model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as (
proxy,
writes,
deleted,
),
without_retries(),
):
registrations: Final = replace(proxy, settle=settled.put).register_models(
tuple(
ModelNewBody(
model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody()
)
for name in names
)
)
last_write: Final = max(writes.get_nowait() for _ in names)
assert writes.empty()
assert settled.get_nowait() >= last_write
assert settled.empty(), "a batch must wait out propagation once, not once per model"
assert tuple((type(r), r.model_name) for r in registrations) == (
(ModelRegistered, "first"),
(ModelNotRegistered, "refused"),
(ModelNotRegistered, "ghost"),
(ModelRegistered, "last"),
)
assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last")
assert deleted.get_nowait() == "id-ghost"
assert deleted.empty()

View file

@ -1,117 +0,0 @@
"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for
every reader and keeps them out, and a reader arriving behind a waiting exclusive holder
queues behind it instead of starving it."""
from __future__ import annotations
import fcntl
import inspect
import os
import subprocess
import sys
import time
from contextlib import ExitStack
from pathlib import Path
from typing import Final
import pytest
import stack_lock
HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent
DEADLINE_SECONDS: Final = 30.0
SETTLE_SECONDS: Final = 0.5
HOLDER_SCRIPT: Final = """
import sys, time
from pathlib import Path
from stack_lock import stack_lock
name, mode, release_path, log_path = sys.argv[1:]
def record(event):
with Path(log_path).open("a") as log:
log.write(f"{name} {event}\\n")
record("waiting")
with stack_lock(exclusive=mode == "exclusive"):
record("enter")
while not Path(release_path).exists():
time.sleep(0.02)
record("exit")
"""
def _events(log_path: Path) -> tuple[str, ...]:
return tuple(log_path.read_text().splitlines()) if log_path.exists() else ()
def _wait_for_event(log_path: Path, event: str) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
while event not in _events(log_path):
if time.monotonic() > deadline:
pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}")
time.sleep(0.02)
def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
with gate_path.open("a") as handle:
while True:
try:
fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB)
except BlockingIOError:
return
fcntl.flock(handle, fcntl.LOCK_UN)
if time.monotonic() > deadline:
pytest.fail("no exclusive holder ever took the gate")
time.sleep(0.02)
def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]:
holder: Final = held.enter_context(
subprocess.Popen(
(
sys.executable,
"-P",
"-c",
HOLDER_SCRIPT,
name,
mode,
str(tmp_path / f"release-{name}"),
str(tmp_path / "events"),
),
cwd=HARNESS_DIR,
env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)},
)
)
held.callback(holder.kill)
return holder
def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None:
lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}"
lock_dir.mkdir()
log_path: Final = tmp_path / "events"
with ExitStack() as held:
first_reader: Final = _start_holder(held, tmp_path, "A", "shared")
_wait_for_event(log_path, "A enter")
second_reader: Final = _start_holder(held, tmp_path, "R", "shared")
_wait_for_event(log_path, "R enter")
(tmp_path / "release-R").touch()
_wait_for_event(log_path, "R exit")
writer: Final = _start_holder(held, tmp_path, "W", "exclusive")
_wait_until_gate_is_held_exclusively(lock_dir / "gate")
late_reader: Final = _start_holder(held, tmp_path, "B", "shared")
_wait_for_event(log_path, "B waiting")
time.sleep(SETTLE_SECONDS)
(tmp_path / "release-A").touch()
_wait_for_event(log_path, "W enter")
(tmp_path / "release-W").touch()
_wait_for_event(log_path, "B enter")
(tmp_path / "release-B").touch()
for holder in (first_reader, second_reader, writer, late_reader):
assert holder.wait(timeout=DEADLINE_SECONDS) == 0
events: Final = _events(log_path)
assert events.index("R enter") < events.index("A exit")
assert events.index("W enter") > events.index("A exit")
assert events.index("B enter") > events.index("W exit")

View file

@ -1,610 +0,0 @@
from __future__ import annotations
import asyncio
import json
import os
import subprocess
import sys
from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from decimal import Decimal
from itertools import chain
from pathlib import Path
from typing import TYPE_CHECKING, Final, TypeAlias
import psycopg
import pytest
import pytest_asyncio
from psycopg import sql
from pydantic import JsonValue
from litellm.accounting import (
CostEvent,
DeliveryState,
DeliveryTarget,
EventKind,
InboundUsageEvent,
Metric,
MetricName,
OutboundUsageEvent,
Outcome,
PricingSnapshot,
)
from litellm.proxy.db import accounting_ledger
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite
from litellm.types.utils import CallTypes
from tests.integration._support.database import scratch_database
from tests.integration._support.prisma_clients import connected_proxy_prisma
if TYPE_CHECKING:
from prisma import Prisma
from litellm.proxy.utils import PrismaClient
_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger]
_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma"
_EVENT_INSERT: Final = (
'INSERT INTO "LiteLLM_AccountingEvent" '
"(event_id, kind, schema_version, payload_hash, occurred_at, payload) "
"VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version"
)
_DELIVERY_INSERT: Final = (
'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) '
"VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts"
)
@pytest.fixture(scope="module")
def ledger_database() -> Iterator[str]:
with scratch_database("accounting") as database_url:
result: Final = subprocess.run(
[sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)],
env={**os.environ, "DATABASE_URL": database_url},
capture_output=True,
text=True,
timeout=300,
)
assert result.returncode == 0, result.stdout + result.stderr
yield database_url
@pytest.fixture
def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]:
with psycopg.connect(ledger_database) as connection:
try:
yield connection
finally:
connection.rollback()
@pytest.mark.parametrize("kind", tuple(EventKind))
def test_ledger_accepts_every_public_event_kind(
ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind
) -> None:
row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone()
assert row == (kind.value, 1)
@pytest.mark.parametrize("target", tuple(DeliveryTarget))
@pytest.mark.parametrize("state", tuple(DeliveryState))
def test_delivery_accepts_every_public_target_and_state(
ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState
) -> None:
row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone()
assert row == (target.value, 1, state.value, 0)
@pytest.mark.parametrize(
("enum_name", "values"),
[
("AccountingEventKind", tuple(value.value for value in EventKind)),
("AccountingDeliveryState", tuple(value.value for value in DeliveryState)),
("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)),
],
)
def test_migrated_enums_match_the_public_contract(
ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...]
) -> None:
rows: Final = ledger.execute(
"SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s",
(enum_name,),
).fetchall()
assert frozenset(row[0] for row in rows) == frozenset(values)
@pytest.mark.parametrize("version", [0, -1])
def test_event_rejects_nonpositive_schema_versions(
ledger: psycopg.Connection[tuple[object, ...]], version: int
) -> None:
with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"):
ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version))
@pytest.mark.parametrize(
("version", "attempts", "constraint"),
[
(0, 0, "AccountingDelivery_version_check"),
(-1, 0, "AccountingDelivery_version_check"),
(1, -1, "AccountingDelivery_attempts_check"),
],
)
def test_delivery_rejects_invalid_versions_and_attempts(
ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str
) -> None:
with pytest.raises(psycopg.errors.CheckViolation, match=constraint):
ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts))
@pytest.mark.parametrize("column", ["state", "target"])
def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None:
ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0))
with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"):
ledger.execute(
sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",)
)
def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None:
with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"):
ledger.execute(_EVENT_INSERT, ("invalid", 1))
def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None:
row: Final = ledger.execute(
'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) '
"VALUES ('event', %s, 1) RETURNING state::text, attempts",
(DeliveryTarget.KEY_SPEND.value,),
).fetchone()
assert row == ("pending", 0)
_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc)
_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND)
def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent:
fields: Final = dict(
event_id=event_id,
request_id="request",
call_type=CallTypes.completion,
started_at=_AT,
ended_at=_AT,
outcome=Outcome.SUCCEEDED,
dimensions=(),
)
if inbound:
return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),))
return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None)
def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent:
return CostEvent(
event_id=event_id,
usage_event_id=source,
occurred_at=_AT,
dimensions=(),
charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),),
snapshot=PricingSnapshot(
calculator_version="fixture",
model="fixture",
provider=None,
pricing_source="cost_map",
priced_model_key="fixture",
total_cost=Decimal(amount),
),
)
def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str:
return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}})
async def _insert(
store: AccountingLedger,
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...] = _TARGETS,
*,
context: Mapping[str, JsonValue] | None = None,
) -> str:
return await store.insert_event(event, targets, envelope=_envelope(event, context=context))
def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]:
with psycopg.connect(database_url) as connection:
cursor: Final = connection.execute(statement, parameters)
return cursor.fetchall() if cursor.description is not None else []
def _complete_and_age(database_url: str) -> None:
with psycopg.connect(database_url) as connection:
connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,))
connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,))
@pytest_asyncio.fixture(loop_scope="function")
async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]:
monkeypatch.setenv("DATABASE_URL", ledger_database)
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
_rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"')
async with connected_proxy_prisma(ledger_database) as client:
yield client, AccountingLedger.for_client(client)
@pytest.mark.asyncio
async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None:
_, store = writer
event: Final = _usage()
assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted"
rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"')
assert rows == [
("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"})))
]
deliveries: Final = _rows(
ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"'
)
assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS}
@pytest.mark.asyncio
async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None:
_, store = writer
event: Final = _usage(inbound=True)
results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8)))
assert results.count("inserted") == 1 and results.count("duplicate") == 7
_complete_and_age(ledger_database)
before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target')
reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2)
assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate"
assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
@pytest.mark.asyncio
@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"])
async def test_conflicting_acceptance_cannot_overwrite_or_add_work(
writer: _Writer, ledger_database: str, change: str
) -> None:
_, store = writer
event: Final = _usage()
await _insert(store, event)
altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event
targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS
context: Final = {"override": "0.5"} if change == "context" else {}
options: Final = {change: 2} if change in ("schema_version", "version") else {}
with pytest.raises(AccountingConflictError, match="conflicts"):
await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options)
assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)]
assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == {
(t.value,) for t in _TARGETS
}
@pytest.mark.asyncio
async def test_competing_cost_identities_accept_only_one_authoritative_charge(
writer: _Writer, ledger_database: str
) -> None:
_, store = writer
await _insert(store, _usage())
results: Final = await asyncio.gather(
_insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True
)
assert sum(result == "inserted" for result in results) == 1
assert sum(isinstance(result, AccountingConflictError) for result in results) == 1
accepted: Final = _rows(
ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'"
)
assert len(accepted) == 1
event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"]))
assert await _insert(store, event) == "duplicate"
assert event.snapshot.total_cost == Decimal("0.125")
with pytest.raises(AccountingConflictError):
await _insert(store, _cost(event_id=event.event_id, amount="0.5"))
@pytest.mark.asyncio
@pytest.mark.parametrize("source", ["missing", "inbound", "cost"])
async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None:
_, store = writer
if source == "inbound":
await _insert(store, _usage("source", inbound=True))
if source == "cost":
await _insert(store, _usage())
await _insert(store, _cost(event_id="source"))
with pytest.raises(AccountingConflictError, match="outbound usage"):
await _insert(store, _cost(source="source", event_id="invalid"))
assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [
(0,)
]
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"])
async def test_interruption_rolls_back_or_retry_recognizes_commit(
writer: _Writer, ledger_database: str, failure: str
) -> None:
client, store = writer
@asynccontextmanager
async def interrupted() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
if failure == "before_commit":
raise RuntimeError("injected pre-commit failure")
if failure == "cancellation":
raise asyncio.CancelledError
raise RuntimeError("injected lost acknowledgement")
error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError
with pytest.raises(error):
await _insert(AccountingLedger(interrupted), _usage())
expected: Final = 1 if failure == "after_commit" else 0
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)]
assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted")
@pytest.mark.asyncio
async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None:
from prisma.errors import RawQueryError
_, store = writer
with psycopg.connect(ledger_database) as connection:
connection.execute(
'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)'
)
try:
with pytest.raises(RawQueryError, match="reject_delivery"):
await _insert(store, _usage())
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
finally:
with psycopg.connect(ledger_database) as connection:
connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery')
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"])
async def test_retention_preserves_incomplete_or_recovery_needed_families(
writer: _Writer, ledger_database: str, blocked: str
) -> None:
_, store = writer
await _insert(store, _usage())
if blocked != "unpriced":
await _insert(store, _cost())
_complete_and_age(ledger_database)
with psycopg.connect(ledger_database) as connection:
if blocked in ("pending", "parked"):
connection.execute(
"UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,)
)
if blocked == "recent_cost":
connection.execute(
"UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,)
)
before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id')
assert (
await store.purge_completed(
redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF
)
== 0
)
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before
@pytest.mark.asyncio
async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage())
await _insert(store, _cost())
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
@pytest.mark.asyncio
async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None:
client, store = writer
for name in ("a", "b", "c"):
await _insert(store, _usage(name, inbound=True))
_complete_and_age(ledger_database)
async with client.writer_db.tx() as tx:
await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE")
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("a",),
("c",),
]
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
@pytest.mark.asyncio
async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage("a", inbound=True))
await _insert(store, _usage("b"))
await _insert(store, _usage("c", inbound=True))
_complete_and_age(ledger_database)
_rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'")
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("a",),
("b",),
]
@pytest.mark.asyncio
async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage(), targets=())
event: Final = _cost(amount="0")
assert await _insert(store, event, targets=()) == "inserted"
assert await _insert(store, event, targets=()) == "duplicate"
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
@pytest.mark.asyncio
async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted(
writer: _Writer, ledger_database: str
) -> None:
client, store = writer
await _insert(store, _usage())
_complete_and_age(ledger_database)
inserted: Final = asyncio.Event()
release: Final = asyncio.Event()
@asynccontextmanager
async def paused_commit() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
inserted.set()
await release.wait()
task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost()))
try:
await asyncio.wait_for(inserted.wait(), timeout=5)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0
finally:
release.set()
await task
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("cost",),
("usage",),
]
@pytest.mark.asyncio
async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None:
client, store = writer
await _insert(store, _usage(inbound=True))
_complete_and_age(ledger_database)
@asynccontextmanager
async def interrupted() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
raise RuntimeError("injected cleanup failure")
with pytest.raises(RuntimeError, match="cleanup failure"):
await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF)
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)]
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1
@pytest.mark.asyncio
async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured(
writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch
) -> None:
with scratch_database("replica") as replica_url:
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url)
async with connected_proxy_prisma(ledger_database) as client:
store: Final = AccountingLedger.for_client(client)
await _insert(store, _usage())
assert await _insert(store, _usage()) == "duplicate"
assert await _insert(store, _cost()) == "inserted"
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)]
def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]:
if isinstance(plan, dict):
own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else ()
children: Final = (plan.get("Plan"), *plan.get("Plans", ()))
return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None))
return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else ()
@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES])
def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None:
explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s")
with psycopg.connect(ledger_database) as connection:
connection.execute("SET enable_seqscan = off")
plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone()
assert plan is not None
conditions: Final = _index_conditions(plan[0])
assert any(
name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition)
for name, condition in conditions
), conditions
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"])
async def test_usage_batch_is_atomic_and_replay_preserves_receipts(
writer: _Writer, ledger_database: str, failure: str
) -> None:
client, store = writer
events: Final = (
_usage("caller", inbound=True),
_usage("attempt").model_copy(
update={"metrics": (Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),)}
),
)
writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events)
@asynccontextmanager
async def interrupted() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
if failure == "before_commit":
raise RuntimeError("injected pre-commit failure")
if failure == "after_commit":
raise RuntimeError("injected lost acknowledgement")
if failure == "conflict":
await store.insert_usage_batch((writes[1],))
_complete_and_age(ledger_database)
rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS))
with pytest.raises(AccountingConflictError):
await store.insert_usage_batch(rejected)
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)]
else:
with pytest.raises(RuntimeError, match="injected"):
await AccountingLedger(interrupted).insert_usage_batch(writes)
count: Final = 2 if failure == "after_commit" else 0
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)]
await store.insert_usage_batch(writes)
recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events))
assert len(recovered) == 1
assert recovered[0].event == events[1]
assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope)
assert set(recovered[0].targets) == set(_TARGETS)
_complete_and_age(ledger_database)
before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target')
await store.insert_usage_batch((*writes, *writes))
assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)]
@pytest.mark.asyncio
async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit(
writer: _Writer, ledger_database: str
) -> None:
from prisma.errors import DataError
from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings
_, store = writer
invalid: Final = _usage("invalid")
valid: Final = _usage("valid")
meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2))
rejected: Final = meter.submit(
(UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),)
)
accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),))
try:
with pytest.raises(DataError):
await rejected
await accepted
finally:
await meter.close()
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)]
assert (meter.pending_bytes, meter.pending_records) == (0, 0)

View file

@ -7,7 +7,7 @@ import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from pydantic import BaseModel, JsonValue
from pydantic import BaseModel, ConfigDict, JsonValue
class CostEstimate(BaseModel):
@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
assert actual == pytest.approx(dict(expectations)), rows
class LoggedBreakdown(BaseModel):
model_config = ConfigDict(frozen=True)
original_cost: float
discount_percent: float | None = None
discount_amount: float | None = None
margin_percent: float | None = None
margin_fixed_amount: float | None = None
margin_total_amount: float | None = None
def _logged_breakdown(request_id: str) -> LoggedBreakdown:
rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,))
assert len(rows) == 1, rows
metadata: Final = rows[0]["metadata"]
parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)
return LoggedBreakdown.model_validate(parsed["cost_breakdown"])
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
response: Final = gateway.request(
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga
(streamed_global_request_id, 0.12),
)
)
discounted: Final = _logged_breakdown(discounted_request_id)
assert discounted.original_cost == pytest.approx(0.06), discounted
assert discounted.discount_percent == pytest.approx(0.05), discounted
assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted
margined: Final = _logged_breakdown(provider_margin_request_id)
assert margined.original_cost == pytest.approx(0.06), margined
assert margined.discount_percent == pytest.approx(0.05), margined
assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined
assert margined.margin_percent == pytest.approx(0.1), margined
assert margined.margin_fixed_amount == pytest.approx(0.001), margined
assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined
rejections: Final = (
(

View file

@ -139,6 +139,7 @@ _PROVIDERS: Final = (
False,
"databricks/databricks-openjev-qwen35-4b",
),
_Provider("empiriolabs", "empiriolabs/aplomb-test", "/v1/decisions", "aplomb-test", _API_KEY, False, None),
)
_LOCAL_PROVIDERS: Final = (
pytest.param(

View file

@ -61,10 +61,10 @@ class Rig:
@property
def exact_tag(self) -> str:
return f"usecase-a-{self.marker}"
return f"app-a-{self.marker}"
def usecase_tag(self, value: str) -> str:
return f"usecase-{self.marker}:{value}"
def app_tag(self, value: str) -> str:
return f"app-{self.marker}:{value}"
def tenant_tag(self, value: str) -> str:
return f"tenant-{self.marker}:{value}"
@ -195,7 +195,7 @@ def _routed_deployment(marker: str, provider: Wire, name: str, tag: str) -> dict
def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonValue]:
routed: Final = (
_routed_deployment(marker, provider, "tagged", f"usecase-a-{marker}"),
_routed_deployment(marker, provider, "tagged", f"app-a-{marker}"),
_routed_deployment(marker, provider, "default", "default"),
)
realtime: Final[dict[str, JsonValue]] = {
@ -221,8 +221,8 @@ def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonV
def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]:
return [
{"key": "entrypoint", "value": "a", "tag": f"usecase-a-{marker}"},
{"key": "entrypoint", "tag_prefix": f"usecase-{marker}:"},
{"key": "service", "value": "a", "tag": f"app-a-{marker}"},
{"key": "service", "tag_prefix": f"app-{marker}:"},
{"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": tenant_required},
{"key": "budgeted", "tag_prefix": f"budgeted-{marker}:", "require_budgeted_tag": True},
{"key": "team+id", "tag_prefix": f"team-{marker}:"},
@ -403,7 +403,7 @@ def test_mapped_baggage_tags_reach_spend_and_callbacks_but_not_requester_metadat
ROUTE_PATHS[route],
_inference_body(route, mapped.chat_model, f"baggage {tenant}", caller_tag),
key=key,
headers={"baggage": f"entrypoint=a, tenant={tenant}"},
headers={"baggage": f"service=a, tenant={tenant}"},
)
assert response.status_code == 200, response.text
assert len(mapped.provider.drain()) == 1
@ -423,7 +423,7 @@ def test_json_string_metadata_gains_exactly_the_mapped_tags(optional: Rig) -> No
key: Final = scenario.key()
optional.provider.drain()
plain: Final = _chat(optional, plain_key, None, fields=fields)
response: Final = _chat(optional, key, f"entrypoint=a,tenant={tenant}", fields=fields)
response: Final = _chat(optional, key, f"service=a,tenant={tenant}", fields=fields)
assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text)
assert len(optional.provider.drain()) == 2
assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags
@ -494,7 +494,7 @@ def test_a2a_spend_row_carries_the_mapped_tags(mapped: Rig) -> None:
},
},
key=key,
headers={"baggage": f"entrypoint=a,tenant={tenant}"},
headers={"baggage": f"service=a,tenant={tenant}"},
)
assert sent.status_code == 200, sent.text
assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key)
@ -514,8 +514,8 @@ def test_spent_mapped_tag_budget_rejects_the_next_request_before_the_provider(
mapped: Rig, key_fields: Mapping[str, JsonValue], fields: Mapping[str, JsonValue], headers: Mapping[str, str]
) -> None:
entry: Final = uuid.uuid4().hex
tag: Final = mapped.usecase_tag(entry)
baggage: Final = f"entrypoint={entry},tenant={entry}"
tag: Final = mapped.app_tag(entry)
baggage: Final = f"service={entry},tenant={entry}"
with mapped.proxy.scenario() as scenario:
key: Final = scenario.key(**key_fields)
mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 0.0001})
@ -528,7 +528,7 @@ def test_spent_mapped_tag_budget_rejects_the_next_request_before_the_provider(
assert blocked.status_code == 422, blocked.text
assert "budget" in blocked.text.lower(), blocked.text
assert mapped.provider.drain() == ()
other: Final = _chat(mapped, key, f"entrypoint=other-{entry},tenant={entry}")
other: Final = _chat(mapped, key, f"service=other-{entry},tenant={entry}")
assert other.status_code == 200, other.text
@ -536,10 +536,10 @@ def test_required_mapping_rejects_a_request_without_its_baggage_value_before_the
with mapped.proxy.scenario() as scenario:
key: Final = scenario.key()
mapped.provider.drain()
rejected: Final = _chat(mapped, key, "entrypoint=a")
rejected: Final = _chat(mapped, key, "service=a")
assert rejected.status_code == 400, rejected.text
assert mapped.provider.drain() == ()
admitted: Final = _chat(mapped, key, f"entrypoint=a,tenant={uuid.uuid4().hex}")
admitted: Final = _chat(mapped, key, f"service=a,tenant={uuid.uuid4().hex}")
assert admitted.status_code == 200, admitted.text
@ -591,7 +591,7 @@ def test_mapped_tags_never_select_a_tag_routed_deployment(mapped: Rig, caller_ta
_chat(
mapped,
key,
f"entrypoint=a,tenant={tenant}",
f"service=a,tenant={tenant}",
model=f"routed-{mapped.marker}",
fields=fields,
**headers,
@ -604,8 +604,8 @@ def test_mapped_tags_never_select_a_tag_routed_deployment(mapped: Rig, caller_ta
def test_mapped_tags_never_trigger_a_tag_scoped_policy_attachment(mapped: Rig) -> None:
entry: Final = uuid.uuid4().hex
tag: Final = mapped.usecase_tag(entry)
baggage: Final = f"entrypoint={entry},tenant={entry}"
tag: Final = mapped.app_tag(entry)
baggage: Final = f"service={entry},tenant={entry}"
proxy: Final = mapped.proxy
with wire_server(lambda _request: Reply(body=GUARDRAIL_VERDICT)) as guardrail, proxy.scenario() as scenario:
key: Final = scenario.key()
@ -662,7 +662,7 @@ def test_required_mapping_refuses_websocket_inference_before_the_upstream(mapped
upstream.get("/__observations").raise_for_status()
status: Final = _handshake_status(
f"ws{proxy_url.removeprefix('http')}/v1/realtime?model=realtime-{mapped.marker}",
{"Authorization": f"Bearer {key}", "baggage": f"entrypoint=a,tenant={uuid.uuid4().hex}"},
{"Authorization": f"Bearer {key}", "baggage": f"service=a,tenant={uuid.uuid4().hex}"},
)
assert status == 403, status
observed: Final = list_value(JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"])
@ -706,7 +706,7 @@ def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> No
path,
{"model": mapped.chat_model, "messages": [{"role": "user", "content": tenant}]},
key=key,
headers={"baggage": f"entrypoint=a,tenant={tenant}"},
headers={"baggage": f"service=a,tenant={tenant}"},
)
assert response.status_code == 200, response.text
assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key)
@ -777,7 +777,7 @@ def test_passthrough_forwards_the_client_body_unchanged_while_its_spend_row_carr
plain: Final = optional.proxy.request("POST", path, body, key=plain_key)
plain_forwarded: Final = _forwarded_bodies(optional)
response: Final = optional.proxy.request(
"POST", path, body, key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"}
"POST", path, body, key=key, headers={"baggage": f"service=a,tenant={tenant}"}
)
assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text)
assert plain_forwarded == (_upstream_body(entrypoint, body),)
@ -929,7 +929,7 @@ def test_reject_clientside_metadata_tags_with_baggage(
) -> None:
tenant: Final = uuid.uuid4().hex
fields, headers = _rejected_client_tag(client_tag, f"caller-{tenant}")
baggage: Final = f"entrypoint=a,tenant={tenant}" if baggage_maps else f"region={tenant}"
baggage: Final = f"service=a,tenant={tenant}" if baggage_maps else f"region={tenant}"
mapped_tags: Final = frozenset({rejecting.exact_tag, rejecting.tenant_tag(tenant)} if baggage_maps else ())
with rejecting.proxy.scenario() as scenario:
plain_key: Final = scenario.key()
@ -952,7 +952,7 @@ def test_plus_sign_rule_key_matches_either_wire_spelling_but_not_both(
optional: Rig, baggage_keys: tuple[str, ...], maps_team: bool
) -> None:
value: Final = uuid.uuid4().hex
baggage: Final = ",".join(("entrypoint=a", *(f"{baggage_key}={value}" for baggage_key in baggage_keys)))
baggage: Final = ",".join(("service=a", *(f"{baggage_key}={value}" for baggage_key in baggage_keys)))
with optional.proxy.scenario() as scenario:
key: Final = scenario.key()
response: Final = _chat(optional, key, baggage)
@ -991,7 +991,7 @@ def test_baggage_mapping_a_deployment_tag_keeps_the_team_tag_guardrail_mode_firs
@pytest.mark.parametrize("metadata", ["run-123", {"tags": "run-123"}], ids=["non_object", "string_tags"])
@pytest.mark.parametrize(
("header_tag", "baggage"),
[(True, "region=sg"), (True, "entrypoint=a"), (False, "region=sg"), (False, "entrypoint=a")],
[(True, "region=sg"), (True, "service=a"), (False, "region=sg"), (False, "service=a")],
ids=["header_tag_unmapped_baggage", "header_tag_mapped_baggage", "unmapped_baggage", "mapped_baggage"],
)
def test_malformed_metadata_keeps_its_status_under_baggage(

View file

@ -0,0 +1,84 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
from integration._support.process import owned_proxy
from pydantic import JsonValue
UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1"
TRANSCRIPT_TURNS: Final = 40
RETRIES: Final = 2
STORED_REQUEST_BUDGET_KB: Final = 64
def _config_storing_prompts(directory: Path) -> Path:
config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
general_settings: Final = object_value(config["general_settings"])
path: Final = directory / "store-prompts.yaml"
path.write_text(
yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}})
)
return path
@pytest.fixture(scope="module")
def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with gateway_from_environment() as parent:
directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot")
with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned:
yield owned
def _transcript(marker: str) -> list[JsonValue]:
turns: Final[list[JsonValue]] = [
{"role": role, "content": f"turn {turn // 2} {role}"}
for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS)
]
return [*turns, {"role": "user", "content": marker}]
def _stored_request(gateway: Gateway, call_id: str) -> JsonValue:
rows: Final = eventually(
lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
lambda values: len(values) == 1,
seconds=40,
)
stored: Final = rows[0]["proxy_server_request"]
return json.loads(stored) if isinstance(stored, str) else stored
def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
primary: Final = f"stored-snapshot-{uuid.uuid4().hex}"
fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}"
for name in (primary, primary, fallback, fallback):
scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0)
key: Final = scenario.key(models=[primary, fallback])
marker: Final = f"stored snapshot {uuid.uuid4().hex}"
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": primary,
"max_tokens": 16,
"messages": _transcript(marker),
"router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]},
},
key=key,
)
assert response.status_code != 200, response.text
assert fallback in response.text, response.text
serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"]))
stored_kb: Final = len(serialized.encode()) / 1024
assert marker in serialized, serialized[:400]
assert stored_kb <= STORED_REQUEST_BUDGET_KB, (
f"one failing request stored a {stored_kb:.0f} KB request snapshot for a "
f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback"
)

View file

@ -0,0 +1,431 @@
import json
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Final
import pytest
from pydantic import TypeAdapter
from litellm.accounting.projection import project
from litellm.accounting.snapshot import build_accounting_payload
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.proxy_server import proxy_logging_obj
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import PrismaClient
from litellm.types.utils import ModelResponse, Usage
_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc)
_FIELDS: Final = TypeAdapter(dict[str, object])
@pytest.mark.parametrize("status", ["success", "failure"])
@pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"])
@pytest.mark.parametrize("spend", [0.0, 0.25])
@pytest.mark.asyncio
async def test_daily_projection_matches_legacy_writer(status: str, origin: str | None, spend: float) -> None:
payload: Final = get_logging_payload(
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "projection-request",
"litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user", "status": status}},
},
response_obj=ModelResponse(
id="projection-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
)
metadata: Final = {
**_FIELDS.validate_json(payload["metadata"]),
**({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}),
**({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin not in (None, "billed_retry") else {}),
}
row: Final = {**payload, "spend": spend, "metadata": json.dumps(metadata), "request_duration_ms": 25}
legacy: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction(
row, PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj), "user"
)
assert legacy is not None
accounting_input: Final[SpendLogsPayload] = {**row}
accounting: Final = build_accounting_payload(
accounting_input,
entities={"api_key": payload["api_key"], "user": payload["user"]},
model_access_groups=(),
windows=(),
router=lambda: None,
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend")
assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments}
assert dict(stored_daily.key)["date"] == legacy["date"]
assert dict(stored_daily.key)["api_key"] == legacy["api_key"]
assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting)
def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None:
row: Final[Mapping[str, object]] = {
"request_id": "saved",
"startTime": _NOW,
"api_key": "key",
"user": "user",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"request_tags": '["tag", "tag"]',
"spend": 0.5,
"metadata": json.dumps(
{
"accounting_savings": {
"compression": 1.0,
"prompt_caching": -2.0,
"gateway_injected_caching": 3.0,
"autorouter": 4.0,
},
"autorouter_savings": 9000,
}
),
}
contributions: Final = project(row)
tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_TagTable")
daily_tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_DailyTagSpend")
assert sum(dict(item.increments)["spend"] for item in tags) == 1.0
assert len(daily_tags) == 1
assert {key: value for key, value in daily_tags[0].increments if key.endswith("savings_spend")} == {
"compression_savings_spend": 1.0,
"prompt_caching_savings_spend": -2.0,
"gateway_injected_caching_savings_spend": 3.0,
"autorouter_savings_spend": 4.0,
}
def test_projection_carries_windows_and_does_not_reprice_guardrails() -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"api_key": "key",
"spend": 0.75,
"metadata": json.dumps(
{
"accounting_windows": [
{
"entity_type": "key",
"entity_id": "key",
"window_duration": "1d",
"window_start": _NOW.isoformat(),
}
],
"guardrail_information": [
{
"guardrail_id": "guard",
"guardrail_usage": {"text": 3, "images": 2},
"guardrail_cost_by_unit": {"text": 0.3, "images": None},
}
],
}
),
}
contributions: Final = project(row)
window: Final = next(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend")
units: Final = {
dict(item.key)["usage_unit"]: dict(item.increments)
for item in contributions
if item.table == "LiteLLM_DailyGuardrailUsageUnits"
}
assert dict(window.increments) == {"spend": 0.75}
assert dict(window.values)["window_start"] == _NOW.isoformat()
assert units == {
"text": {"units": 3, "cost": 0.3, "untracked_units": 0},
"images": {"units": 2, "cost": 0.0, "untracked_units": 2},
}
@pytest.mark.asyncio
async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.responses.streaming_iterator import drain_logging_tasks
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
monkeypatch.setattr(proxy_server, "prisma_client", client)
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
monkeypatch.setattr(client, "spend_log_transactions", [])
monkeypatch.setattr(client, "model_usage_transactions", [])
writer: Final = DBSpendUpdateWriter()
charged: Final = await writer.update_database(
token="sk-projection",
user_id="user",
end_user_id="end-user",
team_id="team",
org_id="org",
project_id="project",
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "entity-request",
"agent_id": "agent",
"user": "end-user",
"litellm_params": {
"metadata": {
"user_api_key": "sk-projection",
"user_api_key_user_id": "user",
"user_api_key_team_id": "team",
"user_api_key_org_id": "org",
"user_api_key_project_id": "project",
"user_api_key_end_user_id": "end-user",
"tags": ["tag", "tag"],
},
"user": "end-user",
},
},
completion_response=ModelResponse(
id="entity-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
response_cost=0.25,
)
await drain_logging_tasks()
assert charged
assert len(client.spend_log_transactions) == 1
queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
logged: Final[SpendLogsPayload] = client.spend_log_transactions[0]
projected: Final = project(
_FIELDS.validate_json(
json.dumps(
build_accounting_payload(
logged,
entities={
"api_key": logged["api_key"],
"user": "user",
"end_user": "end-user",
"team_id": "team",
"organization_id": "org",
"project_id": "project",
"agent_id": "agent",
},
model_access_groups=(),
windows=(),
router=lambda: None,
),
default=str,
)
)
)
legacy: Final = {(item["entity_type"].value, item["entity_id"]): item["response_cost"] for item in queued}
assert legacy[("key", proxy_server.hash_token("sk-projection"))] == next(
dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_VerificationToken"
)
for table, entity, identity in (
("LiteLLM_UserTable", "user", "user"),
("LiteLLM_TeamTable", "team", "team"),
("LiteLLM_OrganizationTable", "organization", "org"),
("LiteLLM_EndUserTable", "end_user", "end-user"),
("LiteLLM_AgentsTable", "agent", "agent"),
("LiteLLM_ProjectTable", "project", "project"),
):
assert legacy[(entity, identity)] == next(
dict(item.increments)["spend"] for item in projected if item.table == table
)
assert sum(dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_TagTable") == sum(
item["response_cost"] or 0 for item in queued if item["entity_type"].value == "tag"
)
daily_queues: Final = (
("LiteLLM_DailyUserSpend", writer.daily_spend_update_queue),
("LiteLLM_DailyTeamSpend", writer.daily_team_spend_update_queue),
("LiteLLM_DailyOrganizationSpend", writer.daily_org_spend_update_queue),
("LiteLLM_DailyEndUserSpend", writer.daily_end_user_spend_update_queue),
("LiteLLM_DailyAgentSpend", writer.daily_agent_spend_update_queue),
("LiteLLM_DailyTagSpend", writer.daily_tag_spend_update_queue),
)
for table, queue in daily_queues:
transactions = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
matching = tuple(item for item in projected if item.table == table)
assert matching, table
contribution = matching[0]
assert len(transactions) == 1
transaction = next(iter(transactions.values()))
assert dict(contribution.increments) == {name: transaction[name] for name, _ in contribution.increments}
assert len(client.model_usage_transactions) == 1
model: Final = client.model_usage_transactions[0]
model_projection: Final = next(item for item in projected if item.table == "LiteLLM_DailyModelUsage")
assert dict(model_projection.increments) == {
"spend": model.spend,
"prompt_tokens": model.prompt_tokens,
"completion_tokens": model.completion_tokens,
"request_count": int(model.counts_as_request),
"successful_requests": int(model.successful),
"failed_requests": int(not model.successful),
}
@pytest.mark.parametrize("served_id", ["requested", "fallback", "missing"])
@pytest.mark.asyncio
async def test_projection_preserves_fallback_narrowed_model_access_groups(served_id: str) -> None:
from litellm import Router
from litellm.proxy.db.db_spend_update_writer import debitable_model_access_groups
router: Final = Router(
model_list=[
{
"model_name": "first",
"litellm_params": {"model": "openai/projection-test", "api_key": "unused"},
"model_info": {"id": "requested", "access_groups": ["first-group"]},
},
{
"model_name": "second",
"litellm_params": {"model": "openai/projection-test", "api_key": "unused"},
"model_info": {"id": "fallback", "access_groups": ["second-group", "unattributed"]},
},
]
)
attributed: Final = ("first-group", "second-group", "first-group")
writer: Final = DBSpendUpdateWriter()
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
await writer._update_model_access_group_db(0.5, attributed, served_id, client, router)
queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
frozen_groups: Final = debitable_model_access_groups(attributed, served_id, router)
projected: Final = project(
{"spend": 0.5, "metadata": json.dumps({"accounting_model_access_groups": frozen_groups})}
)
assert {dict(item.key)["access_group_name"]: dict(item.increments)["spend"] for item in projected} == {
item["entity_id"]: item["response_cost"] for item in queued
}
assert "unattributed" not in frozen_groups
def test_guardrail_projection_skips_unmeasured_counters() -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"metadata": json.dumps(
{
"guardrail_information": [
{"guardrail_usage": {"text": 3}},
{
"guardrail_id": "guard",
"guardrail_usage": {"zero": 0, "negative": -1, "bool": True, "unknown": None, "text": 2},
},
]
}
),
}
units: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyGuardrailUsageUnits")
assert len(units) == 1
assert dict(units[0].key)["usage_unit"] == "text"
assert dict(units[0].increments) == {"units": 2, "cost": 0.0, "untracked_units": 2}
def test_lean_row_without_model_does_not_fabricate_daily_entity_usage() -> None:
projected: Final = project(
{
"request_id": "lean",
"api_key": "key",
"startTime": _NOW,
"spend": 0.5,
"metadata": {"accounting_entities": {"api_key": "key"}},
}
)
assert tuple(item.table for item in projected) == ("LiteLLM_VerificationToken", "LiteLLM_DailyModelUsage")
@pytest.mark.parametrize(
"cost_fields",
[
{"guardrail_cost_in_spend": False, "guardrail_cost_by_unit": {"text": 1.0}},
{"guardrail_cost_by_unit": {"text": -1.0}},
],
)
def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: Mapping[str, object]) -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"request_tags": '["task:debugging"]',
"metadata": json.dumps(
{
"guardrail_information": [
{"guardrail_id": "guard", "guardrail_usage": {"text": 2}, **cost_fields},
]
}
),
}
contributions: Final = project(row)
guardrail: Final = next(item for item in contributions if item.table == "LiteLLM_DailyGuardrailUsageUnits")
model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage")
assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2}
assert dict(model.key)["task_type"] == "debugging"
@pytest.mark.parametrize("organization_id", [None, "", "charged-org"])
def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None:
row: Final[Mapping[str, object]] = {
"startTime": _NOW,
"api_key": "key",
"model": "projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"organization_id": "stored-org",
"spend": 0.25,
"metadata": {"accounting_entities": {"organization_id": organization_id}},
}
daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend")
if organization_id is None:
assert daily == ()
else:
assert len(daily) == 1
assert dict(daily[0].key)["organization_id"] == organization_id
assert dict(daily[0].increments)["spend"] == 0.25
@pytest.mark.parametrize("start", [_NOW, "2026-11-02T12:00:00+00:00", "2026-11-02 12:00:00+00:00"])
def test_every_daily_row_uses_the_calendar_date_of_the_request(start: object) -> None:
row: Final[Mapping[str, object]] = {
"request_id": "dated",
"startTime": start,
"api_key": "key",
"user": "user",
"team_id": "team",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"spend": 0.5,
"metadata": "{}",
}
dates: Final = {
item.table: dict(item.key)["date"] for item in project(row) if item.table.startswith("LiteLLM_Daily")
}
assert {"LiteLLM_DailyUserSpend", "LiteLLM_DailyModelUsage"} <= dates.keys()
assert set(dates.values()) == {"2026-11-02"}
def test_unparseable_start_time_skips_only_guardrail_usage() -> None:
projected: Final = project(
{
"request_id": "bad-time",
"startTime": "yesterday at noon",
"spend": 0.5,
"metadata": {
"accounting_entities": {"api_key": "key"},
"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}],
},
}
)
tables: Final = {item.table for item in projected}
assert "LiteLLM_VerificationToken" in tables
assert "LiteLLM_DailyGuardrailUsageUnits" not in tables
def test_guardrail_usage_is_dated_in_utc() -> None:
projected: Final = project(
{
"request_id": "offset",
"startTime": "2026-11-02T22:30:00-05:00",
"metadata": {"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 1}}]},
}
)
assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [
"2026-11-03"
]

View file

@ -19,7 +19,6 @@ from litellm.accounting import (
Dimension,
DimensionFilter,
DimensionName,
EventKind,
Granularity,
InboundUsageEvent,
Metric,
@ -442,7 +441,6 @@ def test_subset_token_counts_round_trip_without_changing_totals() -> None:
@pytest.mark.parametrize(
("enum_name", "python_enum"),
[
("AccountingEventKind", EventKind),
("AccountingDeliveryState", DeliveryState),
("AccountingDeliveryTarget", DeliveryTarget),
],

View file

@ -2766,9 +2766,9 @@ def test_aggregate_batch_cost_uses_custom_model_info():
assert result.cost == pytest.approx(expected), f"Expected total cost {expected}, got {result.cost}"
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.parametrize("data_residency", ["eu", "us"])
@pytest.mark.parametrize("data_residency", ["eu", "us", "fedramp"])
def test_batch_cost_calculator_applies_data_residency_uplift(data_residency, monkeypatch):
"""batch_cost_calculator should apply the regional uplift multiplier when
"""batch_cost_calculator should apply the regional or FedRAMP uplift multiplier when
data_residency is set and the model carries a configured multiplier."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
prev_model_cost = litellm.model_cost

View file

@ -484,55 +484,6 @@ async def test_bridged_text_completion_is_logged_as_the_text_completion_it_serve
assert logged.payload["response"]["choices"][0]["text"] == "pong"
@pytest.mark.asyncio
async def test_a_bridged_chat_response_waits_on_the_usage_ack_like_any_chat_response(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
database: Final = CommitBarrier()
runtime: Final = UsageCaptureRuntime()
runtime.start(database.write, UsageCaptureSettings(batch_size=2))
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True})
logging_obj: Final = Logging(
model="openai/gpt-5.3-codex",
messages=[],
stream=False,
call_type="acompletion",
start_time=datetime.datetime.now(),
litellm_call_id="bridged-chat",
function_id="bridged-chat",
)
with respx.mock(assert_all_called=True) as router:
router.post(url__startswith="https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(200, json=_searched_responses_body("openai/gpt-5.3-codex"))
)
served: Final = await litellm.acompletion(
model="openai/gpt-5.3-codex",
messages=[{"role": "user", "content": "ping"}],
api_key="test-key",
litellm_logging_obj=logging_obj,
)
response: Final = asyncio.create_task(capture_proxy_response(logging_obj, served))
try:
await database.entered.wait()
assert not response.done()
assert database.committed == ()
database.release.set()
await response
assert "bridged-chat:inbound" in database.committed
finally:
database.release.set()
await runtime.stop(1)
await response
def _served_message_start_id(events: Sequence[bytes]) -> object:
lines: Final = b"".join(events).decode().splitlines()
starts: Final = tuple(

View file

@ -132,6 +132,7 @@ _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = (
"https://openrouter.ai/api/alpha/decisions",
"typesafe/jev-1.13",
),
("empiriolabs", "empiriolabs/aplomb-1", "https://api.empiriolabs.ai/v1/decisions", "aplomb-1"),
)
_OPENAI_RESPONSE: Final[Mapping[str, object]] = {
@ -1549,6 +1550,7 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset(
LlmProviders.LM_STUDIO.value,
LlmProviders.OLLAMA.value,
LlmProviders.OLLAMA_CHAT.value,
LlmProviders.EMPIRIOLABS.value,
}
)
_DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple(
@ -3449,6 +3451,147 @@ async def test_fireworks_ai_bills_cached_input_tokens_at_the_cache_read_rate(
assert recording_logger.standard_logging_object["response_cost"] < 300 * float(prices["input_cost_per_token"])
_EMPIRIOLABS_MODEL: Final = "empiriolabs/aplomb-1"
_EMPIRIOLABS_DEFAULT_URL: Final = "https://api.empiriolabs.ai/v1/decisions"
_EMPIRIOLABS_STATE: Final = "I was charged twice for order 4411 and want a refund today."
_EMPIRIOLABS_QUESTIONS: Final[Mapping[str, object]] = MappingProxyType(
{"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}
)
_EMPIRIOLABS_RESPONSE: Final[Mapping[str, object]] = {
"model": "aplomb-1",
"answers": {"is_refund": {"type": "noul", "noul": 0.97, "confidence": 0.94}},
"usage": {"input_tokens": 24, "output_tokens": 0},
}
@pytest.mark.parametrize(
("api_base", "environment_api_base", "url"),
(
(None, None, _EMPIRIOLABS_DEFAULT_URL),
(None, "https://environment-gateway.example/v1", "https://environment-gateway.example/v1/decisions"),
(
"https://gateway.example/v1/",
"https://environment-gateway.example/v1",
"https://gateway.example/v1/decisions",
),
("https://gateway.example", None, "https://gateway.example/v1/decisions"),
),
ids=("default_base", "environment_base", "explicit_base_over_environment", "explicit_base_without_v1"),
)
def test_empiriolabs_posts_the_systemone_body_to_v1_decisions_on_the_resolved_base(
api_base: str | None,
environment_api_base: str | None,
url: str,
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
if environment_api_base is not None:
monkeypatch.setenv("EMPIRIOLABS_API_BASE", environment_api_base)
route: Final = respx_mock.post(url).respond(json=_EMPIRIOLABS_RESPONSE)
response: Final = litellm.systemone(
model=_EMPIRIOLABS_MODEL,
state=_EMPIRIOLABS_STATE,
questions=_EMPIRIOLABS_QUESTIONS,
api_base=api_base,
api_key="caller-key",
)
assert route.call_count == 1
assert json.loads(route.calls[0].request.content) == {
"model": "aplomb-1",
"state": _EMPIRIOLABS_STATE,
"questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
}
assert response.answers["is_refund"] == NoulAnswer(type="noul", noul=0.97, confidence=0.94)
assert response._hidden_params["custom_llm_provider"] == "empiriolabs"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("api_key", "authorization"),
((None, "Bearer environment-key"), ("caller-key", "Bearer caller-key")),
ids=("environment_key", "caller_key_over_environment"),
)
async def test_empiriolabs_sends_the_caller_key_or_else_the_environment_key(
api_key: str | None,
authorization: str,
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key")
route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond(json=_EMPIRIOLABS_RESPONSE)
response: Final = await litellm.adecisions(
model=_EMPIRIOLABS_MODEL,
input=_EMPIRIOLABS_STATE,
questions=[{"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"}],
api_key=api_key,
)
assert route.call_count == 1
assert route.calls[0].request.headers["authorization"] == authorization
assert isinstance(response, OpenAIDecisionResponse)
assert [answer.model_dump(mode="json") for answer in response.answers] == [
{"type": "predicate", "name": "is_refund", "probability": 0.97}
]
def test_empiriolabs_without_a_key_is_refused_before_http(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False)
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
with pytest.raises(litellm.AuthenticationError, match="Missing API key for Decisions provider 'empiriolabs'"):
litellm.systemone(model=_EMPIRIOLABS_MODEL, state=_EMPIRIOLABS_STATE, questions=_EMPIRIOLABS_QUESTIONS)
assert len(respx_mock.calls) == 0
@pytest.mark.parametrize(
("status", "error_type", "code", "exception"),
(
(401, "authentication_error", "invalid_api_key", litellm.AuthenticationError),
(402, "insufficient_credits", "insufficient_credits", litellm.PaymentRequiredError),
(404, "not_found_error", "model_not_found", litellm.NotFoundError),
(422, "invalid_request_error", "invalid_request", litellm.BadRequestError),
(429, "rate_limit_error", "rate_limit_exceeded", litellm.RateLimitError),
(503, "service_unavailable", "model_unavailable", litellm.ServiceUnavailableError),
),
)
def test_empiriolabs_upstream_errors_keep_their_status_and_message(
status: int,
error_type: str,
code: str,
exception: type[Exception],
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
message: Final = f"upstream refused the decision with {code}"
route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond(
status,
json={"error": {"message": message, "type": error_type, "code": code}},
headers={"retry-after": "1"},
)
with pytest.raises(exception, match=message) as caught:
litellm.systemone(
model=_EMPIRIOLABS_MODEL,
state=_EMPIRIOLABS_STATE,
questions=_EMPIRIOLABS_QUESTIONS,
api_key="caller-key",
num_retries=0,
)
assert getattr(caught.value, "status_code") == status
assert route.call_count == 1
_BOX_QUESTION_REQUESTS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = (
(
"decisions",

View file

@ -33,6 +33,7 @@ from litellm.types.llms.base import CachedTokensDetails
from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
DataResidency,
ImageObject,
ImageResponse,
ImageUsage,
@ -2579,6 +2580,51 @@ def test_data_residency_applies_uplift(data_residency, model, _local_model_cost_
assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9)
def test_fedramp_data_residency_applies_configured_uplift() -> None:
usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
model_info: Final[ModelInfo] = {
"key": "fedramp-fixture",
"max_tokens": None,
"max_input_tokens": None,
"max_output_tokens": None,
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
"litellm_provider": "openai",
"mode": "chat",
"supported_openai_params": None,
"regional_processing_uplift_multiplier_us": 1.5,
"fedramp_uplift_multiplier": 1.1,
}
base: Final = generic_cost_per_token(
model="fedramp-fixture",
usage=usage,
custom_llm_provider="openai",
model_info=model_info,
)
fedramp: Final = generic_cost_per_token(
model="fedramp-fixture",
usage=usage,
custom_llm_provider="openai",
data_residency=DataResidency.FEDRAMP.value,
model_info=model_info,
)
us: Final = generic_cost_per_token(
model="fedramp-fixture",
usage=usage,
custom_llm_provider="openai",
data_residency=DataResidency.US.value,
model_info=model_info,
)
assert sum(base) > 0
assert fedramp[0] == pytest.approx(base[0] * 1.1)
assert fedramp[1] == pytest.approx(base[1] * 1.1)
assert us[0] == pytest.approx(base[0] * 1.5)
assert us[1] == pytest.approx(base[1] * 1.5)
@pytest.mark.parametrize("model", ["gpt-5", "gpt-5-mini", "gpt-5-nano", "gpt-5-pro", "gpt-4o", "gpt-4.1"])
def test_data_residency_no_uplift_for_pre_march_2026_models(model, _local_model_cost_map):
"""Models released before 2026-03-05 must not have the regional uplift."""
@ -2719,6 +2765,22 @@ def test_vertex_uplift_invalid_multiplier_defaults_to_one():
)
@pytest.mark.parametrize(
"model_info, data_residency",
[
({"fedramp_uplift_multiplier": "not-a-number"}, "fedramp"),
({}, "fedramp"),
({"fedramp_uplift_multiplier": 1.1}, "us"),
({"fedramp_uplift_multiplier": 1.1}, None),
],
)
def test_fedramp_uplift_defaults_to_one(model_info: ModelInfo, data_residency: str | None):
"""A malformed or missing FedRAMP multiplier, or a non-FedRAMP host, bills base pricing and never raises."""
from litellm.litellm_core_utils.llm_cost_calc.utils import get_fedramp_uplift_multiplier
assert get_fedramp_uplift_multiplier(model_info, data_residency) == 1.0
def test_service_tier_suffixes_constant_in_sync_with_enum():
from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES
from litellm.types.utils import ServiceTier

View file

@ -340,6 +340,13 @@ class TestGetLitellmParamsDataResidency:
)
assert result["data_residency"] == "us"
def test_fedramp_host_resolves_to_fedramp(self):
result = get_litellm_params(
custom_llm_provider="openai",
api_base="https://gov.api.openai.com/v1",
)
assert result["data_residency"] == "fedramp"
def test_global_host_resolves_to_none(self):
result = get_litellm_params(
custom_llm_provider="openai",

View file

@ -54,9 +54,25 @@ SIGNATURE: Final = "opaque-thought-signature"
def _provider_fields() -> dict[str, object]:
return {
"reasoning_content": SECRET,
"reasoning": SECRET,
"reasoning_text": SECRET,
"reasoning_details": [{"type": "reasoning.text", "text": SECRET, "index": 0, "format": "unknown"}],
"thinking_blocks": [{"type": "thinking", "thinking": SECRET}],
"reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}],
"reasoningContent": {"reasoningText": {"text": SECRET}},
"annotations": [{"type": "file", "file": {"name": "notes.pdf", "content": [{"type": "text", "text": SECRET}]}}],
"citations": [{"cited_text": SECRET}],
"citationsContent": [{"citations": [{"location": {}}], "content": [{"text": SECRET}]}],
"citation": {"type": "char_location", "cited_text": SECRET},
"web_search_calls": [{"type": "web_search_call", "action": {"type": "search", "query": SECRET}}],
"web_search_results": [{"type": "web_search_tool_result", "content": [{"title": SECRET}]}],
"tool_results": [{"type": "web_fetch_tool_result", "content": {"text": SECRET}}],
"code_interpreter_results": [{"stdout": SECRET}],
"search_results": [{"object": "vector_store.search_results.page", "data": [{"content": [{"text": SECRET}]}]}],
"rerank_results": {"results": [{"document": {"text": SECRET}}]},
"executed_tools": [{"type": "search", "search_results": {"results": [{"content": SECRET}]}}],
"custom_outputs": {"answer": SECRET},
"tool_search_items": [{"type": "tool_search_output", "tools": [{"description": SECRET}]}],
"compaction_blocks": [{"type": "compaction", "content": SECRET, "signature": "compaction-signature"}],
"compaction_start": {"type": "compaction", "content": SECRET},
"compaction_delta": {"type": "compaction_delta", "content": SECRET},

View file

View file

@ -0,0 +1,40 @@
from collections.abc import Mapping
import pytest
from litellm.llms.empiriolabs.decisions.transformation import EmpirioLabsDecisionsConfig
@pytest.mark.parametrize(
("environment", "connection"),
(
({}, ("https://api.empiriolabs.ai", None)),
(
{"EMPIRIOLABS_API_BASE": "https://gateway.example/v1", "EMPIRIOLABS_API_KEY": "environment-key"},
("https://gateway.example/v1", "environment-key"),
),
),
ids=("public_api_without_a_key", "environment"),
)
def test_connection_comes_from_the_environment_before_the_public_api(
environment: Mapping[str, str],
connection: tuple[str, str | None],
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False)
for name, value in environment.items():
monkeypatch.setenv(name, value)
assert EmpirioLabsDecisionsConfig().resolve_connection(None, None) == connection
@pytest.mark.parametrize(
"api_base",
("https://api.empiriolabs.ai", "https://api.empiriolabs.ai/v1", "https://api.empiriolabs.ai/v1/"),
)
def test_requests_go_to_v1_decisions_whether_or_not_the_base_ends_in_v1(api_base: str) -> None:
assert (
EmpirioLabsDecisionsConfig().get_complete_url(api_base=api_base, model="aplomb-1")
== "https://api.empiriolabs.ai/v1/decisions"
)

View file

@ -112,6 +112,28 @@ def test_responses_us_api_base_sets_data_residency():
assert logging_obj.litellm_params.get("data_residency") == "us"
def test_responses_fedramp_api_base_sets_data_residency():
mock_client = _make_mock_http_client(_make_responses_api_response_body())
captured, init_spy = _capture_logging_obj()
with (
patch(
"litellm.llms.custom_httpx.llm_http_handler.get_httpx_client",
return_value=mock_client,
),
patch.object(litellm.Logging, "__init__", init_spy),
):
litellm.responses(
model="gpt-4.1",
input="hi",
api_base="https://gov.api.openai.com/v1",
api_key="test-key",
)
logging_obj = captured["logging_obj"]
assert logging_obj.litellm_params.get("data_residency") == "fedramp"
def test_responses_global_api_base_leaves_data_residency_none():
mock_client = _make_mock_http_client(_make_responses_api_response_body())
captured, init_spy = _capture_logging_obj()

View file

@ -2,7 +2,9 @@
import pytest
import litellm
from litellm.llms.openai.data_residency import infer_openai_data_residency
from litellm.types.utils import DataResidency
@pytest.mark.parametrize(
@ -13,6 +15,8 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency
("https://us.api.openai.com/v1", "us"),
("https://us.api.openai.com", "us"),
("https://EU.api.openai.com/v1", "eu"),
("https://gov.api.openai.com/v1", DataResidency.FEDRAMP.value),
("https://GOV.api.openai.com/v1", DataResidency.FEDRAMP.value),
("https://api.openai.com/v1", None),
("https://api.openai.com", None),
("https://example.com/v1", None),
@ -28,7 +32,26 @@ def test_infer_openai_data_residency(api_base, expected):
@pytest.mark.parametrize("custom_llm_provider", [None, "anthropic", "azure", "bedrock"])
def test_infer_openai_data_residency_non_openai_provider(custom_llm_provider):
assert (
infer_openai_data_residency(custom_llm_provider, "https://eu.api.openai.com/v1")
is None
)
assert infer_openai_data_residency(custom_llm_provider, "https://eu.api.openai.com/v1") is None
@pytest.mark.parametrize("env_var", ["OPENAI_BASE_URL", "OPENAI_API_BASE"])
def test_infer_openai_data_residency_from_env_when_api_base_unset(monkeypatch, env_var):
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
monkeypatch.setenv(env_var, "https://gov.api.openai.com/v1")
assert infer_openai_data_residency("openai", None) == DataResidency.FEDRAMP.value
def test_infer_openai_data_residency_explicit_api_base_wins_over_env(monkeypatch):
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setenv("OPENAI_BASE_URL", "https://gov.api.openai.com/v1")
assert infer_openai_data_residency("openai", "https://eu.api.openai.com/v1") == "eu"
def test_infer_openai_data_residency_from_litellm_api_base(monkeypatch):
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
monkeypatch.setattr(litellm, "api_base", "https://us.api.openai.com/v1")
assert infer_openai_data_residency("openai", None) == "us"

View file

@ -2,13 +2,19 @@
Unit tests for the EmpirioLabs OpenAI-like provider.
"""
import json
import os
import sys
from typing import Final
import pytest
import respx
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
)
import litellm
from litellm.llms.openai_like.dynamic_config import create_config_class
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
@ -61,3 +67,50 @@ def test_empiriolabs_complete_url_appends_endpoint():
stream=False,
)
assert url == f"{EMPIRIOLABS_BASE_URL}/chat/completions"
def test_empiriolabs_chat_models_still_route_through_the_json_provider(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key")
monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False)
with respx.mock() as upstream:
route: Final = upstream.post(f"{EMPIRIOLABS_BASE_URL}/chat/completions").respond(
200,
json={
"id": "chatcmpl_empiriolabs",
"object": "chat.completion",
"created": 1_789_550_000,
"model": "qwen3-7-plus",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello from EmpirioLabs AI"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 4, "completion_tokens": 5, "total_tokens": 9},
},
)
response: Final = litellm.completion(
model="empiriolabs/qwen3-7-plus",
messages=[{"role": "user", "content": "Say hello"}],
max_completion_tokens=16,
)
request: Final = route.calls.last.request
assert litellm.get_llm_provider("empiriolabs/qwen3-7-plus") == (
"qwen3-7-plus",
"empiriolabs",
"environment-key",
EMPIRIOLABS_BASE_URL,
)
assert route.call_count == 1
assert request.headers["authorization"] == "Bearer environment-key"
assert json.loads(request.content) == {
"model": "qwen3-7-plus",
"messages": [{"role": "user", "content": "Say hello"}],
"max_tokens": 16,
}
assert response.choices[0].message.content == "Hello from EmpirioLabs AI"

View file

@ -4023,11 +4023,11 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
async def fake_find_many(**kwargs):
if "where" not in kwargs:
return [_tag_registry_row("usecase:spent")]
return [_tag_registry_row("app:spent")]
return [_tag_db_row(name, max_budget=1.0) for name in kwargs["where"]["tag_name"]["in"]]
async def spent_tag_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
return 1.5 if counter_key == "spend:tag:usecase:spent" else fallback_spend
return 1.5 if counter_key == "spend:tag:app:spent" else fallback_spend
async def receive():
return {
@ -4047,7 +4047,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=spent")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=spent")],
},
receive,
)
@ -4055,7 +4055,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
request_body,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -4077,7 +4077,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
valid_token=UserAPIKeyAuth(token="test-token", metadata=key_metadata),
request=request,
)
assert exc_info.value.entity_id == "usecase:spent"
assert exc_info.value.entity_id == "app:spent"
@pytest.mark.asyncio
@ -11789,8 +11789,8 @@ def _client_tag_request(client_tags, baggage):
{
"type": "http",
"headers": [
*([(b"baggage", b"entrypoint=a")] if baggage else []),
*([(b"x-litellm-tags", b"usecase:a")] if client_tags == "header" else []),
*([(b"baggage", b"service=a")] if baggage else []),
*([(b"x-litellm-tags", b"app:a")] if client_tags == "header" else []),
],
}
)
@ -11814,7 +11814,7 @@ def _client_tag_rejection(route, data, request):
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/messages", "/v1/responses"])
@pytest.mark.parametrize("client_tags", [None, [], ["usecase:a"], '{"tags":["usecase:a"]}', "header"])
@pytest.mark.parametrize("client_tags", [None, [], ["app:a"], '{"tags":["app:a"]}', "header"])
@pytest.mark.parametrize("client_slot", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize("baggage", [True, False])
async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_it(
@ -11833,7 +11833,7 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
mapped_request,
mapped,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -11843,19 +11843,19 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_
route, unmapped, unmapped_request
)
if client_tags is None and baggage:
assert "usecase:a" in mapped["metadata" if "chat/completions" in route else "litellm_metadata"]["tags"]
assert "app:a" in mapped["metadata" if "chat/completions" in route else "litellm_metadata"]["tags"]
@pytest.mark.asyncio
async def test_baggage_mapping_still_rejects_client_metadata_tags():
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
request = _client_tag_request(["usecase:b"], baggage=True)
data = {"metadata": {"tags": ["usecase:b"]}}
request = _client_tag_request(["app:b"], baggage=True)
data = {"metadata": {"tags": ["app:b"]}}
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)

View file

@ -2203,7 +2203,7 @@ async def test_websocket_inference_is_denied_when_a_required_baggage_mapping_is_
{
"request_tag_mappings": [
{"key": "region", "tag_prefix": "region:"},
{"key": "entrypoint", "tag_prefix": "usecase:", "required": True},
{"key": "service", "tag_prefix": "app:", "required": True},
]
},
),
@ -2227,7 +2227,7 @@ async def test_websocket_invalid_key_is_rejected_before_the_required_baggage_map
with (
patch.dict(
litellm.proxy.proxy_server.general_settings,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "required": True}]},
),
patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
@ -2248,7 +2248,7 @@ async def test_websocket_inference_is_allowed_when_no_baggage_mapping_is_require
with (
patch.dict(
litellm.proxy.proxy_server.general_settings,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]},
),
patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", new=AsyncMock(return_value=authorized)
@ -2273,8 +2273,8 @@ async def _auth_then_serve_chat(monkeypatch, general_settings, body, headers):
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"], "mock_response": "hi"},
"model_info": {"id": "usecase-deployment"},
"litellm_params": {"model": "gpt-4o", "tags": ["app:a"], "mock_response": "hi"},
"model_info": {"id": "app-deployment"},
},
{
"model_name": "gpt-4",
@ -2321,30 +2321,30 @@ async def test_auth_header_tag_merge_keeps_json_string_metadata_tags_routable(mo
body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"metadata": json.dumps({"tags": ["usecase:a"]}),
"metadata": json.dumps({"tags": ["app:a"]}),
}
cached_body, served_by = await _auth_then_serve_chat(
monkeypatch, general_settings, body, [(b"x-litellm-tags", b"other"), (b"baggage", b"region=sg")]
)
assert served_by == "usecase-deployment"
assert cached_body["metadata"] == {"tags": ["usecase:a", "other"]}
assert served_by == "app-deployment"
assert cached_body["metadata"] == {"tags": ["app:a", "other"]}
@pytest.mark.asyncio
@pytest.mark.parametrize("client_metadata", ["run-123", 42, ["x"]])
@pytest.mark.parametrize(
"baggage",
[[], [(b"baggage", b"region=sg")], [(b"baggage", b"entrypoint=a")]],
[[], [(b"baggage", b"region=sg")], [(b"baggage", b"service=a")]],
ids=["no_baggage", "unmapped_baggage", "mapped_baggage"],
)
async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_baggage_maps(
monkeypatch, client_metadata, baggage
):
body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "metadata": client_metadata}
headers = [(b"x-litellm-tags", b"usecase:a"), *baggage]
headers = [(b"x-litellm-tags", b"app:a"), *baggage]
unset = await _auth_then_serve_chat(monkeypatch, {}, body, headers)
mapped = await _auth_then_serve_chat(
monkeypatch, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, body, headers
monkeypatch, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, body, headers
)
assert mapped == unset
assert mapped[1] == "usecase-deployment"
assert mapped[1] == "app-deployment"

View file

@ -26,7 +26,7 @@ def test_accounting_enrichment_without_request_state_preserves_legacy_metadata()
@pytest.mark.asyncio
@pytest.mark.parametrize("header", [b"", b"entrypoint=a,entrypoint=b"])
@pytest.mark.parametrize("header", [b"", b"service=a,service=b"])
@pytest.mark.parametrize("client_metadata", [None, '{"requester_metadata":{"guardrails":["qa-guard"]}}'])
async def test_mapping_rejection_keeps_string_metadata_safe_for_auth_failure_logging(header, client_metadata):
from litellm.utils import get_request_guardrails
@ -37,7 +37,7 @@ async def test_mapping_rejection_keeps_string_metadata_safe_for_auth_failure_log
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "required": True}]},
None,
MagicMock(),
)
@ -59,7 +59,7 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
return {"type": "http.request", "body": json.dumps(body).encode()}
request = Request(
{"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"baggage", b"entrypoint=a")]},
{"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"baggage", b"service=a")]},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
@ -67,11 +67,11 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
assert "usecase:a" in data["litellm_metadata"]["tags"]
assert "app:a" in data["litellm_metadata"]["tags"]
assert "tags" not in data[slot] or slot == "litellm_metadata"
assert get_request_guardrails(data) == []
assert await read_request_body(request) == body
@ -80,60 +80,60 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
@pytest.mark.parametrize(
"header,value",
[
("entrypoint=a", "a"),
("entrypoint = a", "a"),
("region=sg, entrypoint=customer%20support;version=1", "customer support"),
("service=a", "a"),
("service = a", "a"),
("region=sg, service=customer%20support;version=1", "customer support"),
],
)
def test_variable_mapper_uses_decoded_baggage_value(header, value):
rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:")
assert resolve_baggage_tag_mappings(header, [rule]) == ((f"usecase:{value}", False),)
rule = RequestTagMapping(key="service", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == ((f"app:{value}", False),)
@pytest.mark.parametrize(
"header,expected",
[("entrypoint=a", "department:support"), ("entrypoint=b", "department:sales"), ("entrypoint=x", "usecase:x")],
[("service=a", "department:support"), ("service=b", "department:sales"), ("service=x", "app:x")],
)
def test_exact_rules_override_variable_fallback_independent_of_order(header, expected):
rules = [
RequestTagMapping(key="entrypoint", tag_prefix="usecase:"),
RequestTagMapping(key="entrypoint", value="a", tag="department:support"),
RequestTagMapping(key="entrypoint", value="b", tag="department:sales"),
RequestTagMapping(key="service", tag_prefix="app:"),
RequestTagMapping(key="service", value="a", tag="department:support"),
RequestTagMapping(key="service", value="b", tag="department:sales"),
]
assert resolve_baggage_tag_mappings(header, rules) == ((expected, False),)
assert resolve_baggage_tag_mappings(header, list(reversed(rules))) == ((expected, False),)
@pytest.mark.parametrize("header", ["entrypoint=%00", "entrypoint=a%0Ab", "entrypoint=b,entry%70oint=a"])
@pytest.mark.parametrize("header", ["service=%00", "service=a%0Ab", "service=b,serv%69ce=a"])
def test_variable_mapper_skips_control_characters_and_percent_encoded_duplicate_keys(header):
rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:")
rule = RequestTagMapping(key="service", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == ()
@pytest.mark.parametrize(
"header,expected",
[
("entry+point=a", (("usecase:a", False),)),
("entry%2Bpoint=a", (("usecase:a", False),)),
("entry+point=a,entry%2Bpoint=b", ()),
("serv+ice=a", (("app:a", False),)),
("serv%2Bice=a", (("app:a", False),)),
("serv+ice=a,serv%2Bice=b", ()),
],
)
def test_key_with_a_plus_matches_both_wire_spellings_and_rejects_both_at_once(header, expected):
rule = RequestTagMapping(key="entry+point", tag_prefix="usecase:")
rule = RequestTagMapping(key="serv+ice", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == expected
def test_exact_rules_work_without_a_variable_rule():
rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:manual")
assert resolve_baggage_tag_mappings("entrypoint=a", [rule]) == (("usecase:manual", False),)
assert resolve_baggage_tag_mappings("entrypoint=x", [rule]) == ()
rule = RequestTagMapping(key="service", value="a", tag="app:manual")
assert resolve_baggage_tag_mappings("service=a", [rule]) == (("app:manual", False),)
assert resolve_baggage_tag_mappings("service=x", [rule]) == ()
@pytest.mark.parametrize(
"header", ["", "region=sg", "entrypoint=", "entrypoint=a,entrypoint=b", "entrypoint=b,entry%70oint=a", "entrypoint=x"]
"header", ["", "region=sg", "service=", "service=a,service=b", "service=b,serv%69ce=a", "service=x"]
)
def test_required_mapping_rejects_missing_ambiguous_or_unmapped_values(header):
rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:a", required=True)
rule = RequestTagMapping(key="service", value="a", tag="app:a", required=True)
with pytest.raises(ProxyException) as error:
resolve_baggage_tag_mappings(header, [rule])
assert error.value.code == "400"
@ -142,10 +142,10 @@ def test_required_mapping_rejects_missing_ambiguous_or_unmapped_values(header):
@pytest.mark.parametrize(
"rule",
[
{"key": "entrypoint", "tag": "usecase:a"},
{"key": "entrypoint", "value": "a", "tag_prefix": "usecase:"},
{"key": "entrypoint", "value": "a", "tag": "usecase:a", "tag_prefix": "usecase:"},
{"key": "", "tag_prefix": "usecase:"},
{"key": "service", "tag": "app:a"},
{"key": "service", "value": "a", "tag_prefix": "app:"},
{"key": "service", "value": "a", "tag": "app:a", "tag_prefix": "app:"},
{"key": "", "tag_prefix": "app:"},
],
)
def test_invalid_rule_is_rejected_during_config_validation(rule):
@ -158,8 +158,8 @@ def test_duplicate_fallbacks_are_rejected():
ConfigGeneralSettings.model_validate(
{
"request_tag_mappings": [
{"key": "entrypoint", "tag_prefix": "one:"},
{"key": "entrypoint", "tag_prefix": "two:"},
{"key": "service", "tag_prefix": "one:"},
{"key": "service", "tag_prefix": "two:"},
]
}
)
@ -171,23 +171,23 @@ def test_duplicate_fallbacks_are_rejected():
[("/chat/completions", "metadata"), ("/v1/messages", "litellm_metadata"), ("/v1/responses", "litellm_metadata")],
)
async def test_mapped_tag_merges_with_existing_tags_in_the_budget_metadata_slot(route, slot):
data = {slot: {"tags": ["existing", "usecase:a"]}}
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
data = {slot: {"tags": ["existing", "app:a"]}}
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(data, route)
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
assert data[slot]["tags"] == ["existing", "usecase:a"]
assert data[slot]["tags"] == ["existing", "app:a"]
@pytest.mark.asyncio
async def test_budgeted_mapper_requires_a_database():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException) as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, {}, config, None, MagicMock()
@ -197,10 +197,10 @@ async def test_budgeted_mapper_requires_a_database():
@pytest.mark.asyncio
async def test_budgeted_mapper_answers_503_when_the_tag_budget_lookup_fails():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable"))
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="could not read tag budgets") as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, {"metadata": {}}, config, prisma_client, UserApiKeyCache()
@ -210,16 +210,16 @@ async def test_budgeted_mapper_answers_503_when_the_tag_budget_lookup_fails():
@pytest.mark.asyncio
async def test_budgeted_mapper_serves_through_a_db_outage_when_requests_are_allowed_on_db_unavailable():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable"))
data = {"metadata": {}}
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with patch.dict("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": True}):
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, prisma_client, UserApiKeyCache()
)
assert data["metadata"]["tags"] == ["usecase:a"]
assert data["metadata"]["tags"] == ["app:a"]
async def _cache_with_tag(tag: LiteLLM_TagTable) -> UserApiKeyCache:
@ -234,11 +234,11 @@ async def _cache_with_tag(tag: LiteLLM_TagTable) -> UserApiKeyCache:
[None, LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=None, tpm_limit=10)],
)
async def test_budgeted_mapper_rejects_a_tag_without_a_max_budget(budget_table):
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
data = {"metadata": {}}
cache = await _cache_with_tag(LiteLLM_TagTable(tag_name="usecase:a", litellm_budget_table=budget_table))
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="usecase:a requires an existing tag budget") as error:
cache = await _cache_with_tag(LiteLLM_TagTable(tag_name="app:a", litellm_budget_table=budget_table))
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="app:a requires an existing tag budget") as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, MagicMock(), cache
)
@ -248,23 +248,23 @@ async def test_budgeted_mapper_rejects_a_tag_without_a_max_budget(budget_table):
@pytest.mark.asyncio
async def test_budgeted_mapper_admits_a_tag_with_a_max_budget_and_unbudgeted_rules_need_none():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a,region=sg")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a,region=sg")]})
data = {"metadata": {}}
cache = await _cache_with_tag(
LiteLLM_TagTable(
tag_name="usecase:a", litellm_budget_table=LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=1.0)
tag_name="app:a", litellm_budget_table=LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=1.0)
)
)
config = {
"request_tag_mappings": [
{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True},
{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True},
{"key": "region", "tag_prefix": "region:"},
]
}
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, MagicMock(), cache
)
assert data["metadata"]["tags"] == ["usecase:a", "region:sg"]
assert data["metadata"]["tags"] == ["app:a", "region:sg"]
@pytest.mark.parametrize(
@ -325,7 +325,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
"scheme": "http",
"server": ("localhost", 4000),
"query_string": query,
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a")],
},
receive,
)
@ -337,7 +337,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -350,7 +350,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
proxy_config=ProxyConfig(),
general_settings={},
)
assert enriched[slot]["tags"] == ["usecase:a"]
assert enriched[slot]["tags"] == ["app:a"]
assert "key" not in enriched
assert "api-version" not in enriched
assert "key" not in enriched["proxy_server_request"]["body"]
@ -360,7 +360,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
@pytest.mark.asyncio
@pytest.mark.parametrize(
"client_metadata",
[{"trace_name": "qa"}, {"trace_name": "qa", "tags": ["client:x"]}, {"trace_name": "qa", "tags": ["usecase:a"]}],
[{"trace_name": "qa"}, {"trace_name": "qa", "tags": ["client:x"]}, {"trace_name": "qa", "tags": ["app:a"]}],
)
async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
from litellm.proxy._types import UserAPIKeyAuth
@ -380,14 +380,14 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a")],
},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
assert parse_error is None
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock()
request, data, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, None, MagicMock()
)
enriched = await add_litellm_data_to_request(
data=await read_request_body(request),
@ -397,7 +397,7 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
general_settings={},
)
requester_metadata = enriched["metadata"]["requester_metadata"]
assert "usecase:a" in enriched["metadata"]["tags"]
assert "app:a" in enriched["metadata"]["tags"]
assert requester_metadata["trace_name"] == "qa"
assert requester_metadata.get("tags") == client_metadata.get("tags")
assert BAGGAGE_MAPPED_TAGS_METADATA_KEY not in requester_metadata
@ -407,9 +407,9 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
@pytest.mark.parametrize(
"headers,body_fields",
[
([(b"x-litellm-tags", b"usecase:a")], {}),
([], {"tags": ["usecase:a"]}),
([], {"litellm_metadata": {"tags": ["usecase:a"]}}),
([(b"x-litellm-tags", b"app:a")], {}),
([], {"tags": ["app:a"]}),
([], {"litellm_metadata": {"tags": ["app:a"]}}),
],
ids=["header", "body", "litellm_metadata"],
)
@ -432,14 +432,14 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a"), *headers],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a"), *headers],
},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
assert parse_error is None
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock()
request, data, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, None, MagicMock()
)
enriched = await add_litellm_data_to_request(
data=await read_request_body(request),
@ -452,8 +452,8 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"]},
"model_info": {"id": "usecase-deployment"},
"litellm_params": {"model": "gpt-4o", "tags": ["app:a"]},
"model_info": {"id": "app-deployment"},
},
{
"model_name": "gpt-4",
@ -470,4 +470,4 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
metadata=dict(enriched["metadata"]),
mock_response="hi",
)
assert response._hidden_params["model_id"] == "usecase-deployment"
assert response._hidden_params["model_id"] == "app-deployment"

View file

@ -1,149 +0,0 @@
import json
from datetime import datetime, timezone
from decimal import Decimal
from typing import Final, Literal, NoReturn
import pytest
from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite
from litellm.types.utils import CallTypes
def _no_transaction() -> NoReturn:
raise AssertionError("Invalid accounting input must not open a database transaction")
@pytest.fixture
def event() -> InboundUsageEvent:
return InboundUsageEvent(
event_id="caller",
request_id="caller",
call_type=CallTypes.completion,
started_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
outcome=Outcome.SUCCEEDED,
dimensions=(),
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"])
async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None:
with pytest.raises(ValueError, match=r"Accounting envelope|validation error"):
await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope)
@pytest.mark.asyncio
@pytest.mark.parametrize("version", [0, -1, True, 1.5])
@pytest.mark.parametrize("field", ["version", "schema_version"])
async def test_invalid_versions_fail_before_database_access(
event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"]
) -> None:
with pytest.raises(ValueError, match="positive integers"):
await AccountingLedger(_no_transaction).insert_event(
event,
(),
envelope="{}",
**{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions
)
@pytest.mark.asyncio
async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None:
invalid: Final = event.model_copy(update={"event_id": ""})
with pytest.raises(ValueError, match="event_id"):
await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}")
@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5])
async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None:
with pytest.raises(ValueError, match="batch size"):
await AccountingLedger(_no_transaction).purge_completed(
limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes
)
@pytest.mark.asyncio
@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"])
async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None:
cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
deleted: Final = await AccountingLedger(_no_transaction).purge_completed(
redelivery_before=cutoff if missing == "recovery" else None,
recovery_before=cutoff if missing == "redelivery" else None,
)
assert deleted == 0
@pytest.mark.asyncio
async def test_retention_rejects_ambiguous_naive_watermarks() -> None:
with pytest.raises(ValueError, match="timezone"):
await AccountingLedger(_no_transaction).purge_completed(
redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc)
)
@pytest.mark.asyncio
async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}})
with pytest.raises(ValueError, match="Out of range float"):
await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope)
@pytest.mark.asyncio
async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
envelope: Final = json.dumps({"event": event.model_dump(mode="json")})
with pytest.raises(ValueError, match="not a valid DeliveryTarget"):
await AccountingLedger(_no_transaction).insert_event(
event,
("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target
envelope=envelope,
)
@pytest.mark.asyncio
async def test_empty_usage_batches_do_not_open_a_transaction() -> None:
ledger: Final = AccountingLedger(_no_transaction)
await ledger.insert_usage_batch(())
assert await ledger.load_usage_batch(()) == ()
@pytest.mark.asyncio
async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")}))
changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}}))
with pytest.raises(AccountingConflictError, match="conflicting copies"):
await AccountingLedger(_no_transaction).insert_usage_batch((original, changed))
@pytest.mark.asyncio
async def test_loaded_usage_preserves_decimal_metrics_from_jsonb(event: InboundUsageEvent) -> None:
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import cast
from unittest.mock import AsyncMock
from prisma import Prisma
from litellm.accounting import OutboundUsageEvent
outbound: Final = OutboundUsageEvent(
**event.model_dump(exclude={"metrics"}),
attempt_id="attempt",
retry_of_attempt_id=None,
metrics=(Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),),
)
payload: Final = {"event": outbound.model_dump(mode="json"), "usage": {"prompt_tokens": 120000}}
database: Final = AsyncMock(spec=Prisma)
database.query_raw.return_value = [{"payload": payload, "targets": []}]
@asynccontextmanager
async def transaction() -> AsyncIterator[Prisma]:
yield cast(Prisma, database)
recovered: Final = await AccountingLedger(transaction).load_usage_batch((outbound.event_id,))
assert len(recovered) == 1
assert recovered[0].event == outbound
assert json.loads(recovered[0].envelope) == payload

View file

@ -3981,140 +3981,3 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs(
user_api_key_dict=user_api_key_dict,
)
assert len(writer.rows) == rows_written
@pytest.mark.asyncio
async def test_failure_capture_overload_does_not_skip_budget_release(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _write
database: Final = CommitBarrier()
runtime: Final = UsageCaptureRuntime()
runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2))
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "prisma_client", None)
meter: Final = runtime.meter
assert meter is not None
receipt: Final = meter.submit((_write("occupied-one"), _write("occupied-two")))
auth: Final = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False})
request: Final = {
"call_type": "acompletion",
"litellm_call_id": "failed",
"model": "capture-model",
"start_time": _CALL_START,
"litellm_params": {},
}
try:
await database.entered.wait()
await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), auth)
assert auth.budget_reservation is not None and auth.budget_reservation["finalized"] is True
assert "Usage capture failed on the request failure path" in caplog.text
assert "_accounting_capture" not in request
finally:
database.release.set()
await runtime.stop(1)
await receipt
assert database.committed == ("occupied-one", "occupied-two")
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [False, True])
@pytest.mark.parametrize("complete", [False, True])
@pytest.mark.parametrize("deferred", [False, True])
async def test_capture_and_legacy_share_resolved_key_attribution(
failure: bool, complete: bool, deferred: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
from pydantic import TypeAdapter
from litellm.accounting import DeliveryTarget, DimensionName
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage
key: Final = UserAPIKeyAuth(
api_key="a" * 64,
user_id="owner",
team_id="team",
org_id="org",
project_id="project",
key_alias="key",
team_alias="team-alias",
request_route="/v1/chat/completions",
)
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key="a" * 64, value=key)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", object())
writer: Final = _RecordingSpendWriter()
monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(db_spend_update_writer=writer))
database: Final = StoredUsage()
runtime: Final = UsageCaptureRuntime()
runtime.start(database.insert, UsageCaptureSettings())
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
rejected: Final = failure and deferred and not complete
owner: Final = "original-owner" if deferred and not rejected else "owner"
metadata: Final = {
"user_api_key": "b" * 64 if rejected else "a" * 64,
**(
{
"user_api_key_user_id": "forged-owner" if rejected else owner,
"user_api_key_team_id": "forged-team" if rejected else "team",
"user_api_key_org_id": "forged-org" if rejected else "org",
"user_api_key_project_id": None if deferred and not complete else "project",
"user_api_key_alias": "key",
"user_api_key_team_alias": "team-alias",
}
if complete or deferred
else {}
),
}
kwargs: Final = {
"call_type": "aretrieve_batch" if deferred else "call_mcp_tool",
"litellm_call_id": "attributed",
"model": "MCP: echo",
"response_cost": 0.25,
"start_time": _CALL_START,
"litellm_params": {"metadata": metadata},
}
if rejected:
from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY
kwargs[REJECTED_BEFORE_REQUEST_SETUP_KEY] = REJECTED_BEFORE_REQUEST_SETUP
logger: Final = ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer))
try:
with track_batch_accounting("attributed" if deferred and not rejected else "other"):
if failure:
await logger.async_post_call_failure_hook(
kwargs,
HTTPException(401, "model not allowed"),
key if complete else UserAPIKeyAuth(api_key="a" * 64, request_route="/v1/chat/completions"),
)
else:
await logger.async_log_success_event(kwargs, {"id": "tool"}, _CALL_START, _CALL_START)
finally:
await runtime.stop(1)
assert len(writer.rows) == 1
assert len(database.rows) == 2
dimensions: Final = {item.name: item.value for item in database.rows[-1].event.dimensions}
if rejected:
assert dimensions.get(DimensionName.KEY) == key.api_key
assert {
name: dimensions.get(name)
for name in (DimensionName.USER, DimensionName.TEAM, DimensionName.ORGANIZATION, DimensionName.PROJECT)
} == {
DimensionName.USER: owner,
DimensionName.TEAM: "team",
DimensionName.ORGANIZATION: "org",
DimensionName.PROJECT: None if deferred and not complete and not rejected else "project",
}
assert set(database.rows[-1].targets) >= {DeliveryTarget.DAILY_TEAM_SPEND, DeliveryTarget.DAILY_ORGANIZATION_SPEND}
legacy: Final = TypeAdapter(dict[str, object]).validate_python(writer.rows[0]["kwargs"])
payload: Final = get_logging_payload(legacy, {"id": "tool"}, _CALL_START, _CALL_START)
assert (payload["user"], payload["team_id"], payload["organization_id"]) == (owner, "team", "org")

View file

@ -529,6 +529,7 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo:
("Nace", "Nace", LlmProviders.NACE),
("PARALLEL_AI", "Parallel AI", LlmProviders.PARALLEL_AI),
("EU_JEV", "eu/jev", LlmProviders.EU_JEV),
("EMPIRIOLABS", "EmpirioLabs AI", LlmProviders.EMPIRIOLABS),
),
)
def test_decisions_provider_fields(provider_key: str, display_name: str, provider: LlmProviders) -> None:

View file

@ -92,12 +92,6 @@ def response(cost: float = 0.5) -> ModelResponse:
async def test_header_and_logging_share_one_allowance_allocation(
logging_obj: Logging, used: int, expected: float
) -> None:
import json
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage
from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings
counter: Final = Counter(used)
result: Final = response()
costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3)))
@ -116,24 +110,6 @@ async def test_header_and_logging_share_one_allowance_allocation(
assert payload["cost_finalized"] is True
assert payload["cost_breakdown"] == costs[0].logging_breakdown()
committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
committed.put_nowait(write)
meter: Final = DurableUsageMeter(persist, UsageCaptureSettings())
assert logging_obj.prepared_cost is not None
prepared_usage: Final = logging_obj.prepared_cost.usage.model_dump()
result.usage = Usage(prompt_tokens=999)
capture_logging_usage(meter, logging_obj, result, NOW, NOW)
await meter.close()
committed.get_nowait()
event: Final = json.loads(committed.get_nowait().envelope)
assert event["finalized_cost"]["response_cost"] == pytest.approx(costs[0].amount)
assert event["allocation"]["discount"] == pytest.approx(costs[0].allowance_discount)
assert event["usage"] == prepared_usage
assert counter.calls == 1
@pytest.mark.asyncio

View file

@ -7695,13 +7695,11 @@ def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadat
assert actual == expected
@pytest.mark.parametrize("capture_enabled", [False, True])
@pytest.mark.parametrize("completion_start", [None, "absent", "reported"])
def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, completion_start: str | None) -> None:
def test_spend_payload_preserves_unknown_completion_start(completion_start: str | None) -> None:
start: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc)
end: Final = start + datetime.timedelta(seconds=2)
reported: Final = start + datetime.timedelta(seconds=1)
observations: Final[list[object]] = []
kwargs: Final = {
"call_type": "acompletion",
"litellm_call_id": "timestamp-call",
@ -7713,9 +7711,7 @@ def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool,
{"id": "timestamp-response"},
start,
end,
capture_usage=(lambda row, usage: observations.append(row)) if capture_enabled else None,
)
assert payload["completionStartTime"] == (
end if completion_start == "absent" else reported if completion_start else None
)
assert observations == ([payload] if capture_enabled else [])

File diff suppressed because it is too large Load diff

View file

@ -1,473 +0,0 @@
import asyncio
import json
from datetime import datetime, timezone
from decimal import Decimal
from typing import Final
import pytest
from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome
from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite
from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings
from litellm.types.utils import CallTypes
def _write(identity: str) -> UsageWrite:
event: Final = InboundUsageEvent(
event_id=identity,
request_id=identity,
call_type=CallTypes.acompletion,
started_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
outcome=Outcome.SUCCEEDED,
dimensions=(),
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
)
return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")}))
class CommitBarrier:
def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None:
self.entered: Final = asyncio.Event()
self.release: Final = asyncio.Event()
self.committed: tuple[str, ...] = ()
self.transactions = 0
self.conflict: Final = conflict
self.failure: Final = failure
async def write(self, writes: tuple[UsageWrite, ...]) -> None:
self.transactions += 1
self.entered.set()
await self.release.wait()
if self.failure is not None:
raise self.failure
identities: Final = tuple(write.event.event_id for write in writes)
if self.conflict in identities:
raise AccountingConflictError("Conflicting test event")
self.committed += identities
async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None:
await wakeup.wait()
@pytest.mark.asyncio
async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2))
first: Final = meter.submit((_write("one"),))
second: Final = meter.submit((_write("two"),))
await database.entered.wait()
assert not first.done() and not second.done()
assert meter.pending_records == 2
with pytest.raises(UsageCaptureError, match="capacity"):
meter.submit((_write("three"),))
database.release.set()
await asyncio.gather(first, second)
await meter.close()
assert database.committed == ("one", "two")
assert database.transactions == 1
assert (meter.pending_bytes, meter.pending_records) == (0, 0)
@pytest.mark.asyncio
async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None:
database: Final = CommitBarrier()
write: Final = _write("large" * 300)
size: Final = len(write.envelope.encode())
meter: Final = DurableUsageMeter(
database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush
)
reservation: Final = meter.reserve_producer(size // 2, size // 2)
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.reserve_producer(1, size)
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.resize_producer(reservation, size, 1)
assert meter.pending_bytes == size // 2
meter.release_producer(reservation)
receipt: Final = meter.submit((write,))
with pytest.raises(UsageCaptureError, match="capacity"):
meter.submit((_write("extra"),))
assert meter.pending_bytes == size
database.release.set()
await meter.close()
await receipt
assert database.committed == (write.event.event_id,)
with pytest.raises(UsageCaptureError, match="closed"):
meter.submit((_write("after-close"),))
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.reserve_producer(1, 1)
@pytest.mark.asyncio
async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None:
database: Final = CommitBarrier(conflict="conflict")
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3))
rejected: Final = meter.submit((_write("caller"), _write("conflict")))
accepted: Final = meter.submit((_write("valid"),))
await database.entered.wait()
database.release.set()
with pytest.raises(UsageCaptureError, match="identity"):
await rejected
await accepted
await meter.close()
assert database.committed == ("valid",)
assert meter.pending_records == 0
@pytest.mark.asyncio
async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None:
failure: Final = ValueError("invalid envelope")
database: Final = CommitBarrier(failure=failure)
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
first: Final = meter.submit((_write("one"),))
second: Final = meter.submit((_write("two"),))
await database.entered.wait()
database.release.set()
results: Final = await asyncio.gather(first, second, return_exceptions=True)
await meter.close()
assert results == [failure, failure]
assert meter.pending_bytes == 0
@pytest.mark.asyncio
async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
waiting: Final = asyncio.create_task(meter.record(_write("caller").event))
await database.entered.wait()
waiting.cancel()
with pytest.raises(asyncio.CancelledError):
await waiting
database.release.set()
await meter.close()
assert database.committed == ("caller",)
assert meter.pending_records == 0
@pytest.mark.asyncio
async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
first: Final = meter.submit((_write("in-flight"), _write("in-flight-two")))
await database.entered.wait()
queued: Final = meter.submit((_write("queued"),))
await meter.stop(0)
results: Final = await asyncio.gather(first, queued, return_exceptions=True)
assert all(isinstance(result, UsageCaptureError) for result in results)
assert meter.pending_records == 0
assert meter.pending_bytes == 0
assert database.committed == ()
@pytest.mark.asyncio
async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None:
retrying: Final = asyncio.Event()
blocked: Final = asyncio.Event()
async def unavailable(writes: tuple[UsageWrite, ...]) -> None:
raise ConnectionError("database unavailable")
async def backoff(delay: float) -> None:
retrying.set()
await blocked.wait()
meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff)
receipt: Final = meter.submit((_write("one"), _write("two")))
await retrying.wait()
await meter.stop(0)
with pytest.raises(UsageCaptureError, match="interrupted"):
await receipt
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
@pytest.mark.asyncio
@pytest.mark.parametrize("wrapped", [False, True])
async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None:
import httpx
from prisma.engine.errors import EngineRequestError
from prisma.http import Response
database: Final = CommitBarrier()
attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]()
retrying: Final = asyncio.Event()
recovered: Final = asyncio.Event()
closure: Final = EngineRequestError(
Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}'
)
failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure
if wrapped:
failure.__cause__ = closure
async def persist(writes: tuple[UsageWrite, ...]) -> None:
attempts.put_nowait(writes)
if attempts.qsize() == 1:
raise failure
await database.write(writes)
async def backoff(delay: float) -> None:
retrying.set()
await recovered.wait()
meter: Final = DurableUsageMeter(
persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff
)
writes: Final = (_write("caller"), _write("attempt"))
receipt: Final = meter.submit(writes)
retry_signal: Final = asyncio.create_task(retrying.wait())
reserved: Final = meter.pending_bytes
try:
await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED)
assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry"
assert (meter.pending_records, meter.pending_bytes) == (2, reserved)
with pytest.raises(UsageCaptureError, match="capacity"):
meter.submit((_write("overflow"),))
database.release.set()
recovered.set()
await receipt
assert attempts.get_nowait() == attempts.get_nowait() == writes
assert database.committed == ("caller", "attempt")
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
finally:
retry_signal.cancel()
await asyncio.gather(retry_signal, return_exceptions=True)
await meter.stop(0)
await asyncio.gather(receipt, return_exceptions=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("first_commits", [False, True])
async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None:
isolating: Final = asyncio.Event()
blocked: Final = asyncio.Event()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
if len(writes) > 1:
raise AccountingConflictError("one producer conflicts")
if first_commits and writes[0].event.event_id == "one":
return
isolating.set()
await blocked.wait()
meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2))
first: Final = meter.submit((_write("one"),))
second: Final = meter.submit((_write("two"),))
await isolating.wait()
await meter.stop(0)
results: Final = await asyncio.gather(first, second, return_exceptions=True)
assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError)
assert isinstance(results[1], UsageCaptureError)
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
@pytest.mark.asyncio
async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None:
from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch
database: Final = CommitBarrier()
database.release.set()
async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None:
assert timeout == 5.0
await _wait_for_batch(wakeup, 0)
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush)
receipt: Final = meter.submit((_write("partial-batch"),))
try:
await receipt
assert database.committed == ("partial-batch",)
assert database.transactions == 1
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
finally:
await meter.stop(0)
@pytest.mark.asyncio
async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None:
database: Final = CommitBarrier()
database.release.set()
waits: Final = asyncio.Queue[bool]()
async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None:
waits.put_nowait(wakeup.is_set())
await wakeup.wait()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait)
receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5))
try:
assert await waits.get() is True
await asyncio.gather(*receipts[:2])
assert await waits.get() is True, "A full queued batch must not wait for the flush deadline"
await asyncio.gather(*receipts[2:4])
assert await waits.get() is False
assert database.committed == ("0", "1", "2", "3")
assert meter.pending_records == 1
assert not receipts[-1].done()
await meter.close()
await receipts[-1]
assert database.committed == ("0", "1", "2", "3", "4")
assert database.transactions == 3
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
finally:
await meter.stop(0)
await asyncio.gather(*receipts, return_exceptions=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"])
async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None:
from prisma.errors import DataError
database: Final = CommitBarrier()
database.release.set()
attempts: Final = asyncio.Queue[tuple[str, ...]]()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
identities: Final = tuple(write.event.event_id for write in writes)
attempts.put_nowait(identities)
if "invalid" in identities or sqlstate == "42P01":
raise DataError(
data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}}
)
await database.write(writes)
meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3))
rejected: Final = meter.submit((_write("caller"), _write("invalid")))
neighbor: Final = meter.submit((_write("neighbor"),))
try:
with pytest.raises(DataError):
await rejected
if sqlstate == "42P01":
with pytest.raises(DataError):
await neighbor
assert attempts.qsize() == 1
else:
await neighbor
assert database.committed == ("neighbor",)
assert attempts.qsize() == 3
finally:
await meter.close()
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
@pytest.mark.asyncio
@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"])
async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4))
prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future()
dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,))
neighbor: Final = meter.submit((_write("neighbor"),))
assert meter.pending_records == 3
assert not dependent.done()
with pytest.raises(UsageCaptureError, match="capacity"):
meter.submit((_write("excess-one"), _write("excess-two")))
if dependency_outcome == "commit":
prerequisite.set_result(None)
elif dependency_outcome == "failure":
prerequisite.set_exception(ValueError("dependency failed"))
else:
prerequisite.cancel()
database.release.set()
await meter.close()
await neighbor
if dependency_outcome == "commit":
await dependent
assert database.committed == ("parent", "final", "neighbor")
else:
with pytest.raises(UsageCaptureError, match="dependency"):
await dependent
assert database.committed == ("neighbor",)
assert meter.pending_records == 0
assert meter.pending_bytes == 0
@pytest.mark.asyncio
async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future()
dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,))
await meter.stop(0)
with pytest.raises(UsageCaptureError, match="acknowledgement"):
await dependent
assert not prerequisite.done()
assert database.committed == ()
assert meter.pending_records == 0
assert meter.pending_bytes == 0
prerequisite.set_result(None)
@pytest.mark.asyncio
async def test_large_write_sequence_drains_within_byte_and_record_bounds() -> None:
database: Final = CommitBarrier()
database.release.set()
meter: Final = DurableUsageMeter(
database.write,
UsageCaptureSettings(batch_size=8, max_pending_records=2, max_pending_bytes=1024),
wait_for_batch=_manual_flush,
)
writes: Final = tuple(_write(str(index)) for index in range(20))
pending: Final = meter.submit((_write("pending-a"), _write("pending-b")))
await meter.record_writes(writes)
await pending
with pytest.raises(UsageCaptureError, match="One usage observation"):
await meter.record_writes((_write("oversized" * 300),))
await meter.close()
assert database.committed == ("pending-a", "pending-b", *(write.event.event_id for write in writes))
assert database.transactions > 1
assert (meter.pending_bytes, meter.pending_records) == (0, 0)
with pytest.raises(UsageCaptureError):
await meter.record_writes((_write("closed"),))
@pytest.mark.asyncio
async def test_cancelling_bounded_capture_preserves_admitted_writes_without_admitting_remaining_rows() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(
database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2), wait_for_batch=_manual_flush
)
recording: Final = asyncio.create_task(meter.record_writes(tuple(_write(str(index)) for index in range(8))))
await database.entered.wait()
recording.cancel()
with pytest.raises(asyncio.CancelledError):
await recording
database.release.set()
await meter.close()
assert database.committed == ("0", "1")
assert meter.pending_bytes == meter.pending_records == 0
@pytest.mark.asyncio
async def test_capacity_wait_retries_admission_before_other_producers_finish() -> None:
entered: Final = asyncio.Queue[tuple[str, ...]]()
release: Final = asyncio.Queue[None]()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
entered.put_nowait(tuple(write.event.event_id for write in writes))
await release.get()
meter: Final = DurableUsageMeter(
persist, UsageCaptureSettings(batch_size=2, max_pending_records=3), wait_for_batch=_manual_flush
)
initial: Final = meter.submit((_write("first"), _write("second")))
remaining: Final = meter.submit((_write("remaining"),))
assert await entered.get() == ("first", "second")
recording: Final = asyncio.create_task(meter.record_writes((_write("target"),)))
try:
await asyncio.sleep(0)
release.put_nowait(None)
assert await entered.get() == ("remaining",)
await asyncio.sleep(0)
assert meter.pending_records == 2
release.put_nowait(None)
assert await entered.get() == ("target",)
release.put_nowait(None)
await recording
await asyncio.gather(initial, remaining)
finally:
recording.cancel()
await asyncio.gather(recording, return_exceptions=True)
release.put_nowait(None)
release.put_nowait(None)
await meter.close()

View file

@ -268,76 +268,6 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm
token = MASTER_KEY
@pytest.mark.asyncio
async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from datetime import datetime
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from litellm.responses.streaming_iterator import spawn_logging_task
worker: Final = LoggingWorker()
started: Final = asyncio.Event()
release: Final = asyncio.Event()
first_flush: Final = asyncio.Event()
pending: Final[asyncio.Queue[str]] = asyncio.Queue()
persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
captured: Final[asyncio.Queue[str]] = asyncio.Queue()
async def persist_usage(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
captured.put_nowait(write.event.event_id)
capture: Final = UsageCaptureRuntime()
capture.start(persist_usage, UsageCaptureSettings(flush_interval_seconds=10))
monkeypatch.setattr(litellm.proxy.proxy_server, "usage_capture", capture)
async def enqueue_spend() -> None:
assert capture.meter is not None
capture_callback(
capture.meter,
{"call_type": "acompletion", "litellm_call_id": "shutdown-usage", "litellm_params": {}},
{"id": "shutdown-attempt", "usage": {"total_tokens": 3}},
datetime(2026, 1, 1),
datetime(2026, 1, 1),
)
pending.put_nowait("accounted request")
async def callback() -> None:
started.set()
await release.wait()
spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting")
async def flush_spend() -> None:
while not pending.empty():
persisted.put_nowait(pending.get_nowait())
first_flush.set()
worker.ensure_initialized_and_enqueue(callback())
await started.wait()
shutdown: Final = asyncio.create_task(
litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker)
)
try:
await first_flush.wait()
waited_for_callback: Final = not shutdown.done()
release.set()
await shutdown
assert waited_for_callback, "Shutdown returned while accounting callbacks were still running"
assert persisted.get_nowait() == "accounted request"
assert pending.empty()
assert (captured.get_nowait(), captured.get_nowait()) == ("shutdown-usage:inbound", "shutdown-attempt")
assert capture.meter is None
finally:
release.set()
await worker.flush()
await worker.stop()
await capture.stop(0)
@pytest.mark.asyncio
@ -4377,276 +4307,11 @@ async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() ->
await worker.stop(discard_pending=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("unsupported", ["database", "collector", "remote", "free_tier"])
async def test_capture_enablement_rejects_unrecoverable_writer_modes(
unsupported: str, monkeypatch: pytest.MonkeyPatch
) -> None:
import litellm
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
monkeypatch.setattr(proxy_server, "usage_capture", UsageCaptureRuntime())
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {}})
monkeypatch.setattr(proxy_server, "prisma_client", None if unsupported == "database" else object())
monkeypatch.setattr(litellm, "apply_search_free_tier", unsupported == "free_tier")
monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "true" if unsupported == "collector" else "false")
monkeypatch.delenv("SPEND_LOGS_URL", raising=False)
if unsupported == "remote":
monkeypatch.setenv("SPEND_LOGS_URL", "https://writer.invalid")
with pytest.raises(ValueError, match="requires local Prisma"):
proxy_server._configure_usage_capture()
assert proxy_server.usage_capture.meter is None
@pytest.mark.asyncio
async def test_capture_configuration_requires_restart_without_replacing_the_active_writer(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.utils import PrismaClient
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True})
with pytest.raises(ValueError, match="requires accounting_capture"):
proxy_server._configure_usage_capture()
client: Final = PrismaClient("postgresql://localhost/unused", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(proxy_server, "prisma_client", client)
monkeypatch.setattr(litellm, "apply_search_free_tier", False)
monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false")
monkeypatch.delenv("SPEND_LOGS_URL", raising=False)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 4}})
try:
proxy_server._configure_usage_capture()
active: Final = runtime.meter
assert active is not None
assert active.settings.batch_size == 4
proxy_server._configure_usage_capture()
assert runtime.meter is active
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 6}})
with pytest.raises(ValueError, match="Changing accounting_capture requires a worker restart"):
proxy_server._configure_usage_capture()
monkeypatch.setattr(proxy_server, "general_settings", {})
with pytest.raises(ValueError, match="Disabling accounting_capture requires a worker restart"):
proxy_server._configure_usage_capture()
assert runtime.meter is active
assert active.settings.batch_size == 4
finally:
await runtime.stop(1)
assert runtime.meter is None
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_ack", [False, True])
@pytest.mark.parametrize(
("path", "payload"),
[
(
"/v1/chat/completions",
{"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True},
),
("/v1/responses", {"model": "openai/gpt-5-mini", "input": "hi", "background": True}),
(
"/v1/messages",
{"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1},
),
],
)
async def test_sync_ack_preserves_async_request_support(
sync_ack: bool,
path: str,
payload: dict[str, object],
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
from fastapi import Request
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
persisted.put_nowait(write)
async def authorize(request: Request) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key")
upstream: Final = {
"id": "provider-result",
"object": "chat.completion",
"created": 1,
"model": "gpt-5-mini",
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
if payload.get("stream"):
chunk: Final = {
**upstream,
"object": "chat.completion.chunk",
"choices": [{"index": 0, "finish_reason": "stop", "delta": {"role": "assistant", "content": "hi"}}],
}
respx_mock.post().respond(
200, text=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n", headers={"content-type": "text/event-stream"}
)
elif path == "/v1/responses":
respx_mock.post().respond(
200,
json={
"id": "resp_background",
"object": "response",
"created_at": 1,
"model": "gpt-5-mini",
"status": "in_progress",
"background": True,
"output": [],
"usage": None,
},
)
else:
respx_mock.post(url__regex=r".*/responses$").respond(
200,
json={
"id": "resp_message",
"object": "response",
"created_at": 1,
"model": "gpt-5-mini",
"status": "completed",
"output": [
{
"id": "msg_reply",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "hi", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
},
)
respx_mock.post(url__regex=r".*/chat/completions$").respond(200, json=upstream)
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(
proxy_server, "general_settings", {"accounting_sync_ack": sync_ack, "accounting_sync_ack_timeout_seconds": 0}
)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(
proxy_server,
"llm_router",
litellm.Router(
model_list=[
{
"model_name": payload["model"],
"litellm_params": {"model": payload["model"], "api_key": "test-provider-key"},
}
],
num_retries=0,
),
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()])
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize)
runtime.start(persist, UsageCaptureSettings())
try:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test"
) as client:
response: Final = await client.post(path, json=payload)
finally:
await runtime.stop(1)
assert response.status_code == 200, response.text
assert len(respx_mock.calls) == 1
if payload.get("background"):
assert response.json()["status"] == "in_progress"
else:
assert "hi" in response.text
assert not persisted.empty(), "Terminal usage must still be captured"
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_ack", [False, True])
async def test_text_completion_provider_rejection_captures_unknown_usage(
sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
from fastapi import Request
from litellm.accounting import Outcome
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
persisted.put_nowait(write)
async def authorize(request: Request) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key")
respx_mock.post().respond(
404,
json={
"error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"}
},
)
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack})
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(
proxy_server,
"llm_router",
litellm.Router(
model_list=[
{
"model_name": "instruct",
"litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"},
}
],
num_retries=0,
),
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()])
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize)
runtime.start(persist, UsageCaptureSettings())
try:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test"
) as client:
response: Final = await client.post(
"/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1}
)
finally:
await runtime.stop(1)
assert response.status_code == 404, response.text
assert len(respx_mock.calls) == 1
assert persisted.qsize() == 2
inbound: Final = persisted.get_nowait()
outbound: Final = persisted.get_nowait()
assert outbound.event.outcome is Outcome.FAILED
assert outbound.event.call_type.value == "atext_completion"
assert outbound.event.metrics is None
assert inbound.event.request_id == outbound.event.request_id
envelope: Final = json.loads(outbound.envelope)
assert envelope["usage"] is None
assert envelope["usage_source"] == "unknown"
assert envelope["pricing"]["reported_cost"] is None
@pytest.mark.parametrize(

View file

@ -3735,42 +3735,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest
assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION])
@pytest.mark.asyncio
async def test_a_billed_responses_route_failure_captures_usage_from_the_newest_compaction_item() -> None:
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from litellm.proxy.utils import _failure_fields_to_lift
async def persist(_: tuple[UsageWrite, ...]) -> None:
return None
def captured_prompt_tokens(request_input: list[dict[str, str]]) -> int:
logging_obj: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params())
logging_obj.model_call_details["call_type"] = "aresponses"
logging_obj.model_call_details["messages"] = request_input
lifted: Final = _failure_fields_to_lift(
{"litellm_logging_obj": logging_obj}, provider_billed=True, route_call_type="aresponses"
)
snapshot: Final = lifted["_accounting_usage_snapshot"]
assert isinstance(snapshot, UsageCaptureObservation) and snapshot.usage_json is not None
return json.loads(snapshot.usage_json)["prompt_tokens"]
runtime: Final = UsageCaptureRuntime()
runtime.start(persist, UsageCaptureSettings())
try:
sent: Final = captured_prompt_tokens(
[COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION]
)
since_compaction: Final = captured_prompt_tokens(
[COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]
)
finally:
await runtime.stop(1)
assert sent == since_compaction
@pytest.mark.parametrize(
@ -4405,57 +4369,3 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None:
ProxyException(message="x", type="x", param=None, code="429")
) is True
assert _is_client_error_exception(ValueError("x")) is False
@pytest.mark.asyncio
@pytest.mark.parametrize("tokens", [None, 0, 3])
async def test_failure_capture_preserves_prepared_usage_observation(tokens: int | None) -> None:
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
timestamp: Final = real_datetime.datetime(2026, 1, 1)
committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
committed.put_nowait(write)
logger: Final = Logging(
model=FAILURE_USAGE_MODEL,
messages=[],
stream=False,
call_type="acompletion",
start_time=timestamp,
litellm_call_id="post-response-failure",
function_id="post-response-failure",
)
logger.update_environment_variables(
model=FAILURE_USAGE_MODEL, custom_llm_provider="openai", litellm_params={}, optional_params={}
)
logger.model_call_details["response_cost"] = 0.2
runtime: Final = UsageCaptureRuntime()
runtime.start(persist, UsageCaptureSettings())
try:
prepared: Final = logger.prepare_response_cost(
{"id": "provider", "usage": None if tokens is None else {"prompt_tokens": tokens, "total_tokens": tokens}},
timestamp,
)
assert prepared is not None
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True)
assert lifted["combined_usage_object"] is logger.model_call_details["combined_usage_object"]
assert lifted["response_cost"] == 0.2
assert runtime.meter is not None
capture_callback(
runtime.meter,
{**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"},
None,
timestamp,
timestamp,
)
finally:
await runtime.stop(1)
committed.get_nowait()
envelope: Final = json.loads(committed.get_nowait().envelope)
assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == tokens
assert envelope["usage_source"] == ("unknown" if tokens is None else "measured")

View file

@ -11547,3 +11547,28 @@ async def test_bedrock_strict_tools_limit_rejection_is_not_sent_again_when_no_to
assert error is not None and "Too many strict tools" in str(error)
assert len(sent) == 1
def test_supplied_openai_client_base_url_decides_data_residency(monkeypatch: pytest.MonkeyPatch) -> None:
"""A caller-supplied OpenAI client sends the request to its own base_url, so billing must use that
host and not an env-configured FedRAMP or regional base."""
import openai
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setenv("OPENAI_BASE_URL", "https://gov.api.openai.com/v1")
captured = {}
real_init = litellm.Logging.__init__
def init_spy(self, *args, **kwargs):
real_init(self, *args, **kwargs)
captured["logging_obj"] = self
with patch.object(litellm.Logging, "__init__", init_spy):
litellm.completion(
model="gpt-4.1",
messages=[{"role": "user", "content": "hi"}],
client=openai.OpenAI(api_key="test-key", base_url="https://api.openai.com/v1"),
mock_response="ok",
)
assert captured["logging_obj"].litellm_params.get("data_residency") is None

View file

@ -1019,3 +1019,144 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy
assert while_a_is_down._hidden_params["model_id"] == "B"
assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
assert after_a_recovers._hidden_params["model_id"] == "A"
def _two_deployment_router(model: str) -> Router:
return Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": model, "api_key": "bad", "weight": 1},
"model_info": {"id": "A"},
},
{
"model_name": "test-model",
"litellm_params": {"model": model, "api_key": "good", "weight": 0},
"model_info": {"id": "B"},
},
],
routing_strategy="simple-shuffle",
num_retries=0,
disable_cooldowns=True,
enable_weighted_failover=True,
)
_ENDPOINTS: Final = [
("aembedding", "openai/text-embedding-3-small", {"input": ["hi"]}),
("aimage_generation", "openai/dall-e-3", {"prompt": "a cat"}),
("atranscription", "openai/whisper-1", {"file": b"audio"}),
("aspeech", "openai/tts-1", {"input": "hi", "voice": "alloy"}),
("atext_completion", "openai/gpt-3.5-turbo-instruct", {"prompt": "hi"}),
("arerank", "cohere/rerank-english-v3.0", {"query": "q", "documents": ["d"]}),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(("endpoint", "model", "arguments"), _ENDPOINTS, ids=[e[0] for e in _ENDPOINTS])
async def test_weighted_failover_moves_every_endpoint_to_another_deployment(
monkeypatch: pytest.MonkeyPatch, endpoint: str, model: str, arguments: dict[str, object]
) -> None:
keys_called: list[object] = []
answer: Final = object()
async def provider(**kwargs: object) -> object:
keys_called.append(kwargs.get("api_key"))
if kwargs.get("api_key") == "bad":
raise litellm.InternalServerError(message="A is down", llm_provider="openai", model=model)
return answer
monkeypatch.setattr(litellm, endpoint, provider)
router: Final = _two_deployment_router(model)
response: Final = await getattr(router, endpoint)(model="test-model", **arguments)
assert response is answer
assert keys_called == ["bad", "good"]
@pytest.mark.asyncio
async def test_weighted_failover_skips_the_failed_deployment_for_clientside_credentials(
monkeypatch: pytest.MonkeyPatch,
) -> None:
api_bases_called: list[object] = []
answer: Final = object()
async def provider(**kwargs: object) -> object:
api_bases_called.append(kwargs.get("api_base"))
if kwargs.get("api_base") == "http://a":
raise litellm.InternalServerError(message="A is down", llm_provider="openai", model="m")
return answer
monkeypatch.setattr(litellm, "aembedding", provider)
router: Final = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://a", "weight": 1},
"model_info": {"id": "A"},
},
{
"model_name": "test-model",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://b", "weight": 0},
"model_info": {"id": "B"},
},
],
routing_strategy="simple-shuffle",
num_retries=0,
disable_cooldowns=True,
enable_weighted_failover=True,
)
response: Final = await router.aembedding(model="test-model", input=["hi"], api_key="tenant-key")
assert response is answer
assert api_bases_called == ["http://a", "http://b"]
@pytest.mark.parametrize(
("endpoint", "attempt", "arguments"),
[
("embedding", "_embedding", {"input": ["hi"]}),
("image_generation", "_image_generation", {"prompt": "a cat"}),
],
)
def test_a_sync_attempt_stamps_the_failed_deployment(
monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object]
) -> None:
def provider(**kwargs: object) -> object:
raise litellm.InternalServerError(message="down", llm_provider="openai", model="m")
monkeypatch.setattr(litellm, endpoint, provider)
router: Final = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": "openai/m", "api_key": "k"},
"model_info": {"id": "A"},
}
],
)
with pytest.raises(litellm.InternalServerError) as raised:
getattr(router, attempt)(model="test-model", **arguments)
assert getattr(raised.value, "failed_deployment_id", None) == "A"
@pytest.mark.asyncio
async def test_a_failed_pick_stamps_no_deployment() -> None:
router: Final = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "k"},
"model_info": {"id": "A"},
}
],
)
with pytest.raises(litellm.BadRequestError, match="no healthy deployments") as raised:
await router._aembedding(input=["hi"], model="no-such-model")
assert getattr(raised.value, "failed_deployment_id", None) is None

View file

@ -1082,6 +1082,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
"regional_endpoint_uplift_multiplier": {"type": "number"},
"regional_processing_uplift_multiplier_eu": {"type": "number"},
"regional_processing_uplift_multiplier_us": {"type": "number"},
"fedramp_uplift_multiplier": {"type": "number"},
"input_cost_per_pixel": {"type": "number"},
"input_cost_per_query": {"type": "number"},
"input_cost_per_request": {"type": "number"},
@ -4243,23 +4244,29 @@ def test_speech_stream_is_logged_when_the_audio_ends(
assert audio.rest_sent_at is not None and _logged_end(logged) >= audio.rest_sent_at
class _InlineExecutor:
def submit(self, fn: Callable[..., object], /, *args: object, **kwargs: object) -> Future[object]:
future: Final = Future[object]()
future.set_result(fn(*args, **kwargs))
return future
@pytest.mark.usefixtures("local_model_cost_map")
def test_speech_stream_usage_is_captured_when_the_audio_ends(
def test_speech_stream_success_handler_runs_when_the_audio_ends(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
audio: Final = _SyncAudio(b"first", b"rest")
_mock_elevenlabs_audio(respx_mock, audio)
captured_at: Final[queue.Queue[float]] = queue.Queue()
monkeypatch.setattr(
litellm_logging.Logging, "usage_capture_callback", lambda *_: captured_at.put(time.time())
)
speech_logger: Final = _SyncSpeechLogger()
monkeypatch.setattr(litellm, "callbacks", [speech_logger])
monkeypatch.setattr("litellm.utils.executor", _InlineExecutor())
response: Final = litellm.speech(**_SPEECH_ARGS, stream=True)
captured_before_read: Final = captured_at.qsize()
logged_at_headers: Final = speech_logger.logged_payloads.qsize()
assert response.read() == b"firstrest"
assert captured_before_read == 0
assert audio.rest_sent_at is not None and captured_at.get(timeout=5) >= audio.rest_sent_at
assert logged_at_headers == 0
assert speech_logger.logged_payloads.get_nowait()["response"] == _audio_summary(len(b"firstrest"))
class TestCallbackAsyncSyncSeparation:

View file

@ -5,8 +5,8 @@ import { describe, expect, it, vi } from "vitest";
import { RequestTagMapping, RequestTagMappingsEditor, requestTagMappingsAreValid } from "./RequestTagMappings";
const exact: RequestTagMapping = { header: "baggage", key: "entrypoint", value: "a", tag: "support" };
const variable: RequestTagMapping = { header: "baggage", key: "entrypoint", tag_prefix: "usecase:" };
const exact: RequestTagMapping = { header: "baggage", key: "service", value: "a", tag: "support" };
const variable: RequestTagMapping = { header: "baggage", key: "service", tag_prefix: "app:" };
function Editor({
initial,
@ -47,20 +47,20 @@ describe("RequestTagMappingsEditor", () => {
expect(screen.getByRole("combobox", { name: "Mapping 2 mode" })).toHaveTextContent("Any value");
expect(screen.getByLabelText("Mapping 1 match value")).toHaveValue("a");
expect(screen.getByLabelText("Mapping 1 tag")).toHaveValue("support");
expect(screen.getByLabelText("Mapping 2 tag prefix")).toHaveValue("usecase:");
expect(screen.getByText("entrypoint=a → support")).toBeInTheDocument();
expect(screen.getByText("entrypoint=x → usecase:x")).toBeInTheDocument();
expect(screen.getByLabelText("Mapping 2 tag prefix")).toHaveValue("app:");
expect(screen.getByText("service=a → support")).toBeInTheDocument();
expect(screen.getByText("service=x → app:x")).toBeInTheDocument();
});
it("adds and edits a new exact value mapping", async () => {
const changed = vi.fn();
render(<Editor initial={[]} changed={changed} />);
await userEvent.click(screen.getByRole("button", { name: "Add mapping" }));
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } });
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } });
expect(screen.getByText("entrypoint=b → sales")).toBeInTheDocument();
expect(changed).toHaveBeenLastCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]);
expect(screen.getByText("service=b → sales")).toBeInTheDocument();
expect(changed).toHaveBeenLastCalledWith([{ header: "baggage", key: "service", value: "b", tag: "sales" }]);
});
it("changes an exact mapping to any-value mode and removes its exact match fields", async () => {
@ -70,7 +70,7 @@ describe("RequestTagMappingsEditor", () => {
await userEvent.click(await screen.findByRole("option", { name: "Any value" }));
expect(screen.queryByLabelText("Mapping 1 match value")).not.toBeInTheDocument();
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveAttribute("placeholder", "usecase:");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveAttribute("placeholder", "app:");
expect(changed).toHaveBeenLastCalledWith([
expect.objectContaining({ value: undefined, tag: undefined, tag_prefix: "" }),
]);
@ -81,9 +81,9 @@ describe("RequestTagMappingsEditor", () => {
render(<Editor initial={[]} changed={vi.fn()} saved={saved} />);
await userEvent.click(screen.getByRole("button", { name: "Add mapping" }));
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveValue("");
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveAttribute("placeholder", "entrypoint");
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveAttribute("placeholder", "service");
expect(screen.getByText("Enter a baggage key")).toBeInTheDocument();
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
expect(screen.getByText("Enter a match value")).toBeInTheDocument();
await userEvent.click(screen.getByRole("button", { name: "Update" }));
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } });
@ -93,7 +93,7 @@ describe("RequestTagMappingsEditor", () => {
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } });
expect(screen.queryByText(/^Enter a/)).not.toBeInTheDocument();
await userEvent.click(screen.getByRole("button", { name: "Update" }));
expect(saved).toHaveBeenCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]);
expect(saved).toHaveBeenCalledWith([{ header: "baggage", key: "service", value: "b", tag: "sales" }]);
});
it("blocks saving an any-value mapping until the admin types a tag prefix", async () => {
@ -105,7 +105,7 @@ describe("RequestTagMappingsEditor", () => {
expect(screen.getByRole("button", { name: "Update" })).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 tag prefix"), { target: { value: "team:" } });
await userEvent.click(screen.getByRole("button", { name: "Update" }));
expect(saved).toHaveBeenCalledWith([expect.objectContaining({ key: "entrypoint", tag_prefix: "team:" })]);
expect(saved).toHaveBeenCalledWith([expect.objectContaining({ key: "service", tag_prefix: "team:" })]);
});
it("blocks saving a baggage key that is empty or not a header token", async () => {
@ -139,7 +139,7 @@ describe("RequestTagMappingsEditor", () => {
const changed = vi.fn();
render(<Editor initial={[exact, variable]} changed={changed} />);
await userEvent.click(screen.getByRole("button", { name: "Delete mapping 1" }));
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("usecase:");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("app:");
expect(changed).toHaveBeenLastCalledWith([variable]);
});
});

View file

@ -111,7 +111,7 @@ export function RequestTagMappingsEditor({ value, onChange, disabled = false }:
{t("routerSettings.requestTagMappings.baggageKey")}
<Input
aria-label={t("routerSettings.requestTagMappings.baggageKeyAriaLabel", { number: index + 1 })}
placeholder="entrypoint"
placeholder="service"
value={rule.key}
onChange={(event) => update(index, { key: event.target.value })}
/>
@ -145,7 +145,7 @@ export function RequestTagMappingsEditor({ value, onChange, disabled = false }:
{t("routerSettings.requestTagMappings.tagPrefix")}
<Input
aria-label={t("routerSettings.requestTagMappings.tagPrefixAriaLabel", { number: index + 1 })}
placeholder="usecase:"
placeholder="app:"
value={rule.tag_prefix ?? ""}
onChange={(event) => update(index, { tag_prefix: event.target.value })}
/>

View file

@ -310,12 +310,12 @@ it("keeps a new baggage mapping unsaved until every field the backend requires i
name: "Update",
});
expect(update).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "a" } });
expect(update).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "usecase:support" } });
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "app:support" } });
await user.click(update);
expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([
["token", "request_tag_mappings", [{ header: "baggage", key: "entrypoint", value: "a", tag: "usecase:support" }]],
["token", "request_tag_mappings", [{ header: "baggage", key: "service", value: "a", tag: "app:support" }]],
]);
});

View file

@ -252,6 +252,7 @@ describe("provider_info_helpers", () => {
Providers.DATAROBOT,
Providers.DOCKER_MODEL_RUNNER,
Providers.DOTPROMPT,
Providers.EMPIRIOLABS,
Providers.EMPOWER,
Providers.GALADRIEL,
Providers.GradientAI,
@ -616,6 +617,24 @@ describe("provider_info_helpers", () => {
]);
});
it("should list EmpirioLabs AI chat and decision models for the 'EMPIRIOLABS' provider key", () => {
const modelMap = {
"empiriolabs/aplomb-1": {
litellm_provider: "empiriolabs",
mode: "chat",
supported_endpoints: ["/v1/chat/completions", "/v1/decisions"],
},
"empiriolabs/qwen3-7-plus": { litellm_provider: "empiriolabs", mode: "chat" },
"typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" },
};
expect(getProviderModels("EMPIRIOLABS" as Providers, modelMap)).toEqual([
"empiriolabs/aplomb-1",
"empiriolabs/qwen3-7-plus",
]);
expect(getProviderLogoAndName("empiriolabs").displayName).toBe("EmpirioLabs AI");
expect(getPlaceholder("EMPIRIOLABS")).toBe("empiriolabs/aplomb-1");
});
it("should list only strands_decider models for the 'StrandsDecider' provider key", () => {
const modelMap = {
"strands_decider/strands-decider-2B-hobson-v19": { litellm_provider: "strands_decider", mode: "decisions" },

View file

@ -118,6 +118,7 @@ export enum Providers {
DOTPROMPT = "Dotprompt",
EDENAI = "Eden AI",
ElevenLabs = "ElevenLabs",
EMPIRIOLABS = "EmpirioLabs AI",
EMPOWER = "Empower",
FalAI = "Fal AI",
FEATHERLESS_AI = "Featherless Ai",
@ -252,6 +253,7 @@ export const provider_map: Record<string, string> = {
DOTPROMPT: "dotprompt",
EDENAI: "edenai",
ElevenLabs: "elevenlabs",
EMPIRIOLABS: "empiriolabs",
EMPOWER: "empower",
FalAI: "fal_ai",
FEATHERLESS_AI: "featherless_ai",
@ -499,6 +501,7 @@ const providerPlaceholderMap: Partial<Record<Providers, string>> = {
[Providers.Cursor]: "cursor/claude-4-sonnet",
[Providers.DeepInfra]: "deepinfra/<any-model-on-deepinfra>",
[Providers.EDENAI]: "edenai/openai/gpt-mini-latest",
[Providers.EMPIRIOLABS]: "empiriolabs/aplomb-1",
[Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra",
[Providers.Google_AI_Studio]: "gemini-pro",
[Providers.HOSTINGER_ROUTER]: "hostinger_router/claude-sonnet-5",

View file

@ -30261,27 +30261,6 @@ export interface components {
* @description Documents all the fields supported by `general_settings` in config.yaml
*/
ConfigGeneralSettings: {
/** @description Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes. */
accounting_capture?: components["schemas"]["UsageCaptureSettings"] | null;
/**
* Accounting Sync Ack
* @description Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. Other requests remain supported with asynchronous usage capture.
* @default false
*/
accounting_sync_ack: boolean;
/**
* Accounting Sync Ack Failure Policy
* @description Return 503 or log the persistence failure after the response wait expires
* @default fail_closed
* @enum {string}
*/
accounting_sync_ack_failure_policy: "fail_closed" | "log";
/**
* Accounting Sync Ack Timeout Seconds
* @description Maximum response wait for usage commit, including the batch flush interval
* @default 20
*/
accounting_sync_ack_timeout_seconds: number;
/**
* Admission Queue Timeout Seconds
* @description maximum time a request waits for a worker slot
@ -35973,6 +35952,8 @@ export interface components {
drop_params?: boolean | string | null;
/** Explicit Cache Read Input Token Cost */
explicit_cache_read_input_token_cost?: number | null;
/** Fedramp Uplift Multiplier */
fedramp_uplift_multiplier?: number | null;
/**
* Fireworks Forward User Id
* @description Send the LiteLLM user id of the calling key as the `user` field on Fireworks AI chat, responses and messages requests.
@ -49981,29 +49962,6 @@ export interface components {
*/
version: string | null;
};
/** UsageCaptureSettings */
UsageCaptureSettings: {
/**
* Batch Size
* @default 1000
*/
batch_size: number;
/**
* Flush Interval Seconds
* @default 5
*/
flush_interval_seconds: number;
/**
* Max Pending Bytes
* @default 16777216
*/
max_pending_bytes: number;
/**
* Max Pending Records
* @default 10000
*/
max_pending_records: number;
};
/** UsageChartPoint */
UsageChartPoint: {
/** Blocked */
@ -51677,6 +51635,8 @@ export interface components {
drop_params?: boolean | string | null;
/** Explicit Cache Read Input Token Cost */
explicit_cache_read_input_token_cost?: number | null;
/** Fedramp Uplift Multiplier */
fedramp_uplift_multiplier?: number | null;
/**
* Fireworks Forward User Id
* @description Send the LiteLLM user id of the calling key as the `user` field on Fireworks AI chat, responses and messages requests.