mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge branch 'main' into litellm_tts_true_streaming
This commit is contained in:
commit
5fa3c58bac
98 changed files with 2528 additions and 7319 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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) | ✅ | ✅ | ✅ | | ✅ | | | | | |
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
304
litellm/accounting/projection.py
Normal file
304
litellm/accounting/projection.py
Normal 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
|
||||
|
|
@ -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."""
|
||||
...
|
||||
|
|
|
|||
97
litellm/accounting/snapshot.py
Normal file
97
litellm/accounting/snapshot.py
Normal 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)),
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
10
litellm/llms/empiriolabs/decisions/transformation.py
Normal file
10
litellm/llms/empiriolabs/decisions/transformation.py
Normal 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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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
|
||||
|
|
@ -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})
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -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"),
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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)
|
||||
|
|
@ -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 = (
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
431
tests/unit/accounting/test_projection.py
Normal file
431
tests/unit/accounting/test_projection.py
Normal 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"
|
||||
]
|
||||
|
|
@ -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),
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
0
tests/unit/llms/empiriolabs/__init__.py
Normal file
0
tests/unit/llms/empiriolabs/__init__.py
Normal file
0
tests/unit/llms/empiriolabs/decisions/__init__.py
Normal file
0
tests/unit/llms/empiriolabs/decisions/__init__.py
Normal 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"
|
||||
)
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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()
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 })}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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" }]],
|
||||
]);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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" },
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
48
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
48
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue