mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
Merge b6fff52111 into eb0e3f8c18
This commit is contained in:
todayim
GitHub
commit
5bf3c1b0a7
5 changed files with 761 additions and 0 deletions
|
|
@ -675,6 +675,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/jwt/key/mapping/delete",
|
||||
"/jwt/key/mapping/list",
|
||||
"/jwt/key/mapping/info",
|
||||
"/spend/usage",
|
||||
]
|
||||
+ key_management_routes
|
||||
+ mcp_management_routes
|
||||
|
|
|
|||
|
|
@ -610,6 +610,9 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
|||
router as spend_management_router,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.proxy.spend_tracking.usage_ingestion_endpoints import (
|
||||
router as usage_ingestion_router,
|
||||
)
|
||||
from litellm.proxy.types_utils.utils import get_instance_fn
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
router as ui_crud_endpoints_router,
|
||||
|
|
@ -17681,6 +17684,7 @@ app.include_router(organization_router)
|
|||
app.include_router(customer_router)
|
||||
app.include_router(management_v1_router)
|
||||
app.include_router(spend_management_router)
|
||||
app.include_router(usage_ingestion_router)
|
||||
app.include_router(caching_router)
|
||||
app.include_router(analytics_router)
|
||||
app.include_router(callback_management_endpoints_router)
|
||||
|
|
|
|||
351
litellm/proxy/spend_tracking/usage_ingestion_endpoints.py
Normal file
351
litellm/proxy/spend_tracking/usage_ingestion_endpoints.py
Normal file
|
|
@ -0,0 +1,351 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Final,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
TypeAlias,
|
||||
cast, # noqa: TID251 # untyped tx boundary needs cast for the shim
|
||||
)
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, SpendLogsPayload, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
MAX_RECORDS_PER_REQUEST: Final = 1000
|
||||
|
||||
|
||||
class ExternalUsageRecord(BaseModel):
|
||||
api_key: str | None = Field(
|
||||
default=None, min_length=1, description="Raw virtual key (sk-...) to attribute usage to. Never logged."
|
||||
)
|
||||
api_key_hash: str | None = Field(
|
||||
default=None,
|
||||
min_length=1,
|
||||
description="SHA-256 hash of the virtual key. Use instead of api_key to avoid submitting raw keys.",
|
||||
)
|
||||
model: str = Field(min_length=1)
|
||||
prompt_tokens: int = Field(ge=0)
|
||||
completion_tokens: int = Field(ge=0)
|
||||
start_time: datetime
|
||||
end_time: datetime | None = None
|
||||
cost: float | None = Field(
|
||||
default=None, ge=0, description="Explicit cost in USD. Computed from litellm pricing when omitted."
|
||||
)
|
||||
idempotency_key: str | None = Field(
|
||||
default=None, max_length=255, description="Becomes the spend-log request_id for dedup on retries."
|
||||
)
|
||||
tags: list[str] | None = None # mutable-ok: serialized as a JSON array by spend logs
|
||||
end_user_id: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def end_time_not_before_start_time(self) -> "ExternalUsageRecord":
|
||||
if self.end_time is not None and self.end_time < self.start_time:
|
||||
raise ValueError("end_time must not be before start_time")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def exactly_one_key_identifier(self) -> "ExternalUsageRecord":
|
||||
if (self.api_key is None) == (self.api_key_hash is None):
|
||||
raise ValueError("exactly one of api_key or api_key_hash is required")
|
||||
return self
|
||||
|
||||
|
||||
class UsageIngestRequest(BaseModel):
|
||||
records: tuple[ExternalUsageRecord, ...] = Field(min_length=1, max_length=MAX_RECORDS_PER_REQUEST)
|
||||
|
||||
|
||||
class UsageIngestRecordResult(BaseModel):
|
||||
request_id: str
|
||||
status: Literal["recorded", "duplicate", "error"]
|
||||
spend: float | None = None
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class UsageIngestResponse(BaseModel):
|
||||
results: tuple[UsageIngestRecordResult, ...]
|
||||
|
||||
|
||||
class KeyAttribution(NamedTuple):
|
||||
user_id: str | None
|
||||
team_id: str | None
|
||||
organization_id: str | None
|
||||
|
||||
|
||||
ReservationOutcome: TypeAlias = Literal["reserved", "duplicate", "disabled"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageIngestionDeps:
|
||||
lookup_key: Callable[[str], Awaitable[KeyAttribution | None]]
|
||||
reserve_spend_log: Callable[[ExternalUsageRecord, str, str, KeyAttribution, float], Awaitable[ReservationOutcome]]
|
||||
record_spend: Callable[..., Awaitable[None]]
|
||||
compute_cost: Callable[[litellm.ModelResponse, str], float]
|
||||
generate_request_id: Callable[[], str]
|
||||
|
||||
|
||||
def _build_usage_kwargs(record: ExternalUsageRecord, hashed_token: str) -> Mapping[str, object]:
|
||||
tags: Final = list(record.tags) if record.tags else [] # mutable-ok: real list required by json serializer
|
||||
metadata: Final = MappingProxyType(
|
||||
{
|
||||
"user_api_key": hashed_token,
|
||||
"user_api_key_end_user_id": record.end_user_id,
|
||||
"tags": tags,
|
||||
}
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
"model": record.model,
|
||||
"call_type": "ingest_external_usage",
|
||||
"litellm_params": MappingProxyType({"model": record.model, "metadata": metadata}),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _build_completion_response(record: ExternalUsageRecord, request_id: str) -> litellm.ModelResponse:
|
||||
total_tokens: Final = record.prompt_tokens + record.completion_tokens
|
||||
usage: Final = litellm.Usage(
|
||||
prompt_tokens=record.prompt_tokens,
|
||||
completion_tokens=record.completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return litellm.ModelResponse(
|
||||
id=request_id,
|
||||
model=record.model,
|
||||
created=int(record.start_time.timestamp()),
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
|
||||
def build_spend_log_payload(
|
||||
record: ExternalUsageRecord,
|
||||
request_id: str,
|
||||
hashed_token: str,
|
||||
key: KeyAttribution,
|
||||
cost: float,
|
||||
) -> SpendLogsPayload:
|
||||
payload: Final[SpendLogsPayload] = get_logging_payload(
|
||||
kwargs=_build_usage_kwargs(record, hashed_token),
|
||||
response_obj=_build_completion_response(record, request_id),
|
||||
start_time=record.start_time,
|
||||
end_time=record.end_time or record.start_time,
|
||||
)
|
||||
payload["spend"] = cost
|
||||
if isinstance(payload["startTime"], datetime):
|
||||
payload["startTime"] = payload["startTime"].isoformat()
|
||||
if isinstance(payload["endTime"], datetime):
|
||||
payload["endTime"] = payload["endTime"].isoformat()
|
||||
if key.organization_id is not None and key.organization_id != "":
|
||||
payload["organization_id"] = key.organization_id
|
||||
if key.team_id is not None and key.team_id != "":
|
||||
payload["team_id"] = key.team_id
|
||||
return payload
|
||||
|
||||
|
||||
def _resolve_cost(deps: UsageIngestionDeps, record: ExternalUsageRecord, response: litellm.ModelResponse) -> float:
|
||||
if record.cost is not None:
|
||||
return record.cost
|
||||
return deps.compute_cost(response, record.model)
|
||||
|
||||
|
||||
async def process_external_usage_record(
|
||||
record: ExternalUsageRecord, deps: UsageIngestionDeps
|
||||
) -> UsageIngestRecordResult:
|
||||
request_id: Final = record.idempotency_key or deps.generate_request_id()
|
||||
hashed_token: Final = record.api_key_hash if record.api_key_hash is not None else hash_token(record.api_key or "")
|
||||
|
||||
key: Final = await deps.lookup_key(hashed_token)
|
||||
if key is None:
|
||||
return UsageIngestRecordResult(request_id=request_id, status="error", error="api key not found")
|
||||
|
||||
response: Final = _build_completion_response(record, request_id)
|
||||
|
||||
try:
|
||||
cost: Final = _resolve_cost(deps, record, response)
|
||||
except Exception as e: # noqa: BLE001 # pricing lookup raises arbitrary provider-specific errors; any failure means the model is unpriceable and the record must carry an explicit cost
|
||||
verbose_proxy_logger.info("ingest usage: cost computation failed for model %s: %s", record.model, e)
|
||||
return UsageIngestRecordResult(
|
||||
request_id=request_id,
|
||||
status="error",
|
||||
error="could not compute cost for this model, pass an explicit cost",
|
||||
)
|
||||
|
||||
if record.idempotency_key is not None:
|
||||
try:
|
||||
reservation: Final = await deps.reserve_spend_log(record, request_id, hashed_token, key, cost)
|
||||
except Exception as e: # noqa: BLE001 # booking raises arbitrary persistence errors; an aborted transaction means nothing was booked, so telling the caller to retry is safe
|
||||
verbose_proxy_logger.info("ingest usage: transactional booking failed for %s: %s", request_id, e)
|
||||
return UsageIngestRecordResult(
|
||||
request_id=request_id,
|
||||
status="error",
|
||||
error="booking failed transactionally, nothing was recorded, safe to retry",
|
||||
)
|
||||
if reservation == "duplicate":
|
||||
return UsageIngestRecordResult(request_id=request_id, status="duplicate")
|
||||
if reservation == "disabled":
|
||||
return UsageIngestRecordResult(
|
||||
request_id=request_id,
|
||||
status="error",
|
||||
error="spend updates are disabled on this proxy, nothing was recorded",
|
||||
)
|
||||
return UsageIngestRecordResult(request_id=request_id, status="recorded", spend=cost)
|
||||
|
||||
await deps.record_spend(
|
||||
token=hashed_token,
|
||||
user_id=key.user_id,
|
||||
end_user_id=record.end_user_id,
|
||||
team_id=key.team_id,
|
||||
kwargs=_build_usage_kwargs(record, hashed_token),
|
||||
completion_response=response,
|
||||
start_time=record.start_time,
|
||||
end_time=record.end_time or record.start_time,
|
||||
response_cost=cost,
|
||||
org_id=key.organization_id,
|
||||
)
|
||||
return UsageIngestRecordResult(request_id=request_id, status="recorded", spend=cost)
|
||||
|
||||
|
||||
def _attribution_of(key_row: object) -> KeyAttribution:
|
||||
return KeyAttribution(
|
||||
user_id=getattr(key_row, "user_id", None),
|
||||
team_id=getattr(key_row, "team_id", None),
|
||||
organization_id=getattr(key_row, "organization_id", None),
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma.client import TransactionManager
|
||||
|
||||
|
||||
class _TransactionClientShim:
|
||||
def __init__(self, tx: "TransactionManager") -> None:
|
||||
self.db: Final = tx
|
||||
|
||||
|
||||
async def reserve_spend_log_atomic(
|
||||
record: ExternalUsageRecord,
|
||||
request_id: str,
|
||||
hashed_token: str,
|
||||
key: KeyAttribution,
|
||||
cost: float,
|
||||
) -> ReservationOutcome:
|
||||
from litellm.proxy.proxy_server import litellm_proxy_budget_name, prisma_client, proxy_logging_obj
|
||||
from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
|
||||
if ProxyUpdateSpend.disable_spend_updates() is True:
|
||||
return "disabled"
|
||||
|
||||
spend_payload: Final = build_spend_log_payload(record, request_id, hashed_token, key, cost)
|
||||
payload: Final = prisma_client.jsonify_object(spend_payload)
|
||||
request_tags: Final = spend_payload.get("request_tags")
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
writer: Final = proxy_logging_obj.db_spend_update_writer
|
||||
|
||||
try:
|
||||
async with prisma_client.tx() as tx:
|
||||
shim: Final = cast(PrismaClient, _TransactionClientShim(tx)) # cast-ok: helper uses only .db (untyped)
|
||||
await SpendLogsRepository(shim).table.create(data=payload)
|
||||
await writer._update_key_db(response_cost=cost, hashed_token=hashed_token, prisma_client=shim)
|
||||
await writer._update_user_db(
|
||||
response_cost=cost,
|
||||
user_id=key.user_id,
|
||||
prisma_client=shim,
|
||||
litellm_proxy_budget_name=litellm_proxy_budget_name,
|
||||
end_user_id=record.end_user_id,
|
||||
)
|
||||
await writer._update_team_db(
|
||||
response_cost=cost, team_id=key.team_id, user_id=key.user_id, prisma_client=shim
|
||||
)
|
||||
await writer._update_org_db(response_cost=cost, org_id=key.organization_id, prisma_client=shim)
|
||||
await writer._update_tag_db(response_cost=cost, request_tags=request_tags, prisma_client=shim)
|
||||
except UniqueViolationError:
|
||||
return "duplicate"
|
||||
return "reserved"
|
||||
|
||||
|
||||
def default_ingestion_deps() -> UsageIngestionDeps:
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
raise ProxyException(
|
||||
message="Prisma Client is not initialized",
|
||||
type="internal_error",
|
||||
param="None",
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
async def lookup_key(hashed_token: str) -> KeyAttribution | None:
|
||||
from litellm.proxy.auth.auth_checks import get_key_object
|
||||
|
||||
try:
|
||||
key_row: Final = await get_key_object(
|
||||
hashed_token=hashed_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
except ProxyException:
|
||||
return None
|
||||
return _attribution_of(key_row)
|
||||
|
||||
return UsageIngestionDeps(
|
||||
lookup_key=lookup_key,
|
||||
reserve_spend_log=reserve_spend_log_atomic,
|
||||
record_spend=proxy_logging_obj.db_spend_update_writer.update_database,
|
||||
compute_cost=lambda resp, model: litellm.completion_cost(completion_response=resp, model=model),
|
||||
generate_request_id=lambda: str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/spend/usage",
|
||||
tags=["Budget & Spend Tracking"], # mutable-ok: fastapi decorator contract takes a list
|
||||
dependencies=[Depends(user_api_key_auth)], # mutable-ok: fastapi decorator contract takes a list
|
||||
response_model=UsageIngestResponse,
|
||||
)
|
||||
async def ingest_external_usage(
|
||||
request: UsageIngestRequest,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> UsageIngestResponse:
|
||||
"""
|
||||
PROXY_ADMIN ONLY: record externally measured usage into the same spend pipeline as proxy-routed traffic.
|
||||
|
||||
For inference traffic that legitimately bypasses the proxy (for example async batch processors
|
||||
dispatching directly to model gateways), so budgets and spend stay coherent in litellm as the
|
||||
single metering system.
|
||||
|
||||
Attribution (user/team/org) is derived from the given virtual key, submitted either raw
|
||||
(api_key) or pre-hashed (api_key_hash) to keep raw keys out of request bodies. Records accept an optional
|
||||
idempotency_key, stored as the spend-log request_id: the reservation insert, counter updates and
|
||||
dedup are checked atomically at the database primary key, so overlapping retries are safe. The
|
||||
reservation row is always written (even when disable_spend_logs is set), because it is both the
|
||||
dedup anchor and the audit record for the booked usage. When cost is omitted it is computed from
|
||||
litellm pricing; records whose model cannot be priced are rejected with an error instead of
|
||||
being booked as zero spend.
|
||||
"""
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Only proxy admins ingest spend records here. Use a key with the proxy_admin role.",
|
||||
)
|
||||
|
||||
deps: Final = default_ingestion_deps()
|
||||
results: Final = tuple(
|
||||
await asyncio.gather(*(process_external_usage_record(record, deps) for record in request.records))
|
||||
)
|
||||
return UsageIngestResponse(results=results)
|
||||
|
|
@ -0,0 +1,275 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.spend_tracking.usage_ingestion_endpoints import (
|
||||
ExternalUsageRecord,
|
||||
KeyAttribution,
|
||||
UsageIngestionDeps,
|
||||
build_spend_log_payload,
|
||||
process_external_usage_record,
|
||||
)
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
RAW_KEY = "sk-test-batch-dispatch-key"
|
||||
GENERATED_ID = "generated-uuid-1"
|
||||
DEFAULT_KEY = KeyAttribution(user_id="u-1", team_id="t-1", organization_id="o-1")
|
||||
|
||||
|
||||
class RecordingDeps:
|
||||
def __init__(
|
||||
self,
|
||||
key: KeyAttribution | None = DEFAULT_KEY,
|
||||
existing_ids: frozenset[str] = frozenset(),
|
||||
compute_cost_result: float = 0.05,
|
||||
compute_cost_error: Exception | None = None,
|
||||
reserve_outcome: str = "reserved",
|
||||
reserve_raises: Exception | None = None,
|
||||
):
|
||||
self._key = key
|
||||
self._existing_ids = existing_ids
|
||||
self._compute_cost_result = compute_cost_result
|
||||
self._compute_cost_error = compute_cost_error
|
||||
self._reserve_outcome = reserve_outcome
|
||||
self._reserve_raises = reserve_raises
|
||||
self.spend_calls: list[dict[str, Any]] = []
|
||||
self.reserve_calls: list[dict[str, Any]] = []
|
||||
self.compute_cost_calls: list[tuple[Any, str]] = []
|
||||
|
||||
def as_deps(self) -> UsageIngestionDeps:
|
||||
async def lookup_key(hashed: str) -> KeyAttribution | None:
|
||||
self.looked_up_hashed = hashed
|
||||
return self._key
|
||||
|
||||
async def reserve_spend_log(
|
||||
record: ExternalUsageRecord,
|
||||
request_id: str,
|
||||
hashed_token: str,
|
||||
key: KeyAttribution,
|
||||
cost: float,
|
||||
) -> Any:
|
||||
self.reserve_calls.append(
|
||||
{
|
||||
"record": record,
|
||||
"request_id": request_id,
|
||||
"hashed_token": hashed_token,
|
||||
"key": key,
|
||||
"cost": cost,
|
||||
}
|
||||
)
|
||||
if self._reserve_raises is not None:
|
||||
raise self._reserve_raises
|
||||
if self._reserve_outcome == "disabled":
|
||||
return "disabled"
|
||||
return "duplicate" if request_id in self._existing_ids else "reserved"
|
||||
|
||||
async def record_spend(**kwargs: Any) -> None:
|
||||
self.spend_calls.append(kwargs)
|
||||
|
||||
def compute_cost(response: Any, model: str) -> float:
|
||||
self.compute_cost_calls.append((response, model))
|
||||
if self._compute_cost_error is not None:
|
||||
raise self._compute_cost_error
|
||||
return self._compute_cost_result
|
||||
|
||||
return UsageIngestionDeps(
|
||||
lookup_key=lookup_key,
|
||||
reserve_spend_log=reserve_spend_log,
|
||||
record_spend=record_spend,
|
||||
compute_cost=compute_cost,
|
||||
generate_request_id=lambda: GENERATED_ID,
|
||||
)
|
||||
|
||||
|
||||
def make_record(**overrides: Any) -> ExternalUsageRecord:
|
||||
base: dict[str, Any] = {
|
||||
"api_key": RAW_KEY,
|
||||
"model": "gpt-4o-mini",
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"start_time": datetime(2026, 8, 5, 12, 0, 0, tzinfo=timezone.utc),
|
||||
}
|
||||
base.update(overrides)
|
||||
return ExternalUsageRecord(**base)
|
||||
|
||||
|
||||
def run(coro: Any) -> Any:
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
def test_spend_log_payload_matches_funnel_shape():
|
||||
record = make_record(idempotency_key="batch-1-line-1", tags=["batch:job-42"], end_user_id="tenant-a")
|
||||
payload = build_spend_log_payload(
|
||||
record=record,
|
||||
request_id="batch-1-line-1",
|
||||
hashed_token=hash_token(RAW_KEY),
|
||||
key=DEFAULT_KEY,
|
||||
cost=0.123,
|
||||
)
|
||||
assert payload["request_id"] == "batch-1-line-1"
|
||||
assert payload["spend"] == 0.123
|
||||
assert payload["total_tokens"] == 150
|
||||
assert payload["prompt_tokens"] == 100
|
||||
assert payload["completion_tokens"] == 50
|
||||
assert payload["api_key"] == hash_token(RAW_KEY)
|
||||
assert payload["api_key"] != RAW_KEY
|
||||
assert payload["team_id"] == "t-1"
|
||||
assert payload["organization_id"] == "o-1"
|
||||
assert payload["end_user"] == "tenant-a"
|
||||
|
||||
metadata = json.loads(payload["metadata"]) if isinstance(payload["metadata"], str) else payload["metadata"]
|
||||
assert metadata["user_api_key"] == hash_token(RAW_KEY)
|
||||
|
||||
request_tags = (
|
||||
json.loads(payload["request_tags"]) if isinstance(payload["request_tags"], str) else payload["request_tags"]
|
||||
)
|
||||
assert request_tags == ["batch:job-42"]
|
||||
|
||||
|
||||
def test_idempotent_record_books_through_atomic_reserve_not_funnel():
|
||||
deps = RecordingDeps(compute_cost_error=RuntimeError("pricing must not be consulted"))
|
||||
result = run(
|
||||
process_external_usage_record(make_record(cost=0.123, idempotency_key="batch-1-line-1"), deps.as_deps())
|
||||
)
|
||||
assert result.status == "recorded"
|
||||
assert result.spend == 0.123
|
||||
assert result.request_id == "batch-1-line-1"
|
||||
|
||||
assert len(deps.reserve_calls) == 1
|
||||
reserve_call = deps.reserve_calls[0]
|
||||
assert reserve_call["request_id"] == "batch-1-line-1"
|
||||
assert reserve_call["hashed_token"] == hash_token(RAW_KEY)
|
||||
assert reserve_call["cost"] == 0.123
|
||||
assert reserve_call["key"] == DEFAULT_KEY
|
||||
|
||||
assert deps.spend_calls == []
|
||||
assert deps.compute_cost_calls == []
|
||||
|
||||
|
||||
def test_computed_cost_is_resolved_before_reserving():
|
||||
deps = RecordingDeps(compute_cost_result=0.07)
|
||||
result = run(process_external_usage_record(make_record(idempotency_key="k-2"), deps.as_deps()))
|
||||
assert result.status == "recorded"
|
||||
assert result.spend == 0.07
|
||||
assert len(deps.compute_cost_calls) == 1
|
||||
assert deps.compute_cost_calls[0][1] == "gpt-4o-mini"
|
||||
assert deps.reserve_calls[0]["cost"] == 0.07
|
||||
|
||||
|
||||
def test_unpriceable_model_without_explicit_cost_is_error_and_books_nothing():
|
||||
deps = RecordingDeps(compute_cost_error=ValueError("unknown model"))
|
||||
result = run(process_external_usage_record(make_record(idempotency_key="k-3"), deps.as_deps()))
|
||||
assert result.status == "error"
|
||||
assert result.spend is None
|
||||
assert "explicit cost" in (result.error or "")
|
||||
assert deps.reserve_calls == []
|
||||
assert deps.spend_calls == []
|
||||
|
||||
|
||||
def test_unknown_key_is_rejected_and_books_nothing():
|
||||
deps = RecordingDeps(key=None)
|
||||
result = run(process_external_usage_record(make_record(idempotency_key="k-4"), deps.as_deps()))
|
||||
assert result.status == "error"
|
||||
assert result.error == "api key not found"
|
||||
assert deps.reserve_calls == []
|
||||
assert deps.spend_calls == []
|
||||
|
||||
|
||||
def test_duplicate_idempotency_key_is_skipped_and_never_rebooks():
|
||||
deps = RecordingDeps(existing_ids=frozenset({"k-5"}))
|
||||
result = run(process_external_usage_record(make_record(idempotency_key="k-5"), deps.as_deps()))
|
||||
assert result.status == "duplicate"
|
||||
assert result.request_id == "k-5"
|
||||
assert len(deps.reserve_calls) == 1
|
||||
assert deps.spend_calls == []
|
||||
|
||||
|
||||
def test_missing_idempotency_key_uses_funnel_with_generated_request_id():
|
||||
deps = RecordingDeps()
|
||||
result = run(process_external_usage_record(make_record(cost=0.01), deps.as_deps()))
|
||||
assert result.status == "recorded"
|
||||
assert result.request_id == GENERATED_ID
|
||||
assert deps.reserve_calls == []
|
||||
assert len(deps.spend_calls) == 1
|
||||
|
||||
call = deps.spend_calls[0]
|
||||
assert call["token"] == hash_token(RAW_KEY)
|
||||
assert call["user_id"] == "u-1"
|
||||
assert call["team_id"] == "t-1"
|
||||
assert call["org_id"] == "o-1"
|
||||
assert call["response_cost"] == 0.01
|
||||
response = call["completion_response"]
|
||||
assert response.id == GENERATED_ID
|
||||
assert response.usage.total_tokens == 150
|
||||
metadata = call["kwargs"]["litellm_params"]["metadata"]
|
||||
assert metadata["user_api_key"] == hash_token(RAW_KEY)
|
||||
assert metadata["user_api_key"] != RAW_KEY
|
||||
|
||||
|
||||
def test_end_time_defaults_to_start_time_and_end_before_start_rejected():
|
||||
deps = RecordingDeps()
|
||||
start = datetime(2026, 8, 5, 12, 0, 0, tzinfo=timezone.utc)
|
||||
run(process_external_usage_record(make_record(cost=0.01, start_time=start), deps.as_deps()))
|
||||
assert deps.spend_calls[0]["start_time"] == start
|
||||
assert deps.spend_calls[0]["end_time"] == start
|
||||
|
||||
earlier = datetime(2026, 8, 5, 11, 0, 0, tzinfo=timezone.utc)
|
||||
with pytest.raises(ValidationError):
|
||||
make_record(start_time=start, end_time=earlier)
|
||||
|
||||
|
||||
def test_record_without_idempotency_key_still_flows_tags_to_funnel_kwargs():
|
||||
deps = RecordingDeps()
|
||||
result = run(
|
||||
process_external_usage_record(
|
||||
make_record(cost=0.01, tags=["batch:job-42"], end_user_id="tenant-a"), deps.as_deps()
|
||||
)
|
||||
)
|
||||
assert result.status == "recorded"
|
||||
metadata = deps.spend_calls[0]["kwargs"]["litellm_params"]["metadata"]
|
||||
assert metadata["tags"] == ["batch:job-42"]
|
||||
assert metadata["user_api_key_end_user_id"] == "tenant-a"
|
||||
assert deps.spend_calls[0]["end_user_id"] == "tenant-a"
|
||||
|
||||
|
||||
def test_disabled_spend_updates_reports_error_instead_of_fake_recorded():
|
||||
deps = RecordingDeps(reserve_outcome="disabled")
|
||||
result = run(process_external_usage_record(make_record(cost=0.01, idempotency_key="k-9"), deps.as_deps()))
|
||||
assert result.status == "error"
|
||||
assert "disabled" in (result.error or "")
|
||||
assert result.spend is None
|
||||
assert deps.spend_calls == []
|
||||
|
||||
|
||||
def test_failed_booking_is_retry_safe_error_not_permanent_duplicate():
|
||||
deps = RecordingDeps(reserve_raises=RuntimeError("db gone mid-tx"))
|
||||
result = run(process_external_usage_record(make_record(cost=0.01, idempotency_key="k-10"), deps.as_deps()))
|
||||
assert result.status == "error"
|
||||
assert "safe to retry" in (result.error or "")
|
||||
assert result.spend is None
|
||||
assert deps.spend_calls == []
|
||||
|
||||
|
||||
def test_key_hash_resolves_without_raw_key_in_body():
|
||||
deps = RecordingDeps()
|
||||
record = make_record(cost=0.01, idempotency_key="k-11")
|
||||
record = ExternalUsageRecord(**{**record.model_dump(), "api_key": None, "api_key_hash": hash_token(RAW_KEY)})
|
||||
result = run(process_external_usage_record(record, deps.as_deps()))
|
||||
assert result.status == "recorded"
|
||||
assert deps.looked_up_hashed == hash_token(RAW_KEY)
|
||||
assert deps.reserve_calls[0]["hashed_token"] == hash_token(RAW_KEY)
|
||||
|
||||
|
||||
def test_exactly_one_key_identifier_required():
|
||||
with pytest.raises(ValidationError):
|
||||
make_record(api_key=None)
|
||||
with pytest.raises(ValidationError):
|
||||
make_record(api_key_hash=hash_token(RAW_KEY))
|
||||
130
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
130
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -13812,6 +13812,39 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/spend/usage": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
get?: never;
|
||||
put?: never;
|
||||
/**
|
||||
* Ingest External Usage
|
||||
* @description PROXY_ADMIN ONLY: record externally measured usage into the same spend pipeline as proxy-routed traffic.
|
||||
*
|
||||
* For inference traffic that legitimately bypasses the proxy (for example async batch processors
|
||||
* dispatching directly to model gateways), so budgets and spend stay coherent in litellm as the
|
||||
* single metering system.
|
||||
*
|
||||
* Attribution (user/team/org) is derived from the given virtual key, submitted either raw
|
||||
* (api_key) or pre-hashed (api_key_hash) to keep raw keys out of request bodies. Records accept an optional
|
||||
* idempotency_key, stored as the spend-log request_id: the reservation insert, counter updates and
|
||||
* dedup are checked atomically at the database primary key, so overlapping retries are safe. The
|
||||
* reservation row is always written (even when disable_spend_logs is set), because it is both the
|
||||
* dedup anchor and the audit record for the booked usage. When cost is omitted it is computed from
|
||||
* litellm pricing; records whose model cannot be priced are rejected with an error instead of
|
||||
* being booked as zero spend.
|
||||
*/
|
||||
post: operations["ingest_external_usage_spend_usage_post"];
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/spend/users": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -26272,6 +26305,46 @@ export interface components {
|
|||
/** Updated At */
|
||||
updated_at?: number | null;
|
||||
};
|
||||
/** ExternalUsageRecord */
|
||||
ExternalUsageRecord: {
|
||||
/**
|
||||
* Api Key
|
||||
* @description Raw virtual key (sk-...) to attribute usage to. Never logged.
|
||||
*/
|
||||
api_key?: string | null;
|
||||
/**
|
||||
* Api Key Hash
|
||||
* @description SHA-256 hash of the virtual key. Use instead of api_key to avoid submitting raw keys.
|
||||
*/
|
||||
api_key_hash?: string | null;
|
||||
/** Completion Tokens */
|
||||
completion_tokens: number;
|
||||
/**
|
||||
* Cost
|
||||
* @description Explicit cost in USD. Computed from litellm pricing when omitted.
|
||||
*/
|
||||
cost?: number | null;
|
||||
/** End Time */
|
||||
end_time?: string | null;
|
||||
/** End User Id */
|
||||
end_user_id?: string | null;
|
||||
/**
|
||||
* Idempotency Key
|
||||
* @description Becomes the spend-log request_id for dedup on retries.
|
||||
*/
|
||||
idempotency_key?: string | null;
|
||||
/** Model */
|
||||
model: string;
|
||||
/** Prompt Tokens */
|
||||
prompt_tokens: number;
|
||||
/**
|
||||
* Start Time
|
||||
* Format: date-time
|
||||
*/
|
||||
start_time: string;
|
||||
/** Tags */
|
||||
tags?: string[] | null;
|
||||
};
|
||||
/**
|
||||
* FacetListResponse
|
||||
* @description The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows.
|
||||
|
|
@ -37060,6 +37133,30 @@ export interface components {
|
|||
/** Usage Units Daily */
|
||||
usage_units_daily: components["schemas"]["UsageUnitsDailyPoint"][];
|
||||
};
|
||||
/** UsageIngestRecordResult */
|
||||
UsageIngestRecordResult: {
|
||||
/** Error */
|
||||
error?: string | null;
|
||||
/** Request Id */
|
||||
request_id: string;
|
||||
/** Spend */
|
||||
spend?: number | null;
|
||||
/**
|
||||
* Status
|
||||
* @enum {string}
|
||||
*/
|
||||
status: "recorded" | "duplicate" | "error";
|
||||
};
|
||||
/** UsageIngestRequest */
|
||||
UsageIngestRequest: {
|
||||
/** Records */
|
||||
records: components["schemas"]["ExternalUsageRecord"][];
|
||||
};
|
||||
/** UsageIngestResponse */
|
||||
UsageIngestResponse: {
|
||||
/** Results */
|
||||
results: components["schemas"]["UsageIngestRecordResult"][];
|
||||
};
|
||||
/** UsageLogEntry */
|
||||
UsageLogEntry: {
|
||||
/** Action */
|
||||
|
|
@ -55252,6 +55349,39 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
ingest_external_usage_spend_usage_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody: {
|
||||
content: {
|
||||
"application/json": components["schemas"]["UsageIngestRequest"];
|
||||
};
|
||||
};
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["UsageIngestResponse"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
spend_user_fn_spend_users_get: {
|
||||
parameters: {
|
||||
query?: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue