mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(auto_router): report a session's whole spend and name why savings stay unestimated (#51066)
* feat(auto_router): report a session's whole spend on /auto_router/session Claude Code sends its background calls straight to Haiku under the same session id, so those requests never reached the auto-router session rollup and the status line under-reported the session. Every billed spend row in a caller-supplied session now increments a key-scoped Redis counter whose one-hour expiry each write resets, and the endpoint reports it as total_spend, with spend unchanged. The increment runs where the spend row is accepted, so it counts each row exactly as often as the key's spend counters do. It reads as no total without a shared Redis * fix(auto_router): name the unsupported provider when a turn routes off Anthropic A turn routed to a non-Anthropic tier with no native wire response, such as a Cloudflare DeepSeek model, failed strict native validation and was logged as observation_unavailable, which reads like a broken capture. It now reports unsupported_selected_provider. Turns with a native wire response keep their existing reasons, and no estimate is added, so savings for these turns stay unknown * docs(auto_router): explain when session savings stay unestimated saved_spend on /auto_router/session now says a turn inside the longest prompt-cache TTL of the session's first observed turn, or of its last incomplete observation, is not estimated, so a null there is expected for that window * fix(auto_router): say why a system-only cache prefix has no savings estimate Claude Code puts its cache breakpoint on the system prompt, so the baseline prefix count sends Anthropic a body with no messages, which it rejects, and every turn read as token_count_unavailable. A count that fails only on system or tools prefixes now reports uncountable_system_or_tools_prefix. Gateways that accept the empty body still count it. A live check against claude-opus-5-5 showed a differential count with a placeholder message is off by -2 tokens for system-only and +68 for tools, so these turns stay unestimated and the savings gap for Claude Code is not fixed here * fix(auto_router): declare the session total's Redis key family and test its expiry The session total read and wrote Redis without a service target, so its calls showed up as bare redis spans. Both now run under session_spend_totals. The key no longer hashes the api key, which is already a hashed token, and only digests the session id. A new test drives Redis expiry on a test clock to show an active session keeps its total and an idle one restarts * docs(auto_router): note which turns are still estimated inside the history window Turns that took the identical baseline path, and turns whose whole cached prefix the session had already written, are estimated even inside the first cache TTL window, so saved_spend now says so --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
2605aa1c84
commit
076ded32be
15 changed files with 461 additions and 15 deletions
|
|
@ -1713,6 +1713,7 @@ RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: Final = "responses_websocket_call_id"
|
|||
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
|
||||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
|
||||
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
|
||||
SESSION_SPEND_TOTAL_TTL_SECONDS: Final = 3600
|
||||
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
|
||||
CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id"
|
||||
COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: Final = "_complexity_router_child_session_id"
|
||||
|
|
|
|||
|
|
@ -208,6 +208,7 @@ class UnsupportedCachePlan:
|
|||
"invalid_cache_ttl_order",
|
||||
"unsupported_thinking_cache_semantics",
|
||||
"token_count_unavailable",
|
||||
"uncountable_system_or_tools_prefix",
|
||||
"inconsistent_prefix_token_count",
|
||||
]
|
||||
|
||||
|
|
@ -583,7 +584,10 @@ async def count_cache_plan(
|
|||
counts: Final = tuple(
|
||||
await asyncio.gather(*(token_counter(model, api_key, marker.prefix_body) for marker in plan.breakpoints))
|
||||
)
|
||||
if any(value is None for value in counts):
|
||||
failed: Final = tuple(marker for marker, value in zip(plan.breakpoints, counts, strict=True) if value is None)
|
||||
if failed and all(not marker.prefix_body.get("messages") for marker in failed):
|
||||
return UnsupportedCachePlan("uncountable_system_or_tools_prefix")
|
||||
if failed:
|
||||
return UnsupportedCachePlan("token_count_unavailable")
|
||||
known: Final = tuple(value for value in counts if value is not None)
|
||||
if any(value < 0 for value in (total, *known)) or any(
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ from litellm.proxy.spend_tracking.savings import (
|
|||
extract_cache_read_tokens,
|
||||
marks_gateway_injection,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
|
||||
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
|
||||
from litellm.repositories.prisma_protocols import BatchTable
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
|
|
@ -685,8 +686,10 @@ class DBSpendUpdateWriter:
|
|||
self,
|
||||
redis_cache: RedisCache | None = None,
|
||||
accounting_clock: Callable[[], float] = time.monotonic,
|
||||
session_spend_totals: SessionSpendTotals | None = None,
|
||||
):
|
||||
self.redis_cache = redis_cache
|
||||
self.session_spend_totals: Final = session_spend_totals
|
||||
self.accounting_storage_checks = Cadence(accounting_clock)
|
||||
self.accounting_delay_alerts = Cadence(accounting_clock)
|
||||
self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache)
|
||||
|
|
@ -819,6 +822,7 @@ class DBSpendUpdateWriter:
|
|||
prisma_client=prisma_client,
|
||||
spend_logs_kept=keep_spend_logs,
|
||||
)
|
||||
await self._record_session_spend(hashed_token, billable_cost or 0.0, request_fields)
|
||||
if keep_spend_logs:
|
||||
await self._enqueue_tool_usage_transaction(
|
||||
payload=payload,
|
||||
|
|
@ -874,6 +878,19 @@ class DBSpendUpdateWriter:
|
|||
)
|
||||
return True
|
||||
|
||||
async def _record_session_spend(
|
||||
self, api_key: str | None, spend: float, request_fields: Mapping[str, object]
|
||||
) -> None:
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import caller_session_for_accounting
|
||||
|
||||
totals: Final = self.session_spend_totals
|
||||
if totals is None or not api_key or spend <= 0:
|
||||
return
|
||||
session_id: Final = caller_session_for_accounting(request_fields)
|
||||
if session_id is None:
|
||||
return
|
||||
await totals.record(api_key=api_key, session_id=session_id, spend=spend)
|
||||
|
||||
async def _record_spend_log(
|
||||
self,
|
||||
payload: SpendLogsPayload,
|
||||
|
|
|
|||
|
|
@ -312,6 +312,18 @@ async def _capture_native(
|
|||
) -> CapturedBaselineObservation:
|
||||
capture: Final = context.capture
|
||||
original: Final = capture.observation
|
||||
details: Final = logging_obj.model_call_details
|
||||
if details.get("custom_llm_provider") != "anthropic" and not isinstance(
|
||||
details.get("httpx_response"), httpx.Response
|
||||
):
|
||||
return capture.with_observation(
|
||||
original.model_copy(
|
||||
update={
|
||||
"available_at": max(original.started_at, context.collector.clock()),
|
||||
"reason": context.invalidated or "unsupported_selected_provider",
|
||||
}
|
||||
)
|
||||
)
|
||||
event: Final = _WireEvent.model_validate(logging_obj.model_call_details)
|
||||
wire: Final = event.httpx_response.request
|
||||
usage: Final = _ResponseUsage.model_validate(response_obj).usage
|
||||
|
|
|
|||
|
|
@ -1052,7 +1052,7 @@ async def get_auto_router_session(
|
|||
fills a moment after each turn; a session with no flushed auto-routed turn yet is a 404. The
|
||||
id is bounded the way the writer bounded it, so an oversized client id still finds its row.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
|
@ -1063,6 +1063,11 @@ async def get_auto_router_session(
|
|||
raise HTTPException(
|
||||
status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key"
|
||||
)
|
||||
total: Final = (
|
||||
await proxy_logging_obj.session_spend_totals.read(user_api_key_dict.api_key, session_id)
|
||||
if user_api_key_dict.api_key
|
||||
else None
|
||||
)
|
||||
saved_spend, baseline_spend = _savings_cohort(row.turns, row.savings_estimated_turns, row.spend, row.saved_spend)
|
||||
_, estimated_baseline_spend = _savings_cohort(
|
||||
row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend
|
||||
|
|
@ -1074,6 +1079,7 @@ async def get_auto_router_session(
|
|||
turns=row.turns,
|
||||
last_model=row.last_model,
|
||||
spend=row.spend,
|
||||
total_spend=total,
|
||||
savings_estimated_turns=row.savings_estimated_turns,
|
||||
savings_estimated_actual_spend=row.savings_estimated_actual_spend,
|
||||
saved_spend=saved_spend,
|
||||
|
|
|
|||
59
litellm/proxy/spend_tracking/session_spend_total.py
Normal file
59
litellm/proxy/spend_tracking/session_spend_total.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import hashlib
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.redis_cache import log_redis_failure
|
||||
from litellm.constants import SESSION_SPEND_TOTAL_TTL_SECONDS
|
||||
|
||||
_STORED_SPEND: Final[TypeAdapter[float | str | bytes | None]] = TypeAdapter(float | str | bytes | None)
|
||||
|
||||
|
||||
class SharedCounterStore(Protocol):
|
||||
async def async_increment(
|
||||
self, key: str, value: float, *, ttl: int | None = None, refresh_ttl: bool = False
|
||||
) -> float: ...
|
||||
|
||||
async def async_get_cache(self, key: str) -> object: ...
|
||||
|
||||
|
||||
def _key(api_key: str, session_id: str) -> str:
|
||||
session_digest: Final = hashlib.sha256(session_id.encode("utf-8", errors="surrogatepass")).hexdigest()
|
||||
return f"session_spend_total:{api_key}:{session_digest}"
|
||||
|
||||
|
||||
class SessionSpendTotals:
|
||||
"""Spend billed to one key's caller-supplied session, kept in the shared Redis on a TTL that every write slides."""
|
||||
|
||||
def __init__(self, shared_store: Callable[[], SharedCounterStore | None]) -> None:
|
||||
self.shared_store: Final = shared_store
|
||||
|
||||
@with_service_target("session_spend_totals")
|
||||
async def record(self, api_key: str, session_id: str, spend: float) -> None:
|
||||
store: Final = self.shared_store()
|
||||
if store is None:
|
||||
return
|
||||
try:
|
||||
await store.async_increment(
|
||||
_key(api_key, session_id), spend, ttl=SESSION_SPEND_TOTAL_TTL_SECONDS, refresh_ttl=True
|
||||
)
|
||||
except Exception as record_error: # noqa: BLE001 # a lost session total must never fail the spend write
|
||||
log_redis_failure(verbose_proxy_logger, logging.WARNING, "Session spend total write failed", record_error)
|
||||
|
||||
@with_service_target("session_spend_totals")
|
||||
async def read(self, api_key: str, session_id: str) -> float | None:
|
||||
store: Final = self.shared_store()
|
||||
if store is None:
|
||||
return None
|
||||
try:
|
||||
stored: Final = _STORED_SPEND.validate_python(await store.async_get_cache(_key(api_key, session_id)))
|
||||
except Exception as read_error: # noqa: BLE001 # an unreadable total is reported as no total
|
||||
log_redis_failure(verbose_proxy_logger, logging.WARNING, "Session spend total read failed", read_error)
|
||||
return None
|
||||
if stored is None:
|
||||
return None
|
||||
return float(stored.decode() if isinstance(stored, bytes) else stored)
|
||||
|
|
@ -37,7 +37,7 @@ from litellm.constants import (
|
|||
MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
|
||||
)
|
||||
from litellm.integrations.otel.model.metadata import caller_session_id
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
|
||||
from litellm.litellm_core_utils.batch_utils import (
|
||||
BATCH_COST_REQUEST_ID_SUFFIX,
|
||||
charged_poll_cost,
|
||||
|
|
@ -1050,12 +1050,31 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -
|
|||
return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX)
|
||||
|
||||
|
||||
def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None:
|
||||
def caller_session_for_accounting(kwargs: Mapping[str, object]) -> str | None:
|
||||
params: Final = as_str_mapping(kwargs.get("litellm_params")) or MappingProxyType({})
|
||||
preserved: Final = next(
|
||||
(
|
||||
session
|
||||
for key in ("litellm_metadata", "metadata")
|
||||
if isinstance(session := (as_str_mapping(params.get(key)) or {}).get(CALLER_SESSION_ID_METADATA_KEY), str)
|
||||
),
|
||||
None,
|
||||
)
|
||||
return _caller_session_for_spend_log(
|
||||
kwargs,
|
||||
MappingProxyType({CALLER_SESSION_ID_METADATA_KEY: preserved}) if preserved else None,
|
||||
trace=caller_trace_controls(kwargs),
|
||||
)
|
||||
|
||||
|
||||
def _caller_session_for_spend_log(
|
||||
kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, trace: TraceControls | None = None
|
||||
) -> str | None:
|
||||
beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None
|
||||
caller_session: Final = (
|
||||
beside_trace_header
|
||||
if isinstance(beside_trace_header, str) and beside_trace_header
|
||||
else caller_session_id(kwargs, TraceControls())
|
||||
else caller_session_id(kwargs, trace or TraceControls())
|
||||
)
|
||||
return caller_session if caller_session and "\x00" not in caller_session else None
|
||||
|
||||
|
|
|
|||
|
|
@ -285,6 +285,7 @@ from litellm.proxy.policy_engine.pipeline_executor import (
|
|||
)
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
||||
from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.spend_log_insert import RejectedSpendLogRow
|
||||
|
|
@ -1983,7 +1984,8 @@ class ProxyLogging:
|
|||
)
|
||||
self.premium_user = premium_user
|
||||
self.service_logging_obj = ServiceLogging()
|
||||
self.db_spend_update_writer = DBSpendUpdateWriter()
|
||||
self.session_spend_totals: Final = SessionSpendTotals(lambda: self.internal_usage_cache.dual_cache.redis_cache)
|
||||
self.db_spend_update_writer = DBSpendUpdateWriter(session_spend_totals=self.session_spend_totals)
|
||||
self.proxy_hook_mapping: dict[str, CustomLogger] = {}
|
||||
|
||||
# Guard flags to prevent duplicate background tasks
|
||||
|
|
|
|||
|
|
@ -176,12 +176,24 @@ class AutoRouterSessionResponse(LiteLLMBaseModel):
|
|||
turns: int = Field(description="Auto-routed turns the rollup has recorded for this session so far")
|
||||
last_model: str = Field(description="The deployment model the most recent turn was routed to")
|
||||
spend: float = Field(description="What the session's routed traffic actually cost, classifier calls included")
|
||||
total_spend: float | None = Field(
|
||||
default=None,
|
||||
description="Spend billed to this key under the session id, routed or not, from a shared Redis counter "
|
||||
"whose one-hour expiry every billed request resets. None without a shared Redis or once the counter "
|
||||
"expired. A best-effort figure, not a lifetime total: a session idle past the expiry restarts from zero, "
|
||||
"and a failed Redis write, eviction or flush drops spend without notice",
|
||||
)
|
||||
savings_estimated_turns: int = Field(description="Requests whose savings estimate recorded its baseline cost")
|
||||
savings_estimated_actual_spend: float = Field(
|
||||
description="Actual spend, including classifier cost, for requests whose estimate recorded its baseline cost"
|
||||
)
|
||||
saved_spend: float | None = Field(
|
||||
description="Recorded historical savings plus newer estimates, net of classifier cost"
|
||||
description="Recorded historical savings plus newer estimates, net of classifier cost. None until a turn is "
|
||||
"estimated. A turn that starts within the longest prompt-cache TTL its breakpoints use (5 minutes, or 1 hour "
|
||||
"with a 1h breakpoint) of the session's first observed turn, or of the last turn whose observation was "
|
||||
"incomplete, is not estimated, because the baseline's cache state is unknown then. Two kinds of turn are "
|
||||
"still estimated inside that window: turns that took the identical baseline path, and turns whose whole "
|
||||
"cached prefix the session had already written"
|
||||
)
|
||||
baseline_spend: float | None = Field(description="Estimated single-model cost: spend plus recorded savings")
|
||||
savings_estimated_baseline_spend: float | None = Field(
|
||||
|
|
|
|||
|
|
@ -515,10 +515,15 @@ async def test_cache_plan_count_conserves_total_and_rejects_unknown(
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("section", ["system", "tools"])
|
||||
@pytest.mark.parametrize("rejects_prefix", (False, True))
|
||||
@pytest.mark.parametrize(
|
||||
("rejected", "reason"),
|
||||
[(None, None), (1, "uncountable_system_or_tools_prefix"), (2, "token_count_unavailable")],
|
||||
ids=["counted", "system-or-tools-prefix-rejected", "message-prefix-rejected"],
|
||||
)
|
||||
async def test_native_count_preserves_settings_and_requires_every_prefix(
|
||||
section: str,
|
||||
rejects_prefix: bool,
|
||||
rejected: int | None,
|
||||
reason: str | None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
|
@ -572,9 +577,9 @@ async def test_native_count_preserves_settings_and_requires_every_prefix(
|
|||
endpoint: Final = "https://gateway.example/v1/messages/count_tokens"
|
||||
routes: Final = tuple(
|
||||
upstream.post(endpoint, json={**body, "model": _MODEL}).respond(
|
||||
400 if rejects_prefix and index == 1 else 200,
|
||||
400 if index == rejected else 200,
|
||||
json={"detail": {"error": "messages parameter is required"}}
|
||||
if rejects_prefix and index == 1
|
||||
if index == rejected
|
||||
else {"input_tokens": tokens},
|
||||
)
|
||||
for index, (body, tokens) in enumerate(
|
||||
|
|
@ -588,8 +593,9 @@ async def test_native_count_preserves_settings_and_requires_every_prefix(
|
|||
unexpected: Final = upstream.post(endpoint).respond(200, json={"input_tokens": 1})
|
||||
result: Final = await count_cache_plan(target.model, target.api_key, plan, count)
|
||||
|
||||
if rejects_prefix:
|
||||
assert result == UnsupportedCachePlan("token_count_unavailable")
|
||||
if reason is not None:
|
||||
assert isinstance(result, UnsupportedCachePlan)
|
||||
assert result.reason == reason
|
||||
else:
|
||||
assert isinstance(result, CountedPromptCachePlan)
|
||||
assert result.total_tokens == 6000
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import itertools
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
|
@ -39,6 +39,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
|||
build_window_spend_transaction,
|
||||
)
|
||||
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL
|
||||
from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
|
||||
from litellm.responses.streaming_iterator import drain_logging_tasks
|
||||
from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
|
||||
from tests.unit.proxy.db.fake_prisma_engine import engine_call
|
||||
|
|
@ -5887,3 +5888,126 @@ async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch:
|
|||
assert accumulator.drain() == {RequestErrorKey(
|
||||
date="2026-11-02", api_key=prisma.spend_log_transactions[0]["api_key"], team_id="", user_id="user", model_group="ledger-test", status_code=503,
|
||||
): 1}
|
||||
|
||||
|
||||
class _SessionTotalsRedis:
|
||||
def __init__(self) -> None:
|
||||
self.values: dict[str, float] = {} # mutable-ok: stands in for Redis state
|
||||
|
||||
async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float:
|
||||
self.values[key] = self.values.get(key, 0.0) + value
|
||||
return self.values[key]
|
||||
|
||||
async def async_get_cache(self, key: str) -> object:
|
||||
return self.values.get(key)
|
||||
|
||||
|
||||
def _session_writer() -> tuple[DBSpendUpdateWriter, SessionSpendTotals]:
|
||||
redis: Final = _SessionTotalsRedis()
|
||||
totals: Final = SessionSpendTotals(lambda: redis)
|
||||
writer: Final = DBSpendUpdateWriter(session_spend_totals=totals)
|
||||
writer._batch_database_updates = AsyncMock() # pyright: ignore[reportPrivateUsage] # the batch rollup needs a database this test does not stand up
|
||||
return writer, totals
|
||||
|
||||
|
||||
async def _record_spend_row(writer: DBSpendUpdateWriter, spend: float, litellm_params: Mapping[str, object]) -> None:
|
||||
payload: Final[dict[str, object]] = {
|
||||
**_minimal_spend_payload(),
|
||||
"request_id": "req-1",
|
||||
"call_type": "anthropic_messages",
|
||||
"api_key": "hashed-key",
|
||||
"spend": spend,
|
||||
"status": "success",
|
||||
}
|
||||
with (
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.disable_spend_logs", False
|
||||
),
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.prisma_client", _tool_usage_prisma()
|
||||
),
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
|
||||
),
|
||||
patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam
|
||||
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
|
||||
return_value=payload,
|
||||
),
|
||||
):
|
||||
await writer.update_database(
|
||||
token="hashed-key",
|
||||
user_id="test-user",
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
org_id=None,
|
||||
kwargs={"model": "claude-haiku-5-5", "litellm_params": dict(litellm_params)},
|
||||
completion_response=None,
|
||||
start_time=datetime(2026, 11, 10, 10, 0, tzinfo=timezone.utc),
|
||||
end_time=datetime(2026, 11, 10, 10, 0, 1, tzinfo=timezone.utc),
|
||||
response_cost=spend,
|
||||
)
|
||||
|
||||
|
||||
_CALLER_SESSION: Final[Mapping[str, object]] = {
|
||||
"litellm_session_id": "cc-session-1",
|
||||
"metadata": {"session_id": "cc-session-1"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params",
|
||||
[
|
||||
pytest.param(_CALLER_SESSION, id="direct-call"),
|
||||
pytest.param(
|
||||
{
|
||||
"litellm_session_id": "cc-session-1",
|
||||
"metadata": {"session_id": "cc-session-1", "trace_id": "cc-session-1"},
|
||||
},
|
||||
id="claude-code-session-header",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
**_CALLER_SESSION,
|
||||
"metadata": {"session_id": "cc-session-1", "internal_call_origin": "autorouter_classifier"},
|
||||
},
|
||||
id="classifier-sub-call",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_a_billed_row_in_a_caller_session_adds_its_spend_to_that_sessions_total(
|
||||
litellm_params: Mapping[str, object],
|
||||
):
|
||||
writer, totals = _session_writer()
|
||||
|
||||
await _record_spend_row(writer, 0.25, litellm_params)
|
||||
|
||||
assert await totals.read("hashed-key", "cc-session-1") == 0.25
|
||||
assert await totals.read("another-key", "cc-session-1") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("litellm_params", "spend"),
|
||||
[
|
||||
pytest.param(
|
||||
{
|
||||
"litellm_session_id": "minted-1",
|
||||
"metadata": {"session_id": "minted-1", "litellm_session_id_generated": True},
|
||||
},
|
||||
0.25,
|
||||
id="generated-session",
|
||||
),
|
||||
pytest.param({"metadata": {}}, 0.25, id="no-session"),
|
||||
pytest.param(_CALLER_SESSION, 0.0, id="free-row"),
|
||||
],
|
||||
)
|
||||
async def test_a_row_without_a_caller_session_or_cost_adds_nothing(
|
||||
litellm_params: Mapping[str, object], spend: float
|
||||
):
|
||||
writer, totals = _session_writer()
|
||||
|
||||
await _record_spend_row(writer, spend, litellm_params)
|
||||
|
||||
assert await totals.read("hashed-key", "cc-session-1") is None
|
||||
assert await totals.read("hashed-key", "minted-1") is None
|
||||
|
|
|
|||
|
|
@ -1098,3 +1098,55 @@ async def test_tier_pins_never_enter_the_caller_snapshot_on_later_routing_passes
|
|||
assert context is not None and context.baseline_body is not None
|
||||
assert context.baseline_body.get("max_tokens") == 16
|
||||
assert "output_config" not in context.baseline_body and "thinking" not in context.baseline_body
|
||||
|
||||
|
||||
_NON_NATIVE_MODELS: Final = _MESSAGES.validate_json("""[
|
||||
{"model_name":"test-router","litellm_params":{"model":"auto_router/complexity_router",
|
||||
"complexity_router_config":{"tiers":{"SIMPLE":"flash","MEDIUM":"flash","COMPLEX":"flash",
|
||||
"REASONING":"opus"},"session_affinity":false}}},
|
||||
{"model_name":"flash","litellm_params":{"model":"cloudflare/@cf/deepseek-ai/deepseek-v4.1-flash","api_key":"test-selected",
|
||||
"api_base":"https://chat.example.test/v1"},"model_info":{"id":"selected"}},
|
||||
{"model_name":"opus","litellm_params":{"model":"anthropic/claude-opus-5","api_key":"test-selected"},
|
||||
"model_info":{"id":"baseline"}}]""")
|
||||
|
||||
|
||||
def _chat_upstream(request: httpx.Request) -> httpx.Response:
|
||||
stream: Final = _JSON_OBJECT.validate_json(request.content).get("stream") is True
|
||||
chunk: Final = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "deepseek-v4.1-flash",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 12, "completion_tokens": 2, "total_tokens": 14},
|
||||
}
|
||||
completion: Final = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "deepseek-v4.1-flash",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 12, "completion_tokens": 2, "total_tokens": 14},
|
||||
}
|
||||
content: Final = (
|
||||
f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() if stream else json.dumps(completion).encode()
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=content,
|
||||
request=request,
|
||||
headers=MappingProxyType({"content-type": "text/event-stream" if stream else "application/json"}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True))
|
||||
async def test_a_turn_routed_to_a_non_anthropic_tier_says_its_provider_is_unsupported(
|
||||
monkeypatch: pytest.MonkeyPatch, stream: bool
|
||||
) -> None:
|
||||
rig: Final = _Rig(monkeypatch, models=_NON_NATIVE_MODELS)
|
||||
with respx.mock() as transport:
|
||||
transport.route(host="chat.example.test").mock(side_effect=_chat_upstream)
|
||||
await _call(rig.router, rig.logging(stream))
|
||||
captured: Final = _observation(await rig.capture.payload())
|
||||
assert captured.observation.outcome == "uncertain"
|
||||
assert captured.observation.reason == "unsupported_selected_provider"
|
||||
|
|
|
|||
|
|
@ -1339,6 +1339,7 @@ class TestAutoRouterSession:
|
|||
"turns": turns,
|
||||
"last_model": "anthropic/claude-sonnet-5",
|
||||
"spend": spend,
|
||||
"total_spend": None,
|
||||
"saved_spend": 0.24,
|
||||
"savings_estimated_turns": 3 if estimated else 0,
|
||||
"savings_estimated_actual_spend": 0.14 if estimated else 0.0,
|
||||
|
|
@ -1350,6 +1351,42 @@ class TestAutoRouterSession:
|
|||
"baseline_models": {"anthropic/claude-opus-5": 3},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("recorded", "expected"), [((0.5, 0.25), 0.75), ((), None)], ids=["counted", "none"])
|
||||
async def test_the_session_reports_the_spend_billed_to_its_key_under_the_session(
|
||||
self, monkeypatch: pytest.MonkeyPatch, recorded: tuple[float, ...], expected: float | None
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session
|
||||
from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
|
||||
|
||||
class _Redis:
|
||||
def __init__(self) -> None:
|
||||
self.values: dict[str, float] = {} # mutable-ok: stands in for Redis state
|
||||
|
||||
async def async_increment(
|
||||
self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False
|
||||
) -> float:
|
||||
self.values[key] = self.values.get(key, 0.0) + value
|
||||
return self.values[key]
|
||||
|
||||
async def async_get_cache(self, key: str) -> object:
|
||||
return self.values.get(key)
|
||||
|
||||
redis: Final = _Redis()
|
||||
totals: Final = SessionSpendTotals(lambda: redis)
|
||||
caller = UserAPIKeyAuth(api_key="sk-caller")
|
||||
for spend in recorded:
|
||||
await totals.record(caller.api_key or "", "sess-1", spend)
|
||||
await totals.record(UserAPIKeyAuth(api_key="sk-other").api_key or "", "sess-1", 9.0)
|
||||
monkeypatch.setattr(proxy_server.proxy_logging_obj, "session_spend_totals", totals)
|
||||
self._rig(monkeypatch, [{**self.ROW, "api_key": caller.api_key, "session_id": "sess-1"}])
|
||||
|
||||
response = await get_auto_router_session(user_api_key_dict=caller, session_id="sess-1")
|
||||
|
||||
assert response.total_spend == expected
|
||||
assert response.spend == self.ROW["spend"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_another_keys_session_is_a_404_even_for_an_admin(self, monkeypatch: pytest.MonkeyPatch):
|
||||
# The scope is the caller's own key hash, exactly what the spend writer keyed the row under;
|
||||
|
|
|
|||
90
tests/unit/proxy/spend_tracking/test_session_spend_total.py
Normal file
90
tests/unit/proxy/spend_tracking/test_session_spend_total.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
|
||||
|
||||
|
||||
class _Redis:
|
||||
"""The two RedisCache calls the session total makes, over a dict with Redis expiry on a test clock."""
|
||||
|
||||
def __init__(self, failing: bool = False) -> None:
|
||||
self.failing = failing
|
||||
self.now = 0.0
|
||||
self.values: dict[str, tuple[float, float]] = {} # mutable-ok: stands in for Redis state
|
||||
|
||||
def _live(self, key: str) -> tuple[float, float] | None:
|
||||
entry: Final = self.values.get(key)
|
||||
return entry if entry is not None and entry[1] > self.now else None
|
||||
|
||||
async def async_increment(
|
||||
self, key: str, value: float, *, ttl: int | None = None, refresh_ttl: bool = False
|
||||
) -> float:
|
||||
if self.failing:
|
||||
raise ConnectionError("Redis dropped the write")
|
||||
live: Final = self._live(key)
|
||||
expires_at: Final = self.now + (ttl or 0) if live is None or refresh_ttl else live[1]
|
||||
total: Final = (live[0] if live is not None else 0.0) + value
|
||||
self.values[key] = (total, expires_at)
|
||||
return total
|
||||
|
||||
async def async_get_cache(self, key: str) -> str | None:
|
||||
if self.failing:
|
||||
raise ConnectionError("Redis dropped the read")
|
||||
live: Final = self._live(key)
|
||||
return str(live[0]) if live is not None else None
|
||||
|
||||
|
||||
def _totals(redis: _Redis) -> SessionSpendTotals:
|
||||
return SessionSpendTotals(lambda: redis)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_session_total_sums_its_own_keys_rows_only() -> None:
|
||||
totals: Final = _totals(_Redis())
|
||||
|
||||
await totals.record("key-a", "session-1", 0.5)
|
||||
await totals.record("key-a", "session-1", 0.25)
|
||||
await totals.record("key-b", "session-1", 4.0)
|
||||
await totals.record("key-a", "session-2", 8.0)
|
||||
|
||||
assert await totals.read("key-a", "session-1") == 0.75
|
||||
assert await totals.read("key-b", "session-1") == 4.0
|
||||
assert await totals.read("key-a", "session-3") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_redis_outage_never_fails_the_spend_write_and_reads_as_no_total() -> None:
|
||||
totals: Final = _totals(_Redis(failing=True))
|
||||
|
||||
await totals.record("key", "session", 0.5)
|
||||
|
||||
assert await totals.read("key", "session") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_a_shared_redis_there_is_no_total() -> None:
|
||||
totals: Final = SessionSpendTotals(lambda: None)
|
||||
|
||||
await totals.record("key", "session", 0.5)
|
||||
|
||||
assert await totals.read("key", "session") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_active_session_keeps_its_total_and_an_idle_one_restarts() -> None:
|
||||
redis: Final = _Redis()
|
||||
totals: Final = _totals(redis)
|
||||
|
||||
await totals.record("key", "session", 0.5)
|
||||
redis.now = 3000.0
|
||||
await totals.record("key", "session", 0.25)
|
||||
redis.now = 6000.0
|
||||
active: Final = await totals.read("key", "session")
|
||||
redis.now = 6000.0 + 3601.0
|
||||
idle: Final = await totals.read("key", "session")
|
||||
await totals.record("key", "session", 0.125)
|
||||
|
||||
assert active == 0.75
|
||||
assert idle is None
|
||||
assert await totals.read("key", "session") == 0.125
|
||||
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27922,7 +27922,7 @@ export interface components {
|
|||
router_type: string;
|
||||
/**
|
||||
* Saved Spend
|
||||
* @description Recorded historical savings plus newer estimates, net of classifier cost
|
||||
* @description Recorded historical savings plus newer estimates, net of classifier cost. None until a turn is estimated. A turn that starts within the longest prompt-cache TTL its breakpoints use (5 minutes, or 1 hour with a 1h breakpoint) of the session's first observed turn, or of the last turn whose observation was incomplete, is not estimated, because the baseline's cache state is unknown then. Two kinds of turn are still estimated inside that window: turns that took the identical baseline path, and turns whose whole cached prefix the session had already written
|
||||
*/
|
||||
saved_spend: number | null;
|
||||
/**
|
||||
|
|
@ -27947,6 +27947,11 @@ export interface components {
|
|||
* @description What the session's routed traffic actually cost, classifier calls included
|
||||
*/
|
||||
spend: number;
|
||||
/**
|
||||
* Total Spend
|
||||
* @description Spend billed to this key under the session id, routed or not, from a shared Redis counter whose one-hour expiry every billed request resets. None without a shared Redis or once the counter expired. A best-effort figure, not a lifetime total: a session idle past the expiry restarts from zero, and a failed Redis write, eviction or flush drops spend without notice
|
||||
*/
|
||||
total_spend?: number | null;
|
||||
/**
|
||||
* Turns
|
||||
* @description Auto-routed turns the rollup has recorded for this session so far
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue