Merge remote-tracking branch 'origin/main' into litellm_db_general_settings_stale_yaml_keys

This commit is contained in:
mateo-berri 2026-10-28 16:57:50 -07:00
commit 54b90ef6a5
157 changed files with 12035 additions and 763 deletions

View file

@ -1,6 +1,8 @@
---
name: database-access
description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
metadata:
internal: true
---
# Database access

View file

@ -4,10 +4,6 @@ on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the P
pull_request_target:
branches:
- main
pull_request_review:
types:
- submitted
- dismissed
permissions:
contents: read

View file

@ -137,6 +137,7 @@ jobs:
include:
- shard: core-utils
test-path: |-
tests/unit/accounting
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction

View file

@ -2,9 +2,12 @@
needs an approving review from the spend-log owner before it can merge.
CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies,
so a path rule would put every migration on the owner. This gate reads the diff instead: when no
added line in the watched files mentions the table it passes at once, otherwise it passes only once
the owner's latest review on the pull request is an approval. The workflow re-runs on every review.
so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma
hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration
hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes
only once the owner's latest review on the pull request is an approval that still covers every spend-log line
in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw).
Reviews do not trigger the workflow, so the check is re-run by hand after approving.
"""
from __future__ import annotations
@ -13,8 +16,10 @@ import argparse
import re
import subprocess
import sys
from collections.abc import Sequence
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -25,7 +30,10 @@ WATCHED: Final = (
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
)
MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/"
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs", re.IGNORECASE)
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs(?![A-Za-z0-9])", re.IGNORECASE)
DECIDING_STATES: Final = frozenset({"APPROVED", "CHANGES_REQUESTED", "DISMISSED"})
MODEL_START: Final = re.compile(r"^model\s+LiteLLM_SpendLogs\s*\{")
HUNK: Final = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@[^\n]*\n", re.MULTILINE)
class Reviewer(BaseModel):
@ -37,31 +45,77 @@ class Review(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
state: str
user: Reviewer
commit_id: str | None = None
_REVIEWS: Final = TypeAdapter(list[Review])
@dataclass(frozen=True, slots=True)
class Block:
"""Lines of `model LiteLLM_SpendLogs { ... }` in one schema.prisma copy at the merge base and at the head."""
base: range
head: range
@dataclass(frozen=True, slots=True)
class Hunk:
removed: tuple[tuple[int, str], ...]
added: tuple[tuple[int, str], ...]
def watched(path: str) -> bool:
return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql"))
def _added_lines(section: str) -> tuple[str, ...]:
body: Final = section.split("\n@@", 1)
return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else ()
def model_block(schema: str) -> range:
"""1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none."""
lines: Final = schema.splitlines()
starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line))
if not starts:
return range(0)
ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}"))
return range(starts[0], (ends[0] if ends else len(lines)) + 1)
def spend_log_hits(diff: str) -> tuple[str, ...]:
"""`path: added line` for every added line in a watched file that names the table."""
def _hunks(section: str) -> tuple[Hunk, ...]:
headers: Final = tuple(HUNK.finditer(section))
ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),)
return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends))
def _hunk(header: re.Match[str], body: str) -> Hunk:
lines: Final = body.splitlines()
return Hunk(
removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))),
added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))),
)
def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]:
"""`path: +line` or `path: -line` for every changed line that touches the table.
A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines,
merge base for removed ones), a migration line when it is added and names the table."""
sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE)
return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line))
def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]:
pairs: Final = tuple((_path_of(section), section) for section in sections)
paths: Final = tuple((_path_of(section), section) for section in sections)
return tuple(
(path, line) for path, section in pairs if watched(path) for line in _added_lines(section)
) # comprehension-ok: path then its lines
hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks)
) # comprehension-ok: path then its hits
def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]:
if path not in blocks:
return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line))
block: Final = blocks[path]
removed: Final = tuple(
f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base
) # comprehension-ok: hunk then its lines
added: Final = tuple(
f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head
) # comprehension-ok: hunk then its lines
return removed + added
def _path_of(section: str) -> str:
@ -70,11 +124,33 @@ def _path_of(section: str) -> str:
def approved_by(reviews: Sequence[Review], login: str) -> bool:
"""True when `login`'s most recent review on the pull request is an approval."""
states: Final = tuple(review.state for review in reviews if review.user.login == login)
"""True when `login`'s most recent review that changes GitHub's approval state is an approval.
COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too.
"""
own: Final = tuple(review.state for review in reviews if review.user.login == login)
states: Final = tuple(state for state in own if state in DECIDING_STATES)
return bool(states) and states[-1] == "APPROVED"
def latest_approval(reviews: Sequence[Review], login: str) -> Review | None:
"""`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes."""
own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED")
return own[-1] if own else None
def approval_covers(
approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None]
) -> bool:
"""True when the approval is on the head, or the head adds no spend-log line the approved commit did not."""
if approval.commit_id is None:
return False
if approval.commit_id == head:
return True
reviewed: Final = hits_at(approval.commit_id)
return reviewed is not None and set(hits) <= set(reviewed)
def added_diff(base: str, head: str) -> str:
result: Final = subprocess.run(
("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX),
@ -85,6 +161,27 @@ def added_diff(base: str, head: str) -> str:
return result.stdout
def file_at(revision: str, path: str) -> str:
"""Contents of `path` at `revision`, empty when the file does not exist there."""
result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True)
return result.stdout if result.returncode == 0 else ""
def schema_blocks(base: str, head: str) -> Mapping[str, Block]:
return MappingProxyType(
{path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED}
)
def fetched(commit: str) -> bool:
result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False)
return result.returncode == 0
def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None:
return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None
class Arguments(BaseModel):
model_config = ConfigDict(frozen=True)
base: str
@ -100,18 +197,28 @@ def main(argv: Sequence[str]) -> int:
parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge")
parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews")
args: Final = Arguments.model_validate(vars(parser.parse_args(argv)))
hits: Final = spend_log_hits(added_diff(args.base, args.head))
hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(args.base, args.head))
if not hits:
sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n")
sys.stdout.write(
"No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; "
"nothing to review.\n"
)
return 0
if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver):
reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes())
approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None
if approval is not None and approval_covers(
approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit)
):
sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n")
return 0
listed: Final = "".join(f" {hit}\n" for hit in hits)
sys.stdout.write(
reason: Final = (
f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n"
f"{listed}The check re-runs when a review is submitted.\n"
if approval is None
else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, "
f"a new approving review is needed:\n"
)
sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n")
return 1

View file

@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
from litellm.types.guardrails import GuardrailEventHooks
@ -753,7 +754,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
_masked_entity_count.set({})
# Covers multimodal list content + Responses-API input.
inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data)
inspected = walk_user_text(
data,
self.redact_text,
self.redact_tool_payload,
skip_tool_message=effective_skip_tool_message_for_guardrail(self),
) + self._redact_prompt(data)
if inspected == 0:
# Image-only, empty-text, and unsupported payloads inspected

View file

@ -28,7 +28,7 @@ import httpx
from fastapi import HTTPException
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -53,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_metadata,
)
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
@ -309,6 +309,7 @@ def _is_xai_batch_output(
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
model_mappings: Mapping[str, str]
flat_model_file_ids: Sequence[str]
storage_backend: Optional[str]
storage_url: Optional[str]
@ -362,6 +363,8 @@ class _ManagedResourceDatabase(Protocol):
@property
def litellm_managedfiletable(self) -> _ManagedFileTableActions: ...
async def query_raw(self, query: str, *args: object) -> object: ...
@property
def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ...
@ -385,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False):
custom_llm_provider: ReadOnly[str | None]
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
return prisma_client.db
def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions:
return prisma_client.db.litellm_managedfiletable
class _ListedFileRow(Protocol):
@property
def unified_file_id(self) -> str: ...
@property
def file_object(self) -> object: ...
@property
def model_mappings(self) -> Mapping[str, object]: ...
@property
def storage_backend(self) -> str | None: ...
class _ScopedFileRow(BaseModel):
model_config = ConfigDict(frozen=True)
unified_file_id: str
file_object: JsonValue = None
model_mappings: Mapping[str, JsonValue]
storage_backend: str | None = None
_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...])
def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]:
match owner_filter:
case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}:
return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id)
case {"created_by": str() as user_id}:
return '"created_by" = $1', (user_id,)
case {"team_id": str() as team_id}:
return '"team_id" = $1', (team_id,)
case _:
return ("FALSE", ()) if owner_filter else ("TRUE", ())
async def _scoped_file_rows(
database: _ManagedResourceDatabase,
owner_filter: Mapping[str, object],
model_mapping_keys: frozenset[str],
cursor_id: str | None,
direction: Literal["asc", "desc"],
take: int,
) -> tuple[_ScopedFileRow, ...]:
owner_condition, owner_params = _owner_condition(owner_filter)
keys_param: Final = len(owner_params) + 1
comparison: Final = "<" if direction == "desc" else ">"
cursor_condition: Final = (
f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" '
f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) '
if cursor_id
else ""
)
order: Final = direction.upper()
rows: Final = await database.query_raw(
'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" '
f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) '
f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} '
f"LIMIT ${keys_param + (2 if cursor_id else 1)}",
*owner_params,
sorted(model_mapping_keys),
*((cursor_id,) if cursor_id else ()),
take,
)
return _SCOPED_FILE_ROWS.validate_python(rows)
def _iter_provider_file_id_pairs(
rows: Sequence[_ManagedFileRow],
requested_provider_file_ids: frozenset[str],
@ -2265,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
limit: Optional[int] = None,
after: Optional[str] = None,
model_scope: Optional[FileListModelScope] = None,
before: Optional[str] = None,
**data: Dict,
) -> FileListPage:
"""List the managed files the caller owns, newest first.
Pagination is keyset based on ``unified_file_id`` so a key that owns
every file on the proxy still reads one bounded page at a time.
``purpose`` is applied after parsing, because the managed file table
keeps it inside the ``file_object`` blob instead of a column, and rows
``after`` pages toward older files and ``before`` toward newer ones,
each page still ordered newest first.
``model_scope`` is applied in the query, so only rows whose
``model_mappings`` name one of its deployment ids or model names, and
storage-backed rows, leave the database; a storage-backed row then
stays only when its model name is one the router sends to the listed
provider. ``purpose`` is applied after parsing, because the managed
file table keeps it inside the ``file_object`` blob instead of a
column. Rows
whose blob will not parse drop out there too, so a chunk of rows can
yield fewer matches than the page holds. Successive chunks are read
until the page is full or the caller's rows run out, which keeps
@ -2281,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
scan has to continue past it, the chunk widens to
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
that many rows instead of one per page. That bound is per query, not
per request: the work is still linear in the rows the caller owns, and
a filter matching nothing reads every one of them, with no index
covering either the owner filter or the sort.
per request: a ``purpose`` filter matching nothing still reads every
row that passes the query.
"""
validate_file_list_limit(limit)
validate_file_list_purpose(purpose)
@ -2292,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if owner_filter is None:
return FileListPage.model_validate(build_list_page([]))
if after:
cursor_row = await _managed_file_table(self.prisma_client).find_first(
where={**owner_filter, "unified_file_id": after}
if after and before:
raise ProxyException(
message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.",
type="invalid_request_error",
param="before_id",
code=400,
openai_code="invalid_value",
)
cursor_param: Final = "before_id" if before else "after"
cursor: Final = before or after
database: Final = _managed_database(self.prisma_client)
if cursor:
cursor_row = await database.litellm_managedfiletable.find_first(
where={**owner_filter, "unified_file_id": cursor}
)
if cursor_row is None:
raise ProxyException(
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.",
type="invalid_request_error",
param="after",
param=cursor_param,
code=400,
openai_code="invalid_value",
)
if model_scope is not None and not model_scope.model_mapping_keys:
return FileListPage.model_validate(build_list_page([]))
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
matches: Final[List[OpenAIFileObject]] = []
cursor_id = after
cursor_id = cursor
chunk_size = page_size + 1
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
while len(matches) <= page_size:
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_file_table(self.prisma_client).find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
**cursor_args,
chunk: Sequence[_ListedFileRow] = (
await _scoped_file_rows(
database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size
)
if model_scope is not None
else await database.litellm_managedfiletable.find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": direction}, {"unified_file_id": direction}],
**cursor_args,
)
)
matches.extend(
_public_file_object(parsed_file_object, row.unified_file_id)
for row in chunk
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
if (
model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None)
)
and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
and (purpose is None or parsed_file_object.purpose == purpose)
)
if len(chunk) < chunk_size:
@ -2329,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
cursor_id = chunk[-1].unified_file_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size))
page: Final = matches[:page_size]
return FileListPage.model_validate(
build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size)
)
def _is_batch_polling_enabled(self) -> bool:
"""

View file

@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
"""The bytes a value takes in the index entry: its length, or its stored size when
Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
from psycopg import sql
if column.type_name in _OCTET_LENGTH_TYPES:
return sql.SQL("octet_length({})").format(identifier)
return sql.SQL("octet_length({}::text)").format(identifier)
length: Final = (
sql.SQL("octet_length({})").format(identifier)
if column.type_name in _OCTET_LENGTH_TYPES
else sql.SQL("octet_length({}::text)").format(identifier)
)
return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":

View file

@ -0,0 +1,2 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx"
ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC);

View file

@ -0,0 +1,47 @@
SET LOCAL lock_timeout = '10s';
CREATE TYPE "AccountingEventKind" AS ENUM ('inbound_usage', 'outbound_usage', 'cost');
CREATE TYPE "AccountingDeliveryState" AS ENUM ('pending', 'applied', 'parked');
CREATE TYPE "AccountingDeliveryTarget" AS ENUM ('LiteLLM_VerificationToken', 'LiteLLM_UserTable', 'LiteLLM_TeamTable', 'LiteLLM_TeamMembership', 'LiteLLM_OrganizationTable', 'LiteLLM_OrganizationMembership', 'LiteLLM_EndUserTable', 'LiteLLM_AgentsTable', 'LiteLLM_TagTable', 'LiteLLM_ProjectTable', 'LiteLLM_ModelAccessGroupBudgetTable', 'LiteLLM_BudgetTable', 'LiteLLM_BudgetWindowSpend', 'LiteLLM_DailyUserSpend', 'LiteLLM_DailyTeamSpend', 'LiteLLM_DailyOrganizationSpend', 'LiteLLM_DailyEndUserSpend', 'LiteLLM_DailyAgentSpend', 'LiteLLM_DailyTagSpend', 'LiteLLM_DailyModelUsage', 'LiteLLM_DailyToolSpend', 'LiteLLM_AutoRouterDailySpend');
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingEvent" (
"event_id" TEXT NOT NULL,
"kind" "AccountingEventKind" NOT NULL,
"source_event_id" TEXT,
"schema_version" INTEGER NOT NULL,
"payload_hash" TEXT NOT NULL,
"occurred_at" TIMESTAMP(3) NOT NULL,
"recorded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"payload" JSONB NOT NULL,
CONSTRAINT "LiteLLM_AccountingEvent_schema_version_check" CHECK ("schema_version" > 0),
CONSTRAINT "LiteLLM_AccountingEvent_pkey" PRIMARY KEY ("event_id")
);
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingDelivery" (
"event_id" TEXT NOT NULL,
"target" "AccountingDeliveryTarget" NOT NULL,
"version" INTEGER NOT NULL,
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
"attempts" INTEGER NOT NULL DEFAULT 0,
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"applied_at" TIMESTAMP(3),
"last_error" TEXT,
CONSTRAINT "LiteLLM_AccountingDelivery_version_check" CHECK ("version" > 0),
CONSTRAINT "LiteLLM_AccountingDelivery_attempts_check" CHECK ("attempts" >= 0),
CONSTRAINT "LiteLLM_AccountingDelivery_pkey" PRIMARY KEY ("event_id","target","version")
);
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_recorded_at_idx" ON "LiteLLM_AccountingEvent"("recorded_at");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_kind_source_event_id_idx" ON "LiteLLM_AccountingEvent"("kind", "source_event_id");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingDelivery_target_state_next_attempt_at_idx" ON "LiteLLM_AccountingDelivery"("target", "state", "next_attempt_at");

View file

@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
@@index([unified_file_id])
@@index([flat_model_file_ids], type: Gin)
@@index([team_id, created_at(sort: Desc)])
@@index([created_by, created_at(sort: Desc)])
}
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport {
@@index([window_end])
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
}

View file

@ -1,6 +1,8 @@
---
name: rust-error
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
metadata:
internal: true
---
# Rust errors

View file

@ -1,6 +1,8 @@
---
name: rust-gateway
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
metadata:
internal: true
---
# Rust gateway

View file

@ -1,6 +1,8 @@
---
name: rust-string-enums
description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas
metadata:
internal: true
---
# Rust string enums

View file

@ -1,6 +1,8 @@
---
name: rust-tracing
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
metadata:
internal: true
---
# Rust tracing

View file

@ -13,6 +13,7 @@ const FINISH_REASONS: &[(&str, &str)] = &[
("max_tokens", "length"),
("refusal", "content_filter"),
("compaction", "length"),
("model_context_window_exceeded", "length"),
("guardrail_intervened", "content_filter"),
("content_filtered", "content_filter"),
("content_filter", "content_filter"),
@ -74,6 +75,7 @@ mod tests {
#[case::end_turn("end_turn", "stop")]
#[case::stop_sequence("stop_sequence", "stop")]
#[case::max_tokens("max_tokens", "length")]
#[case::model_context_window_exceeded("model_context_window_exceeded", "length")]
#[case::refusal("refusal", "content_filter")]
#[case::guardrail_intervened("guardrail_intervened", "content_filter")]
#[case::content_filtered("content_filtered", "content_filter")]

View file

@ -0,0 +1,62 @@
"""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.types import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
DeliveryState,
DeliveryTarget,
Dimension,
DimensionFilter,
DimensionName,
EventKind,
Granularity,
InboundUsageEvent,
Measurement,
Metric,
MetricName,
Multiplier,
NamedCost,
OutboundUsageEvent,
Outcome,
PricedQuantity,
PricingSnapshot,
Rate,
Unpriced,
validate_aggregate_rows,
)
__all__ = [
"AdmissionDecision",
"AggregateQuery",
"AggregateRow",
"BudgetEnforcer",
"CostAggregator",
"CostCompiler",
"CostEvent",
"DeliveryState",
"DeliveryTarget",
"Dimension",
"DimensionFilter",
"DimensionName",
"EventKind",
"Granularity",
"InboundUsageEvent",
"Measurement",
"Metric",
"MetricName",
"Multiplier",
"NamedCost",
"OutboundUsageEvent",
"Outcome",
"PricedQuantity",
"PricingSnapshot",
"Rate",
"Unpriced",
"UsageMeter",
"measurements_for",
"validate_aggregate_rows",
]

View file

@ -0,0 +1,13 @@
from typing import Final
from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent
def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]:
"""Return the complete contribution set for one event, ready for one apply call."""
metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics
occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at
return tuple(
Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric)
for metric in metrics or ()
)

View file

@ -0,0 +1,85 @@
"""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 must survive request teardown. Cost-map rates are resolved at compilation, not capture.
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 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 by event identity and cost-map rates
at compilation. 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.
"""
...
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 query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]:
"""Return disjoint [start, end) rows per group using existing daily projections.
Filters combine with AND, values within a filter with IN. DAY emits daily
buckets; None emits whole-interval totals. Reject unsupported metric/scope
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."""
...

428
litellm/accounting/types.py Normal file
View file

@ -0,0 +1,428 @@
"""Accounting values use UTC instants, decimal strings on the wire and USD charges."""
from collections.abc import Mapping
from datetime import datetime, timezone
from decimal import Decimal
from enum import Enum
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
from litellm.types.utils import CallTypes
_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True)
Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")]
Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)]
def _utc(value: datetime) -> datetime:
return value.astimezone(timezone.utc)
Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)]
class MetricName(str, Enum):
"""Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens."""
REQUESTS = "requests"
SUCCESSFUL_REQUESTS = "successful_requests"
FAILED_REQUESTS = "failed_requests"
PROMPT_TOKENS = "prompt_tokens"
COMPLETION_TOKENS = "completion_tokens"
REASONING_TOKENS = "reasoning_tokens"
TOTAL_TOKENS = "total_tokens"
CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens"
CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens"
CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr"
SPEND = "spend"
class DimensionName(str, Enum):
KEY = "key"
USER = "user"
TEAM = "team"
ORGANIZATION = "organization"
END_USER = "end_user"
MODEL = "model"
MODEL_GROUP = "model_group"
DEPLOYMENT = "deployment"
PROVIDER = "provider"
TOOL = "tool"
AGENT = "agent"
PROJECT = "project"
MODEL_ACCESS_GROUP = "model_access_group"
API_BASE = "api_base"
SERVICE_TIER = "service_tier"
TAG = "tag"
class Outcome(str, Enum):
SUCCEEDED = "succeeded"
FAILED = "failed"
CANCELLED = "cancelled"
class Granularity(str, Enum):
DAY = "day"
class EventKind(str, Enum):
INBOUND_USAGE = "inbound_usage"
OUTBOUND_USAGE = "outbound_usage"
COST = "cost"
class DeliveryState(str, Enum):
PENDING = "pending"
APPLIED = "applied"
PARKED = "parked"
class DeliveryTarget(str, Enum):
"""One member per existing table the projector writes; each is applied and receipted independently."""
KEY_SPEND = "LiteLLM_VerificationToken"
USER_SPEND = "LiteLLM_UserTable"
TEAM_SPEND = "LiteLLM_TeamTable"
TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership"
ORGANIZATION_SPEND = "LiteLLM_OrganizationTable"
ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership"
END_USER_SPEND = "LiteLLM_EndUserTable"
AGENT_SPEND = "LiteLLM_AgentsTable"
TAG_SPEND = "LiteLLM_TagTable"
PROJECT_SPEND = "LiteLLM_ProjectTable"
MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable"
BUDGET_SPEND = "LiteLLM_BudgetTable"
BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend"
DAILY_USER_SPEND = "LiteLLM_DailyUserSpend"
DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend"
DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend"
DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend"
DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend"
DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend"
DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage"
DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend"
AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend"
_SUBSET_PAIRS: Final = (
(MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS),
(MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS),
)
def _check_subsets(values: Mapping[str, Decimal]) -> None:
for part, whole in _SUBSET_PAIRS:
if part.value in values and whole.value in values and values[part.value] > values[whole.value]:
raise ValueError(f"{part.value} cannot exceed {whole.value}")
class Metric(BaseModel):
model_config = _CONFIG
name: MetricName
value: Quantity
unit: Literal["count", "USD"] = "count"
@model_validator(mode="after")
def validate_unit(self) -> Self:
if (self.name is MetricName.SPEND) != (self.unit == "USD"):
raise ValueError("SPEND requires USD; usage metrics require count")
if self.unit == "count" and self.value != self.value.to_integral_value():
raise ValueError("Count metrics require integral values")
return self
class Dimension(BaseModel):
model_config = _CONFIG
name: DimensionName
value: Identifier
_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP})
def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]:
if len(frozenset(values)) != len(values):
raise ValueError("Duplicate dimension values")
single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS)
if len(frozenset(single_names)) != len(single_names):
raise ValueError("Only tag and model access group dimensions may have multiple values")
return values
Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)]
_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS))
def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]:
if len(frozenset(value.name for value in values)) != len(values):
raise ValueError("Duplicate metric names")
return values
Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)]
def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None:
if end < start or (end == start and not allow_empty):
raise ValueError("End must follow start")
class InboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
request_id: Identifier
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics
@model_validator(mode="after")
def validate_caller_activity(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if any(metric.name not in _CALLER_METRICS for metric in self.metrics):
raise ValueError("Inbound events own only caller request metrics")
if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics):
raise ValueError("Inbound events count exactly one caller request")
for metric in self.metrics:
if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED):
raise ValueError("Successful request count must match outcome")
if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED):
raise ValueError("Failed request count must match outcome")
return self
class OutboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
attempt_id: Identifier
request_id: Identifier | None
retry_of_attempt_id: Identifier | None
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics | None
@model_validator(mode="after")
def validate_attempt_usage(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if self.retry_of_attempt_id == self.attempt_id:
raise ValueError("An attempt cannot retry itself")
if self.metrics is not None and any(
metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics
):
raise ValueError("Outbound events own only attempt usage metrics")
if self.metrics is not None:
_check_subsets({metric.name.value: metric.value for metric in self.metrics})
return self
PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"]
class Rate(BaseModel):
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
usd_per_unit: Quantity
class PricedQuantity(BaseModel):
"""A pricing input in its billing unit. Only seconds may be fractional."""
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
value: Quantity
@model_validator(mode="after")
def validate_discrete_quantity(self) -> Self:
if self.unit != "second" and self.value != self.value.to_integral_value():
raise ValueError("Discrete priced quantities require integral values")
return self
class Multiplier(BaseModel):
"""A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5."""
model_config = _CONFIG
name: Identifier
value: Identifier
factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)]
class NamedCost(BaseModel):
"""A free-form extra charge by name, e.g. azure_model_router_flat_cost."""
model_config = _CONFIG
name: Identifier
amount: Quantity
class PricingSnapshot(BaseModel):
"""Retained inputs and result of authoritative compilation, reproducible after cost-map rates change.
Request-scoped inputs and overrides come from the usage row's versioned payload. Cost-map rates are
resolved at compilation, not frozen at capture. Retries reuse the committed charge and snapshot.
Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge.
Component sums and float tolerances belong to reconciliation, not this structural contract.
"""
model_config = _CONFIG
schema_version: Literal[1] = 1
calculator_version: Identifier
model: Identifier
provider: Identifier | None
pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"]
priced_model_key: Identifier
multipliers: tuple[Multiplier, ...] = ()
rates: tuple[Rate, ...] = ()
usage: tuple[PricedQuantity, ...] = ()
input_cost: Quantity | None = None
output_cost: Quantity | None = None
cache_read_cost: Quantity | None = None
cache_creation_cost: Quantity | None = None
reasoning_cost: Quantity | None = None
tool_usage_cost: Quantity | None = None
guardrail_cost: Quantity | None = None
additional_costs: tuple[NamedCost, ...] = ()
original_cost: Quantity | None = None
discount_percent: Quantity | None = None
discount_amount: Quantity | None = None
margin_percent: Quantity | None = None
margin_fixed_amount: Quantity | None = None
margin_total_amount: Quantity | None = None
upstream_reported_cost: Quantity | None = None
total_cost: Quantity
@model_validator(mode="after")
def validate_snapshot(self) -> Self:
rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates)
quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage)
if len(rate_keys) != len(self.rates):
raise ValueError("Duplicate rate measures and units")
if len(quantity_keys) != len(self.usage):
raise ValueError("Duplicate priced quantity measures and units")
if not rate_keys.issubset(quantity_keys):
raise ValueError("Every rate requires one priced quantity with the same measure and unit")
if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs):
raise ValueError("Duplicate additional cost names")
if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers):
raise ValueError("Duplicate multiplier names")
_check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"})
return self
class CostEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
usage_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
charges: Metrics
snapshot: PricingSnapshot
@model_validator(mode="after")
def validate_charge(self) -> Self:
if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND:
raise ValueError("A cost event requires one authoritative USD spend metric")
if self.event_id == self.usage_event_id:
raise ValueError("Cost and usage events require distinct identities")
if self.snapshot.total_cost != self.charges[0].value:
raise ValueError("Snapshot total must equal the authoritative spend")
return self
class Unpriced(BaseModel):
model_config = _CONFIG
usage_event_id: Identifier
reason: Identifier
class Measurement(BaseModel):
model_config = _CONFIG
source_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
metric: Metric
class DimensionFilter(BaseModel):
model_config = _CONFIG
name: DimensionName
values: Annotated[tuple[Identifier, ...], Field(min_length=1)]
class AggregateQuery(BaseModel):
model_config = _CONFIG
metric: MetricName
filters: tuple[DimensionFilter, ...]
group_by: tuple[DimensionName, ...]
start: Instant
end: Instant
granularity: Granularity | None
@model_validator(mode="after")
def validate_query(self) -> Self:
_interval(self.start, self.end)
if len(frozenset(item.name for item in self.filters)) != len(self.filters):
raise ValueError("Use one IN filter per dimension; filters combine with AND")
if len(frozenset(self.group_by)) != len(self.group_by):
raise ValueError("Duplicate group-by dimensions")
if any(
(instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0)
for instant in (self.start, self.end)
):
raise ValueError("Initial reporting queries require UTC day boundaries")
return self
class AggregateRow(BaseModel):
model_config = _CONFIG
start: Instant
end: Instant
dimensions: Dimensions
value: Quantity
@model_validator(mode="after")
def validate_bounds(self) -> Self:
_interval(self.start, self.end)
return self
def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]:
"""Validate [start, end) buckets per group; overlapping tag groups are not additive."""
groups: Final = frozenset(frozenset(row.dimensions) for row in rows)
for group in groups:
ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start)
if any(left.end > right.start for left, right in zip(ordered, ordered[1:])):
raise ValueError("Aggregate buckets overlap within a dimension group")
return rows
class AdmissionDecision(BaseModel):
model_config = _CONFIG
allowed: bool
reason: Identifier | None = None

View file

@ -274,7 +274,9 @@ class InMemoryCache(BaseCache):
) -> list[float] | None:
results: Final = []
for increment in increment_list:
result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs)
result = await self.async_increment(
increment["key"], increment["increment_value"], ttl=increment["ttl"], **kwargs
)
results.append(result)
return results

View file

@ -129,6 +129,7 @@ MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
MAX_FILE_LIST_LIMIT: Final = 10000
ANTHROPIC_DEFAULT_FILE_LIST_LIMIT: Final = 20
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10

View file

@ -1,9 +1,9 @@
import asyncio
import contextvars
import importlib
from collections.abc import Coroutine, Mapping
from collections.abc import Collection, Coroutine, Mapping
from functools import partial
from typing import TYPE_CHECKING, Final, Literal, cast, overload
from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
from pydantic import TypeAdapter, ValidationError
@ -61,6 +61,7 @@ from litellm.types.utils import (
ProviderRoutingHints,
)
from litellm.utils import (
EXTRA_BODY_ROUTING_KEYS,
ImageResponse,
ProviderConfigManager,
get_llm_provider,
@ -466,7 +467,7 @@ def image_generation(
litellm.LlmProviders.QWENCLOUD,
litellm.LlmProviders.QWEN_AI_PLATFORM,
litellm.LlmProviders.EDENAI,
):
) or (custom_llm_provider == litellm.LlmProviders.TENCENT and image_generation_config is not None):
if image_generation_config is None:
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
@ -807,6 +808,32 @@ def image_variation(
return response
_IMAGE_EDIT_OPENAI_PARAMS: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams))
_DROP_PARAM_KWARGS: Final = frozenset({"drop_params", "additional_drop_params"})
_ADDITIONAL_DROP_PARAMS: Final = TypeAdapter[tuple[str, ...]](tuple[str, ...])
def _listed_drop_params(additional_drop_params: object) -> tuple[str, ...]:
try:
return _ADDITIONAL_DROP_PARAMS.validate_python(additional_drop_params)
except ValidationError:
return ()
def _tokenhub_passthrough_edit_params(
non_default_params: Mapping[str, object],
extra_body: Mapping[str, object] | None,
mapped_params: Collection[str],
additional_drop_params: Collection[str] | None,
) -> Mapping[str, object]:
dropped: Final = frozenset(additional_drop_params or ())
unmapped_excluded: Final = _DROP_PARAM_KWARGS | _IMAGE_EDIT_OPENAI_PARAMS | frozenset(mapped_params) | dropped
return {
**{key: value for key, value in non_default_params.items() if key not in unmapped_excluded},
**{key: value for key, value in (extra_body or {}).items() if key not in dropped | EXTRA_BODY_ROUTING_KEYS},
}
@client
def image_edit(
image: FileTypes | list[FileTypes] | None = None,
@ -935,22 +962,24 @@ def image_edit(
raise ValueError(f"image edit is not supported for {custom_llm_provider}")
local_vars.update(kwargs)
provider_supported_params: Final[frozenset[str]] = frozenset(
image_edit_provider_config.get_supported_openai_params(model)
)
# Get ImageEditOptionalRequestParams with only valid parameters
image_edit_optional_params: Final[ImageEditOptionalRequestParams] = (
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
local_vars,
provider_supported_params=frozenset(
image_edit_provider_config.get_supported_openai_params(model)
).intersection(non_default_params),
provider_supported_params=provider_supported_params.intersection(non_default_params),
)
)
# Get optional parameters for the responses API
additional_drop_params: Final = _listed_drop_params(kwargs.get("additional_drop_params"))
image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit(
model=model,
image_edit_provider_config=image_edit_provider_config,
image_edit_optional_params=image_edit_optional_params,
drop_params=kwargs.get("drop_params"),
additional_drop_params=kwargs.get("additional_drop_params"),
additional_drop_params=additional_drop_params,
)
if image_edit_provider_config.use_multipart_form_data() and (
@ -1039,7 +1068,19 @@ def image_edit(
image=images,
prompt=prompt,
image_edit_provider_config=image_edit_provider_config,
image_edit_optional_request_params=image_edit_request_params,
image_edit_optional_request_params={
**image_edit_request_params,
**(
_tokenhub_passthrough_edit_params(
non_default_params,
extra_body if isinstance(extra_body, dict) else None,
provider_supported_params,
additional_drop_params,
)
if custom_llm_provider == litellm.LlmProviders.TENCENT.value
else {}
),
},
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,

View file

@ -1,6 +1,9 @@
import base64
import os
from collections.abc import Collection, Mapping
from io import BufferedReader, BytesIO
from typing import Any, Final, cast, get_type_hints
from pathlib import Path
from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable
import litellm
from litellm.litellm_core_utils.token_counter import get_image_type
@ -16,7 +19,7 @@ class ImageEditRequestUtils:
image_edit_provider_config: BaseImageEditConfig,
image_edit_optional_params: ImageEditOptionalRequestParams,
drop_params: bool | None = None,
additional_drop_params: list[str] | None = None,
additional_drop_params: Collection[str] | None = None,
) -> dict:
"""
Get optional parameters for the image edit API.
@ -145,3 +148,36 @@ class ImageEditRequestUtils:
except Exception:
# If anything goes wrong, default to PNG
return FILE_MIME_TYPES[FileType.PNG]
@runtime_checkable
class _SeekableBinaryReader(Protocol):
def tell(self) -> int: ...
def seek(self, offset: int) -> int: ...
def read(self) -> bytes: ...
def _read_image_bytes(image: object) -> bytes:
if isinstance(image, bytes):
return image
if isinstance(image, tuple):
return _read_image_bytes(image[1])
if isinstance(image, os.PathLike):
return Path(image).read_bytes()
if isinstance(image, _SeekableBinaryReader):
position: Final = image.tell()
image.seek(0)
data: Final = image.read()
image.seek(position)
return data
raise ValueError(f"Unsupported image type for image edit: {type(image).__name__}")
def image_to_data_url(image: object) -> str:
if isinstance(image, str):
return image
image_bytes: Final = _read_image_bytes(image)
mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes)
return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"

View file

@ -338,6 +338,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_ToolTable",
"LiteLLM_TelemetryReport",
"LiteLLM_AccessGroupTable",
"LiteLLM_AccountingDelivery",
"LiteLLM_AccountingEvent",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_MemoryTable",
"LiteLLM_AdaptiveRouterState",
@ -434,6 +436,7 @@ _POSTGRES_OPERATION_BY_CALL_TYPE: Final[Mapping[str, PostgresOperation]] = Mappi
"reset_budget_windows": PostgresOperation("select", None),
"write_budget_windows": PostgresOperation("update", None),
"roll_window_spend_row": PostgresOperation("update", None),
"open_window_spend_row": PostgresOperation("upsert", None),
"seed_window_spend": PostgresOperation("select", None),
"select_window_spend_rows": PostgresOperation("select", None),
"commit_window_spend_updates": PostgresOperation("upsert", None),

View file

@ -97,6 +97,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
_RESPONSE_CONTENT_FIELD: Final = "content"
_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:"
_ResponseT: Final = TypeVar("_ResponseT")
@ -363,7 +365,8 @@ class WebSearchInterceptionLogger(CustomLogger):
get_last_user_message,
)
query: Final = get_last_user_message(cast(list[AllMessageValues], messages))
last_user_message: Final = get_last_user_message(cast(list[AllMessageValues], messages)) or ""
query: Final = last_user_message.removeprefix(_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX).strip()
if not query:
return None

View file

@ -266,6 +266,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = {
"tool_use": "tool_calls",
"refusal": "content_filter",
"compaction": "length",
"model_context_window_exceeded": "length",
# Cohere
"COMPLETE": "stop",
"ERROR_TOXIC": "content_filter",

View file

@ -75,6 +75,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
"rpm",
"itpm",
"otpm",
"tpd",
"rpd",
"use_xai_oauth",
"chatgpt_auth_file",
"github_copilot_auth_type",

View file

@ -842,6 +842,7 @@ class Logging(LiteLLMLoggingBaseClass):
self._async_success_enqueued: bool = False
self._sync_success_scheduled: bool = False
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self._stream_failure_claimed: bool = False
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
@ -2561,6 +2562,12 @@ class Logging(LiteLLMLoggingBaseClass):
},
)
def claim_stream_failure_report(self) -> bool:
if self._stream_failure_claimed:
return False
self._stream_failure_claimed = True
return True
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
usage: Final = getattr(assembled, "usage", None)
@ -2751,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass):
if entry is None or not is_str_keyed_dict(litellm_params):
return
metadata: Final[object] = litellm_params.get("metadata")
bucket: Final = metadata if is_str_keyed_dict(metadata) else {}
bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {}
litellm_params["metadata"] = bucket
with_guardrail_entry(bucket, entry)

View file

@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from datetime import datetime, timezone, tzinfo
from types import MappingProxyType
from typing import Final, Literal, TypedDict, cast
from typing import Final, Literal, Protocol, TypedDict, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from typing_extensions import ReadOnly
@ -2085,6 +2085,41 @@ def calculate_image_response_web_search_cost(
)
class _ModelPricedImageCostCalculator(Protocol):
def __call__(self, model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: ...
def _model_priced_image_cost_calculator(
custom_llm_provider: str | None, model: str
) -> _ModelPricedImageCostCalculator | None:
match custom_llm_provider:
case litellm.LlmProviders.RECRAFT.value:
from litellm.llms.recraft.cost_calculator import cost_calculator as recraft_image_cost_calculator
return recraft_image_cost_calculator
case litellm.LlmProviders.AIML.value:
from litellm.llms.aiml.image_generation.cost_calculator import cost_calculator as aiml_image_cost_calculator
return aiml_image_cost_calculator
case litellm.LlmProviders.COMETAPI.value:
from litellm.llms.cometapi.image_generation.cost_calculator import (
cost_calculator as cometapi_image_cost_calculator,
)
return cometapi_image_cost_calculator
case litellm.LlmProviders.RUNWAYML.value:
from litellm.llms.runwayml.cost_calculator import cost_calculator as runwayml_image_cost_calculator
return runwayml_image_cost_calculator
case litellm.LlmProviders.TENCENT.value:
from litellm.llms.tencent.cost_calculator import image_cost_calculator as tencent_image_cost_calculator
from litellm.llms.tencent.image_generation.transformation import is_tokenhub_image_model
return tencent_image_cost_calculator if is_tokenhub_image_model(model) else None
case _:
return None
class CostCalculatorUtils:
@staticmethod
def call_type_has_image_response(call_type: str) -> bool:
@ -2131,9 +2166,6 @@ class CostCalculatorUtils:
from litellm.llms.gemini.image_generation.cost_calculator import (
cost_calculator as gemini_image_cost_calculator,
)
from litellm.llms.recraft.cost_calculator import (
cost_calculator as recraft_image_cost_calculator,
)
from litellm.llms.vertex_ai.image_generation.cost_calculator import (
cost_calculator as vertex_ai_image_cost_calculator,
)
@ -2146,7 +2178,10 @@ class CostCalculatorUtils:
)
resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0)
pricing: Final = deployment_pricing(model_info)
model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model)
if model_priced_calculator is not None:
return model_priced_calculator(model=model, image_response=completion_response, model_info=pricing)
if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value:
if isinstance(completion_response, ImageResponse):
return vertex_ai_image_cost_calculator(
@ -2164,36 +2199,6 @@ class CostCalculatorUtils:
optional_params=optional_params,
)
raise TypeError("completion_response must be of type ImageResponse for bedrock image cost calculation")
elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value:
from litellm.llms.recraft.cost_calculator import (
cost_calculator as recraft_image_cost_calculator,
)
return recraft_image_cost_calculator(
model=model,
image_response=completion_response,
model_info=pricing,
)
elif custom_llm_provider == litellm.LlmProviders.AIML.value:
from litellm.llms.aiml.image_generation.cost_calculator import (
cost_calculator as aiml_image_cost_calculator,
)
return aiml_image_cost_calculator(
model=model,
image_response=completion_response,
model_info=pricing,
)
elif custom_llm_provider == litellm.LlmProviders.COMETAPI.value:
from litellm.llms.cometapi.image_generation.cost_calculator import (
cost_calculator as cometapi_image_cost_calculator,
)
return cometapi_image_cost_calculator(
model=model,
image_response=completion_response,
model_info=pricing,
)
elif custom_llm_provider == litellm.LlmProviders.GEMINI.value:
if call_type in (
CallTypes.image_edit.value,
@ -2237,16 +2242,6 @@ class CostCalculatorUtils:
optional_params=optional_params,
model_info=pricing,
)
elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value:
from litellm.llms.runwayml.cost_calculator import (
cost_calculator as runwayml_image_cost_calculator,
)
return runwayml_image_cost_calculator(
model=model,
image_response=completion_response,
model_info=pricing,
)
elif (
custom_llm_provider == litellm.LlmProviders.OPENAI.value
or custom_llm_provider == litellm.LlmProviders.AZURE.value

View file

@ -268,6 +268,7 @@ async def caller_within_model_budget(
user_id=user_id,
user_model_max_budget=user_model_max_budget,
model=model,
team_id=user_api_key_auth.team_id,
)
except litellm.BudgetExceededError:
return "denied"
@ -309,6 +310,7 @@ async def caller_within_model_budget(
end_user_id=end_user_id,
end_user_model_max_budget=end_user_model_max_budget,
model=model,
team_id=user_api_key_auth.team_id,
)
except litellm.BudgetExceededError:
return "denied"

View file

@ -2,15 +2,16 @@
## Translates OpenAI call to Anthropic `/v1/messages` format
import asyncio
from collections import deque
from collections.abc import AsyncGenerator, Iterator, Mapping
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final
import anyio
from pydantic import ConfigDict, field_validator
from litellm import verbose_logger
from litellm._logging import redact_internal_details_from_client_message
from litellm._uuid import uuid
from litellm.exceptions import MidStreamFallbackError
from litellm.exceptions import APIError, MidStreamFallbackError
from litellm.litellm_core_utils.prompt_templates.common_utils import (
encrypted_reasoning_signature,
)
@ -22,6 +23,7 @@ from litellm.llms.anthropic.messages.native.utils import (
from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults
from litellm.llms.anthropic.messages.utils import anthropic_sse_frame
from litellm.responses.streaming_iterator import (
map_stream_error_to_exception,
partial_responses_incomplete_event_on_disconnect,
stream_error_status_and_message,
)
@ -130,10 +132,12 @@ class AnthropicResponsesStreamWrapper:
) -> None:
self.responses_stream = responses_stream
self.model = model
self.litellm_logging_obj = litellm_logging_obj
self.safeguards_evaluator: Final = safeguards_evaluator
self._message_id: str = f"msg_{uuid.uuid4()}"
if litellm_logging_obj is not None:
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
self._unreported_failure: Exception | None = None
self._current_block_index: int = -1
# Map item_id -> content_block_index so we can stop the right block later
self._item_id_to_block_index: dict[str, int] = {}
@ -377,16 +381,24 @@ class AnthropicResponsesStreamWrapper:
if event_type == "response.failed":
failed: Final = _FailedResponseEvent.model_validate(event)
status_code, message = stream_error_status_and_message(
failed.response.error if failed.response is not None else None
)
upstream_error: Final = failed.response.error if failed.response is not None else None
status_code, message = stream_error_status_and_message(upstream_error)
verbose_logger.error(
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s",
self.model,
status_code,
message,
)
self._fail_stream(status_code, message)
self._fail_stream(
status_code,
message,
failure=map_stream_error_to_exception(
upstream_error,
self.model,
getattr(self.responses_stream, "custom_llm_provider", None) or "",
logging_obj=self.litellm_logging_obj,
),
)
return
# ---- response completed -> message_delta + message_stop ----
@ -439,10 +451,36 @@ class AnthropicResponsesStreamWrapper:
self._sent_message_stop = True
return
def _fail_stream(self, status_code: int, message: str) -> None:
def _fail_stream(self, status_code: int, message: str, failure: Exception | None = None) -> None:
self._stream_failed = True
self._unreported_failure = failure or APIError(
status_code=status_code,
message=message,
llm_provider=getattr(self.responses_stream, "custom_llm_provider", None),
model=self.model,
)
self._chunk_queue.append(_anthropic_error_chunk(status_code, message))
async def _report_unreported_failure(self) -> None:
failure: Final = self._unreported_failure
if failure is None:
return
self._unreported_failure = None
on_failure: Final[Callable[[Exception], Awaitable[None]] | None] = getattr(
self.litellm_logging_obj, "_on_detached_stream_failure", None
)
if on_failure is None:
return
try:
with anyio.CancelScope(shield=True):
await on_failure(failure)
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not hide the error event from the client
verbose_logger.warning(
"AnthropicResponsesStreamWrapper: stream failure hook raised: %s(%s)",
type(hook_failure).__name__,
hook_failure,
)
def __aiter__(self) -> "AnthropicResponsesStreamWrapper":
return self
@ -450,6 +488,7 @@ class AnthropicResponsesStreamWrapper:
if self._chunk_queue:
return self._chunk_queue.popleft()
if self._stream_failed:
await self._report_unreported_failure()
raise StopAsyncIteration
if not self._sent_message_start:
@ -478,7 +517,7 @@ class AnthropicResponsesStreamWrapper:
verbose_logger.exception(
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model
)
self._fail_stream(*_failure_status_and_message(e))
self._fail_stream(*_failure_status_and_message(e), failure=_original_failure(e))
if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed:
verbose_logger.error(
@ -498,8 +537,14 @@ class AnthropicResponsesStreamWrapper:
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator)
async for chunk in self:
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
try:
async for chunk in self:
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
finally:
from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt
if not closing_superseded_attempt():
await self._report_unreported_failure()
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)

View file

@ -1,6 +1,7 @@
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Union
from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Union
import httpx
from openai.types.file_deleted import FileDeleted
@ -38,6 +39,22 @@ else:
UserAPIKeyAuth = Any
def _routes_nowhere(_model_name: str) -> bool:
return False
@dataclass(frozen=True, slots=True)
class FileListModelScope:
model_mapping_keys: frozenset[str]
routes_stored_model_name: Callable[[str], bool] = _routes_nowhere
def holds(self, model_mapping_keys: Iterable[str], stored: bool) -> bool:
mapped_keys: Final = frozenset(model_mapping_keys)
return not self.model_mapping_keys.isdisjoint(mapped_keys) or (
stored and any(map(self.routes_stored_model_name, mapped_keys))
)
class BaseFileUploadStream(ABC):
"""Re-iterable request body that yields an upload's bytes lazily.
@ -274,6 +291,8 @@ class BaseFileEndpoints(ABC):
user_api_key_dict: UserAPIKeyAuth,
limit: int | None = None,
after: str | None = None,
model_scope: FileListModelScope | None = None,
before: str | None = None,
**data: dict,
) -> FileListPage:
pass

View file

@ -6,6 +6,7 @@ import ssl
import sys
import typing
import urllib.request
import weakref
from collections.abc import Callable, Generator
from typing import ClassVar, Final
@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
raise mapped_exc(message) from exc
def _release_transport_anchor(_transport: object) -> None:
"""Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None:
"""Keep the transport alive for as long as the stream can still read.
The transport finalizer closes the session it owns, and nothing on the path from a
live response back to that session references the transport, so a client dropped
mid-stream would close the session under the body. ``weakref.finalize`` holds the
transport in its own registry rather than on the stream, which keeps the transport
out of the response's reference cycle (an exception carrying the response, for one)
so it is finalized by refcount once the anchor drops and can still schedule an async
close, instead of being collected in the same cyclic pass as the session it closes.
"""
weakref.finalize(stream, _release_transport_anchor, transport)
class AiohttpResponseStream(httpx.AsyncByteStream):
CHUNK_SIZE = 1024 * 16
@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
# Re-raise if it's a different RuntimeError
raise
stream: Final = AiohttpResponseStream(response)
_anchor_transport_to(stream, self)
return httpx.Response(
status_code=response.status,
headers=response.headers,
stream=AiohttpResponseStream(response),
stream=stream,
request=request,
)

View file

@ -4,13 +4,14 @@ from typing import Final
from httpx._types import RequestFiles
from litellm.images.utils import image_to_data_url
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes
from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig, to_data_url
from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig
FLUX_LORA_DEPTH_ENDPOINT: Final[str] = "fal-ai/flux-lora-depth"
SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("n", "size")
@ -68,7 +69,7 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig):
)
request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict
"prompt": prompt,
"image_url": to_data_url(next(iter(images))),
"image_url": image_to_data_url(next(iter(images))),
**provider_params,
}
return request_body, ()

View file

@ -1,15 +1,12 @@
import base64
import os
from collections.abc import Mapping
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Final
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, TypeAdapter
from litellm.images.utils import ImageEditRequestUtils
from litellm.images.utils import image_to_data_url
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.fal_ai.image_generation.gpt_image_2_transformation import (
map_gpt_image_quality,
@ -39,39 +36,6 @@ PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType(
)
@runtime_checkable
class _SeekableBinaryReader(Protocol):
def tell(self) -> int: ...
def seek(self, offset: int) -> int: ...
def read(self) -> bytes: ...
def _read_image_bytes(image: object) -> bytes:
if isinstance(image, bytes):
return image
if isinstance(image, tuple):
return _read_image_bytes(image[1])
if isinstance(image, os.PathLike):
return Path(image).read_bytes()
if isinstance(image, _SeekableBinaryReader):
position: Final = image.tell()
image.seek(0)
data: Final = image.read()
image.seek(position)
return data
raise ValueError(f"Unsupported image type for Fal AI image edit: {type(image).__name__}")
def to_data_url(image: object) -> str:
if isinstance(image, str):
return image
image_bytes: Final = _read_image_bytes(image)
mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes)
return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"
def _first(value: object) -> object:
return value[0] if isinstance(value, list) and value else value
@ -146,14 +110,14 @@ class FalAIImageEditConfig(BaseImageEditConfig):
raise ValueError("Fal AI image edit requires at least one input image")
mask: Final = _first(image_edit_optional_request_params.get("mask"))
mask_field: Final[Mapping[str, str]] = (
MappingProxyType({"mask_url": to_data_url(mask)}) if mask is not None else MappingProxyType({})
MappingProxyType({"mask_url": image_to_data_url(mask)}) if mask is not None else MappingProxyType({})
)
provider_params: Final[Mapping[str, object]] = MappingProxyType(
{key: value for key, value in image_edit_optional_request_params.items() if key != "mask"}
)
request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict
"prompt": prompt,
"image_urls": tuple(to_data_url(img) for img in images),
"image_urls": tuple(image_to_data_url(img) for img in images),
**mask_field,
**provider_params,
}

View file

@ -1,6 +1,49 @@
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import Usage
from typing import Final
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_image_response_cost_from_usage,
flat_image_cost,
generic_cost_per_token,
resolve_image_model_info,
)
from litellm.types.utils import ImageResponse, ModelInfo, Usage
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="tencent")
def _deployment_image_token_pricing(model_info: ModelInfo | None) -> ModelInfo | None:
if model_info is None or model_info.get("output_cost_per_image_token") is not None:
return model_info
output_cost_per_token: Final = model_info.get("output_cost_per_token")
if output_cost_per_token is None:
return model_info
priced: Final[ModelInfo] = {**model_info, "output_cost_per_image_token": output_cost_per_token}
return priced
def _deployment_prices_per_image_only(model_info: ModelInfo | None) -> bool:
return (
model_info is not None
and model_info.get("output_cost_per_image") is not None
and model_info.get("output_cost_per_image_token") is None
and model_info.get("output_cost_per_token") is None
)
def image_cost_calculator(model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float:
resolved_model_info: Final = resolve_image_model_info(
model=model, custom_llm_provider="tencent", model_info=_deployment_image_token_pricing(model_info)
)
if _deployment_prices_per_image_only(model_info):
return flat_image_cost(resolved_model_info, image_response)
token_based_cost: Final = calculate_image_response_cost_from_usage(
model=model,
image_response=image_response,
custom_llm_provider="tencent",
model_info=resolved_model_info,
)
if token_based_cost is not None:
return token_based_cost
return flat_image_cost(resolved_model_info, image_response)

View file

@ -0,0 +1,147 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
import httpx
from httpx._types import RequestFiles
from pydantic import TypeAdapter
from typing_extensions import override
from litellm.images.utils import image_to_data_url
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.tencent.image_generation.transformation import (
TOKENHUB_IMAGE_PARAMS,
tokenhub_image_headers,
tokenhub_image_params,
tokenhub_image_response,
tokenhub_image_url,
)
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_IMAGES: Final = TypeAdapter(tuple[object, ...])
_FORM_BOOLEANS: Final = MappingProxyType({"true": True, "false": False})
_OPENAI_PARAMS: Final = frozenset(TOKENHUB_IMAGE_PARAMS)
_FormEntries: TypeAlias = tuple[tuple[tuple[str, ...], object], ...]
def _form_scalar(value: object) -> object:
match value:
case str() if value.lower() in _FORM_BOOLEANS:
return _FORM_BOOLEANS[value.lower()]
case str() if value.removeprefix("-").isdigit():
return int(value)
case _:
return value
def _form_path(key: str) -> tuple[str, ...]:
root, bracket, rest = key.partition("[")
if not bracket or not rest.endswith("]"):
return (key,)
return (root, *rest.removesuffix("]").split("]["))
def _nested_form_value(entries: _FormEntries) -> object:
match entries:
case (((), value),):
return _form_scalar(value)
case _:
return _nested_form_fields(tuple(entry for entry in entries if entry[0]))
def _nested_form_fields(entries: _FormEntries) -> Mapping[str, object]:
roots: Final = tuple(dict.fromkeys(path[0] for path, _ in entries))
return {
root: _nested_form_value(tuple((path[1:], value) for path, value in entries if path[0] == root))
for root in roots
}
def _tokenhub_form_params(params: Mapping[str, object]) -> Mapping[str, object]:
tokenhub_entries: Final = tuple(
(_form_path(key), value) for key, value in params.items() if key not in _OPENAI_PARAMS
)
return {
**{key: value for key, value in params.items() if key in _OPENAI_PARAMS},
**_nested_form_fields(tokenhub_entries),
}
class TencentImageEditConfig(BaseImageEditConfig):
@override
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: mirrors override contract
return list(TOKENHUB_IMAGE_PARAMS)
@override
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: mirrors override contract
return dict(tokenhub_image_params(image_edit_optional_params, model, drop_params))
@override
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None = None,
litellm_params: Mapping[str, object] | None = None,
api_base: str | None = None,
) -> dict[str, str]: # mutable-ok: mirrors override contract
return dict(tokenhub_image_headers(headers, api_key, model))
@override
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: Mapping[str, object],
) -> str:
return tokenhub_image_url(api_base)
@override
def use_multipart_form_data(self) -> bool:
return False
@override
def transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: object,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict[str, object], RequestFiles]: # mutable-ok: mirrors override contract
images: Final = _IMAGES.validate_python(image) if isinstance(image, list) else (image,)
return {
"model": model,
"prompt": prompt,
"images": tuple(image_to_data_url(img) for img in images if img is not None),
**_tokenhub_form_params(image_edit_optional_request_params),
}, ()
@override
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> ImageResponse:
image_response: Final = tokenhub_image_response(raw_response)
if image_response is None:
raise BaseLLMException(
status_code=raw_response.status_code,
message=raw_response.text,
headers=raw_response.headers,
)
return image_response

View file

@ -0,0 +1,200 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import override
from litellm.exceptions import AuthenticationError, UnsupportedParamsError
from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import (
ImageObject,
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
LlmProviders,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
TOKENHUB_API_BASE: Final = "https://tokenhub-intl.tencentcloudmaas.com/v1"
TOKENHUB_IMAGE_PATH: Final = "/wand/si-image/generation"
TOKENHUB_IMAGE_MODEL_PREFIX: Final = "seedream-"
TOKENHUB_IMAGE_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
"n",
"size",
"response_format",
"output_format",
"background",
"user",
)
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
class _TokenHubImage(BaseModel):
model_config = ConfigDict(frozen=True)
url: str | None = None
b64_json: str | None = None
size: str | None = None
output_format: str | None = None
class _TokenHubUsage(BaseModel):
model_config = ConfigDict(frozen=True)
total_tokens: int
class _TokenHubImageResponse(BaseModel):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
created: int
data: tuple[_TokenHubImage, ...]
tokenhub_usage: _TokenHubUsage | None = None
def is_tokenhub_image_model(model: str) -> bool:
return model.removeprefix(f"{LlmProviders.TENCENT.value}/").startswith(TOKENHUB_IMAGE_MODEL_PREFIX)
def tokenhub_image_url(api_base: str | None) -> str:
base: Final = api_base or get_secret_str("TENCENT_API_BASE") or TOKENHUB_API_BASE
return join_api_base(ensure_api_route(base, "v1"), TOKENHUB_IMAGE_PATH)
def tokenhub_image_headers(headers: Mapping[str, str], api_key: str | None, model: str) -> Mapping[str, str]:
resolved_api_key: Final = api_key or get_secret_str("TENCENT_API_KEY")
if not resolved_api_key:
raise AuthenticationError(
message="TENCENT_API_KEY is not set. Set it in the environment or pass api_key.",
llm_provider="tencent",
model=model,
)
return {**headers, "Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"}
def tokenhub_image_params(params: Mapping[str, object], model: str, drop_params: bool) -> Mapping[str, object]:
n: Final = params.get("n")
if n is not None and n != 1 and not drop_params:
raise UnsupportedParamsError(
status_code=400,
message=f"Tencent TokenHub returns one image per request, so n={n} is not supported for {model}. "
+ "To send n=1 instead, set `litellm.drop_params = True`.",
)
return {
key: value
for key, value in params.items()
if key != "n" and value is not None and not (key == "background" and value == "auto")
}
def tokenhub_image_response(raw_response: httpx.Response) -> ImageResponse | None:
try:
parsed: Final = _TokenHubImageResponse.model_validate_json(raw_response.content)
except ValidationError:
return None
total_tokens: Final = parsed.tokenhub_usage.total_tokens if parsed.tokenhub_usage is not None else 0
first_image: Final = parsed.data[0] if parsed.data else _TokenHubImage()
return ImageResponse(
created=parsed.created,
data=[ImageObject(url=image.url, b64_json=image.b64_json) for image in parsed.data],
size=first_image.size,
output_format=first_image.output_format,
usage=ImageUsage(
input_tokens=0,
input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0),
output_tokens=total_tokens,
total_tokens=total_tokens,
),
)
class TencentImageGenerationConfig(BaseImageGenerationConfig):
@override
def get_supported_openai_params(
self, model: str
) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: mirrors override contract
return list(TOKENHUB_IMAGE_PARAMS)
@override
def map_openai_params(
self,
non_default_params: Mapping[str, object],
optional_params: Mapping[str, object],
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: mirrors override contract
return {**optional_params, **tokenhub_image_params(non_default_params, model, drop_params)}
@override
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
stream: bool | None = None,
) -> str:
return tokenhub_image_url(api_base)
@override
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, str]: # mutable-ok: mirrors override contract
return dict(tokenhub_image_headers(headers, api_key, model))
@override
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
headers: Mapping[str, str],
) -> dict[str, object]: # mutable-ok: mirrors override contract
extra_body: Final = _JSON_OBJECT.validate_python(optional_params.get("extra_body") or {})
params: Final = {key: value for key, value in optional_params.items() if key != "extra_body"}
return {"model": model, "prompt": prompt, **params, **extra_body}
@override
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: Mapping[str, object],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
image_response: Final = tokenhub_image_response(raw_response)
if image_response is None:
raise BaseLLMException(
status_code=raw_response.status_code,
message=raw_response.text,
headers=raw_response.headers,
)
return image_response

View file

@ -2743,10 +2743,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
chunk_indices = support.get("groundingChunkIndices", [])
if start_index is not None and end_index is not None and chunk_indices:
# Use the first chunk's URL for the annotation
first_chunk_idx = chunk_indices[0]
if first_chunk_idx in chunk_to_uri_map:
uri_info = chunk_to_uri_map[first_chunk_idx]
for chunk_index in chunk_indices:
if chunk_index not in chunk_to_uri_map:
continue
uri_info = chunk_to_uri_map[chunk_index]
url_citation: ChatCompletionAnnotationURLCitation = {
"start_index": start_index,

View file

@ -21589,7 +21589,13 @@
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
"supports_video_input": true,
"input_cost_per_token_priority": 5.6252e-07,
"output_cost_per_token_priority": 3.37498e-06,
"cache_read_input_token_cost_priority": 5.628e-08,
"cache_creation_input_token_cost_priority": 5.6252e-07,
"input_cost_per_audio_token_priority": 1.12497e-06,
"cache_read_input_audio_token_cost_priority": 1.1249e-07
},
"databricks/databricks-gemini-3-1-flash-image": {
"cache_creation_input_token_cost": 6.2503e-07,
@ -21676,7 +21682,14 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.50002e-06,
"output_cost_per_token_priority": 2.699998e-05,
"cache_read_input_token_cost_priority": 4.5003e-07,
"cache_creation_input_token_cost_priority": 4.50002e-06,
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
},
"databricks/databricks-gemini-3-flash": {
"cache_creation_input_token_cost": 6.2503e-07,
@ -21710,7 +21723,13 @@
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
"supports_video_input": true,
"input_cost_per_token_priority": 1.12497e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.1249e-07,
"cache_creation_input_token_cost_priority": 1.12497e-06,
"input_cost_per_audio_token_priority": 2.25001e-06,
"cache_read_input_audio_token_cost_priority": 2.2498e-07
},
"databricks/databricks-gemini-3-pro": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -21757,7 +21776,11 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-7-flash": {
"cache_creation_input_token_cost": 7.4998e-07,
@ -21785,10 +21808,14 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-6-flash": {
"cache_creation_input_token_cost": 1.87502e-06,
"cache_creation_input_token_cost": 7.4998e-07,
"cache_read_input_token_cost": 7.497e-08,
"input_cost_per_token": 7.4998e-07,
"input_dbu_cost_per_token": 1.0714e-05,
@ -21816,7 +21843,11 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-5-flash": {
"cache_creation_input_token_cost": 1.87502e-06,
@ -21846,7 +21877,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 3.37498e-06,
"output_cost_per_token_priority": 2.025002e-05,
"cache_read_input_token_cost_priority": 3.3747e-07,
"cache_creation_input_token_cost_priority": 3.37498e-06
},
"databricks/databricks-gemini-3-5-flash-lite": {
"cache_creation_input_token_cost": 3.7499e-07,
@ -21878,7 +21913,11 @@
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 6.7501e-07,
"output_cost_per_token_priority": 5.62499e-06,
"cache_read_input_token_cost_priority": 6.748e-08,
"cache_creation_input_token_cost_priority": 6.7501e-07
},
"databricks/databricks-gemma-3-12b": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -21921,7 +21960,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
"supports_vision": false,
"input_cost_per_token_priority": 2.45e-06,
"output_cost_per_token_priority": 7.7e-06,
"cache_read_input_token_cost_priority": 4.55e-07,
"cache_creation_input_token_cost_priority": 2.45e-06
},
"databricks/databricks-glm-5-3": {
"cache_creation_input_token_cost": 1.4e-06,
@ -21951,7 +21994,11 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 2.45e-06,
"output_cost_per_token_priority": 7.7e-06,
"cache_read_input_token_cost_priority": 4.55e-07,
"cache_creation_input_token_cost_priority": 2.45e-06
},
"databricks/databricks-glm-5-3-flash": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -21980,7 +22027,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 2.625e-07,
"output_cost_per_token_priority": 8.75e-07,
"cache_read_input_token_cost_priority": 5.25e-08,
"cache_creation_input_token_cost_priority": 2.625e-07
},
"databricks/databricks-gpt-5": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -21999,7 +22050,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 2.49998e-06,
"output_cost_per_token_priority": 1.999998e-05,
"cache_read_input_token_cost_priority": 2.4997e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06
},
"databricks/databricks-gpt-5-1": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -22018,7 +22073,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 2.49998e-06,
"output_cost_per_token_priority": 1.999998e-05,
"cache_read_input_token_cost_priority": 2.4997e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06
},
"databricks/databricks-gpt-5-2": {
"cache_creation_input_token_cost": 1.75e-06,
@ -22037,7 +22096,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 3.5e-06,
"output_cost_per_token_priority": 2.8e-05,
"cache_read_input_token_cost_priority": 3.5e-07,
"cache_creation_input_token_cost_priority": 3.5e-06
},
"databricks/databricks-gpt-5-3-codex": {
"cache_creation_input_token_cost": 1.75e-06,
@ -22056,7 +22119,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 3.5e-06,
"output_cost_per_token_priority": 2.8e-05,
"cache_read_input_token_cost_priority": 3.5e-07,
"cache_creation_input_token_cost_priority": 3.5e-06
},
"databricks/databricks-gpt-5-4": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -22087,7 +22154,14 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.00003e-06,
"output_cost_per_token_priority": 2.999997e-05,
"cache_read_input_token_cost_priority": 5.0001e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
},
"databricks/databricks-gpt-5-4-mini": {
"cache_creation_input_token_cost": 7.4998e-07,
@ -22115,7 +22189,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.50003e-06,
"output_cost_per_token_priority": 8.99997e-06,
"cache_read_input_token_cost_priority": 1.5001e-07,
"cache_creation_input_token_cost_priority": 1.50003e-06
},
"databricks/databricks-gpt-5-4-nano": {
"cache_creation_input_token_cost": 1.9999e-07,
@ -22175,7 +22253,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 8.00002e-06,
"output_cost_per_token_priority": 4.000003e-05,
"cache_read_input_token_cost_priority": 8.0003e-07,
"cache_creation_input_token_cost_priority": 9.99999e-06,
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
},
"databricks/databricks-gpt-5-6-terra": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -22207,7 +22293,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 2.399999e-05,
"cache_read_input_token_cost_priority": 3.9998e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-5-6-luna": {
"cache_creation_input_token_cost": 2.4997e-07,
@ -22239,7 +22333,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 3.9998e-07,
"output_cost_per_token_priority": 2.40002e-06,
"cache_read_input_token_cost_priority": 3.997e-08,
"cache_creation_input_token_cost_priority": 5.0001e-07,
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
},
"databricks/databricks-gpt-5-5": {
"cache_creation_input_token_cost": 5.00003e-06,
@ -22270,7 +22372,14 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.249997e-05,
"output_cost_per_token_priority": 7.50001e-05,
"cache_read_input_token_cost_priority": 1.24999e-06,
"cache_creation_input_token_cost_priority": 1.249997e-05,
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
},
"databricks/databricks-gpt-5-5-pro": {
"cache_creation_input_token_cost": 2.999997e-05,
@ -22319,7 +22428,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 4.5003e-07,
"output_cost_per_token_priority": 3.60003e-06,
"cache_read_input_token_cost_priority": 4.501e-08,
"cache_creation_input_token_cost_priority": 4.5003e-07
},
"databricks/databricks-gpt-5-nano": {
"cache_creation_input_token_cost": 4.998e-08,
@ -22469,7 +22582,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.25e-06,
"output_cost_per_token_priority": 2.625e-05,
"cache_read_input_token_cost_priority": 5.25e-07,
"cache_creation_input_token_cost_priority": 5.25e-06
},
"databricks/system.ai.kimi-k3": {
"litellm_provider": "databricks",
@ -22551,7 +22668,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 3.85e-07,
"output_cost_per_token_priority": 3.85e-06,
"cache_read_input_token_cost_priority": 3.85e-07,
"cache_creation_input_token_cost_priority": 3.85e-07
},
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -61865,6 +61986,39 @@
"supports_reasoning": true,
"supports_vision": false
},
"tencent/seedream-image-v5.0-flash": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.018,
"output_cost_per_image_token": 1.6e-06,
"source": "https://cloud.tencent.com/document/product/1823/130055",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"tencent/seedream-image-v5.0-lite": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.035,
"output_cost_per_image_token": 1.6e-06,
"source": "https://www.tencentcloud.com/document/product/1300/78937",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"tencent/seedream-image-v5.0-pro": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.09,
"output_cost_per_image_token": 1.6e-06,
"source": "https://www.tencentcloud.com/document/product/1300/78937",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"cognition/swe-1.6": {
"input_cost_per_token": 5e-07,
"output_cost_per_token": 2.5e-06,
@ -77774,7 +77928,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.25e-07,
"output_cost_per_token_priority": 2.1e-06,
"cache_read_input_token_cost_priority": 5.25e-08,
"cache_creation_input_token_cost_priority": 5.25e-07
},
"databricks/databricks-deepseek-v4-1-flash-1017": {
"cache_creation_input_token_cost": 3.0002e-07,
@ -77897,7 +78055,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.50002e-06,
"output_cost_per_token_priority": 2.250003e-05,
"cache_read_input_token_cost_priority": 2.2498e-07,
"cache_creation_input_token_cost_priority": 4.50002e-06
},
"databricks/databricks-gpt-5-4-pro": {
"cache_creation_input_token_cost": 2.999997e-05,
@ -77959,7 +78121,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 1.99997e-05,
"cache_read_input_token_cost_priority": 1.9999e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-6-1-terra": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -77991,7 +78161,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.99997e-06,
"output_cost_per_token_priority": 9.99999e-06,
"cache_read_input_token_cost_priority": 1.0003e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06,
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
},
"databricks/databricks-gpt-6-astra": {
"cache_creation_input_token_cost": 1.249997e-05,
@ -78023,7 +78201,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.999998e-05,
"output_cost_per_token_priority": 9.99999e-05,
"cache_read_input_token_cost_priority": 1.99997e-06,
"cache_creation_input_token_cost_priority": 2.500001e-05,
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
},
"databricks/databricks-gpt-6-luna": {
"cache_creation_input_token_cost": 1.2502e-07,
@ -78055,7 +78241,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.9999e-07,
"output_cost_per_token_priority": 1.0003e-06,
"cache_read_input_token_cost_priority": 2.002e-08,
"cache_creation_input_token_cost_priority": 2.4997e-07,
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
},
"databricks/databricks-gpt-6-sol": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -78087,7 +78281,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 1.99997e-05,
"cache_read_input_token_cost_priority": 3.9998e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-image-1": {
"cache_creation_input_token_cost": 5.00003e-06,

View file

@ -3793,8 +3793,8 @@
"responses": true,
"text_completion": true,
"embeddings": false,
"image_generations": false,
"image_edits": false,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": false,
"audio_speech": false,
"rerank": false,

View file

@ -3487,6 +3487,13 @@
"title": "Entities",
"type": "object"
},
"guardrails": {
"additionalProperties": {
"$ref": "#/components/schemas/GuardrailMetric"
},
"title": "Guardrails",
"type": "object"
},
"mcp_servers": {
"additionalProperties": {
"$ref": "#/components/schemas/MetricWithMetadata"
@ -3802,6 +3809,74 @@
"title": "ExportType",
"type": "string"
},
"GuardrailKeyMetric": {
"properties": {
"metadata": {
"$ref": "#/components/schemas/KeyMetadata"
},
"metrics": {
"$ref": "#/components/schemas/GuardrailSpendMetrics"
}
},
"required": [
"metrics",
"metadata"
],
"title": "GuardrailKeyMetric",
"type": "object"
},
"GuardrailMetadata": {
"properties": {
"guardrail_name": {
"title": "Guardrail Name",
"type": "string"
}
},
"required": [
"guardrail_name"
],
"title": "GuardrailMetadata",
"type": "object"
},
"GuardrailMetric": {
"properties": {
"api_key_breakdown": {
"additionalProperties": {
"$ref": "#/components/schemas/GuardrailKeyMetric"
},
"title": "Api Key Breakdown",
"type": "object"
},
"metadata": {
"$ref": "#/components/schemas/GuardrailMetadata"
},
"metrics": {
"$ref": "#/components/schemas/GuardrailSpendMetrics"
}
},
"required": [
"metrics",
"metadata"
],
"title": "GuardrailMetric",
"type": "object"
},
"GuardrailSpendMetrics": {
"properties": {
"spend": {
"default": 0.0,
"title": "Spend",
"type": "number"
},
"untracked_units": {
"default": 0,
"title": "Untracked Units",
"type": "integer"
}
},
"title": "GuardrailSpendMetrics",
"type": "object"
},
"HTTPAuthSecurityScheme": {
"description": "Defines a security scheme using HTTP authentication.",
"properties": {

View file

@ -370,7 +370,7 @@ class _KeyModelBudgetLimiter(Protocol):
class _UserModelBudgetLimiter(Protocol):
async def is_user_within_model_budget(
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
) -> bool: ...
@ -388,11 +388,15 @@ class _EveryModelBudgetLimiter(Protocol):
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
async def is_user_within_model_budget(
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
) -> bool: ...
async def is_end_user_within_model_budget(
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
self,
end_user_id: str,
end_user_model_max_budget: Mapping[str, object],
model: str,
team_id: str | None = None,
) -> bool: ...
async def is_team_within_model_budget(
@ -503,6 +507,7 @@ async def _check_user_model_budget(
user_id=valid_token.user_id,
user_model_max_budget=user_model_max_budget,
model=model_name,
team_id=valid_token.team_id,
)
@ -547,6 +552,7 @@ async def _is_within_every_model_budget(
user_id=valid_token.user_id,
user_model_max_budget=user_model_max_budget,
model=model,
team_id=valid_token.team_id,
)
if (
valid_token.end_user_id is not None
@ -557,6 +563,7 @@ async def _is_within_every_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_model_max_budget,
model=model,
team_id=valid_token.team_id,
)
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
await model_max_budget_limiter.is_team_within_model_budget(
@ -3154,6 +3161,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=model_name,
team_id=valid_token.team_id,
)
# Check 6: Additional Common Checks across jwt + key auth
@ -4733,6 +4741,7 @@ async def _run_post_custom_auth_checks(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=model_name,
team_id=valid_token.team_id,
)
# team / user / end_user / project context objects are fetched by

View file

@ -4,6 +4,7 @@ import json
import logging
import math
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from functools import lru_cache
from itertools import chain
@ -136,7 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
_check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
check_and_merge_model_level_guardrails,
)
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
@ -265,6 +266,8 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
AnthropicStreamModelRestamper,
)
from litellm.proxy.litellm_pre_call_utils import (
ROUTER_SCHEDULED_ROUTE_TYPES,
SCHEDULER_PRIORITY_HEADER,
LiteLLMProxyRequestSetup,
add_litellm_data_to_request,
get_metadata_variable_name,
@ -2006,6 +2009,49 @@ def _content_text(fields: Mapping[str, object]) -> str:
return str(fields.get("content", ""))
async def _no_stream_failure_to_release() -> None:
return None
@dataclass(slots=True)
class _ClientStreamFailureReport:
logging_obj: LiteLLMLoggingObj
report: Callable[[Exception], Awaitable[None]]
held: Exception | None = None
@classmethod
def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]:
if logging_obj is None or logging_obj.on_detached_stream_failure is None:
return _no_stream_failure_to_release
return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding()
def _start_holding(self) -> Callable[[], Awaitable[None]]:
self.logging_obj.on_detached_stream_failure = self.hold
return self.release
async def hold(self, failure: Exception) -> None:
if self.held is None:
self.held = failure
async def release(self) -> None:
self.logging_obj.on_detached_stream_failure = self.report
if self.held is None:
return
await asyncio.shield(
spawn_logging_task(
asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report"
)
)
async def _report_held(self, failure: Exception) -> None:
try:
await self.report(failure)
except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup
verbose_proxy_logger.warning(
"Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure)
)
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict, served_model: str | None = None):
self.data = data
@ -2434,14 +2480,20 @@ class ProxyBaseLLMRequestProcessing:
self.data["model"] = alias_target
aliased_model: Final = self.data.get("model")
if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str):
self.data.update(
llm_router.priority_header_update(
model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
)
)
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
self.data["litellm_call_id"] = litellm_call_id
bind_telemetry_call_id(litellm_call_id)
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
DDSpanTagger.tag_call_id(litellm_call_id)
DDSpanTagger.tag_request(
user_api_key_dict=user_api_key_dict,
requested_model=self.data.get("model"),
requested_model=aliased_model,
)
### AUTO STREAM USAGE TRACKING ###
@ -3498,6 +3550,8 @@ class ProxyBaseLLMRequestProcessing:
request_data: Final = self.data
async def _on_detached_stream_failure(exc: Exception) -> None:
if not logging_obj.claim_stream_failure_report():
return
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
@ -4448,11 +4502,15 @@ class ProxyBaseLLMRequestProcessing:
client_disconnected = False
delivered_chunk = False
recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes
stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below
guarded_stream: Final[AsyncGenerator[object, None]] = proxy_logging_obj.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
)
release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming(
stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None
)
try:
str_so_far = ""
async for chunk in guarded_stream:
@ -4523,10 +4581,18 @@ class ProxyBaseLLMRequestProcessing:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e
)
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=request_data,
already_reported: Final = (
isinstance(stream_logging_obj, LiteLLMLoggingObj)
and not stream_logging_obj.claim_stream_failure_report()
)
transformed_exception: Final = (
None
if already_reported
else await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=request_data,
)
)
if transformed_exception is not None:
e = transformed_exception
@ -4549,15 +4615,18 @@ class ProxyBaseLLMRequestProcessing:
seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail)
yield seal + error_frame if seal else error_frame
finally:
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
request=request,
request_data=request_data,
response=response,
stream_completed=stream_completed,
client_disconnected=client_disconnected,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
try:
await release_stream_failure()
finally:
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
request=request,
request_data=request_data,
response=response,
stream_completed=stream_completed,
client_disconnected=client_disconnected,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod
def async_sse_data_generator(

View file

@ -64,6 +64,16 @@ _ROLL_WINDOW_SPEND_SQL: Final = (
"AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')"
)
_OPEN_WINDOW_SPEND_SQL: Final = (
'INSERT INTO "LiteLLM_BudgetWindowSpend" '
"(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) "
"VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), 0, "
"($4::timestamptz AT TIME ZONE 'UTC'), ($4::timestamptz AT TIME ZONE 'UTC')) "
"ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET "
"window_start = EXCLUDED.window_start, spend = 0, updated_at = EXCLUDED.updated_at "
"WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')"
)
_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = (
"SELECT COALESCE(SUM(spend), 0.0) AS total, "
"COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch "
@ -389,3 +399,29 @@ async def roll_window_spend_row(
to_naive_utc(new_window_start),
to_naive_utc(datetime.now(timezone.utc)),
)
async def open_window_spend_row(
prisma_client: "PrismaClient",
entity_type: str,
entity_id: str,
window_duration: str,
window_start: datetime,
replaces_window_start: datetime | None,
) -> None:
"""Start a window at `window_start` unless another writer already started one.
The row is created when it does not exist, and moved onto the new window only
while it still holds `replaces_window_start`, the ended window the caller saw.
A writer that lost the race leaves the winner's window in place, so reading
the row back gives every writer the same window start.
"""
async with db_span("open_window_spend_row", "LiteLLM_BudgetWindowSpend"):
await prisma_client.db.execute_raw(
_OPEN_WINDOW_SPEND_SQL,
entity_type,
entity_id,
window_duration,
to_naive_utc(window_start),
None if replaces_window_start is None else to_naive_utc(replaces_window_start),
)

View file

@ -305,6 +305,11 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
yield from _iter_text_parts_in_content(message.get("content"))
def _is_tool_message(message: object) -> bool:
fields: Final = _text_object(message)
return fields is not None and str(fields.get("role") or "").lower() == "tool"
def _call_item(item: object) -> tuple[str, object] | None:
"""Return ``(field, value)`` for a Responses tool-call input item, else None."""
fields: Final = _text_object(item)
@ -316,11 +321,14 @@ def walk_user_text(
data: dict[str, object],
visit: Callable[[str], str],
visit_tool: Callable[[JsonValue], JsonValue] | None = None,
*,
skip_tool_message: bool = False,
) -> int:
"""Rewrite every text fragment in place via ``visit``.
Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and
``data["instructions"]``. Returns the number
``data["instructions"]``. ``skip_tool_message`` leaves ``role: tool`` messages and ``tool_result``
blocks as sent. Returns the number
of fragments visited so callers can short-circuit when nothing was
inspected.
"""
@ -353,6 +361,8 @@ def walk_user_text(
def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]:
nonlocal visited
if skip_tool_message and part.get("type") == "tool_result":
return part
text: Final = _part_text(part)
rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})}
if text is not None:
@ -384,18 +394,17 @@ def walk_user_text(
return content
messages: Final = data.get("messages")
if isinstance(messages, list):
for message in messages:
if not isinstance(message, dict):
continue
if "content" in message:
message["content"] = _rewrite_content(message["content"])
calls = message.get("tool_calls")
if isinstance(calls, list):
for call in calls:
function = call.get("function") if isinstance(call, dict) else None
if isinstance(function, dict) and "arguments" in function:
function["arguments"] = _rewrite_arguments(function["arguments"])
for message in _object_list(messages) or ():
if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict):
continue
if "content" in message:
message["content"] = _rewrite_content(message["content"])
calls = message.get("tool_calls")
if isinstance(calls, list):
for call in calls:
function = call.get("function") if isinstance(call, dict) else None
if isinstance(function, dict) and "arguments" in function:
function["arguments"] = _rewrite_arguments(function["arguments"])
if "system" in data:
data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract

View file

@ -0,0 +1,176 @@
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.db.budget_window_spend_writer import open_window_spend_row
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendTransaction,
build_window_spend_transaction,
)
from litellm.repositories.table_repositories import BudgetWindowSpendRepository
from litellm.router_strategy.budget_limiter import budget_window_seconds
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
MODEL_BUDGET_WINDOW_ENTITY_TYPE: Final = "model_max_budget"
_READ_TTL_SECONDS: Final = 5
@dataclass(frozen=True, slots=True)
class SharedWindow:
started_at: float
spend: float
def is_open(self, budget_duration: str, now: float) -> bool:
return now < self.started_at + budget_window_seconds(self.started_at, budget_duration)
class LocalTTLCache(Protocol):
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: object, *, ttl: float) -> None: ...
def delete_cache(self, key: str) -> None: ...
class SharedWindows(Protocol):
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None: ...
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: ...
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: ...
PrismaClientSource: TypeAlias = Callable[[], "PrismaClient | None"]
WindowSpendSink: TypeAlias = Callable[[WindowSpendTransaction], Awaitable[None]]
def _proxy_prisma_client() -> "PrismaClient | None":
try:
from litellm.proxy.proxy_server import prisma_client
except ImportError:
return None
return prisma_client
async def _queue_on_proxy_spend_writer(transaction: WindowSpendTransaction) -> None:
from litellm.proxy.proxy_server import proxy_logging_obj
await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(transaction)
@dataclass(frozen=True, slots=True)
class _Reading:
window: SharedWindow | None
class DatabaseSharedWindows:
def __init__(
self,
prisma_client_source: PrismaClientSource = _proxy_prisma_client,
sink: WindowSpendSink = _queue_on_proxy_spend_writer,
read_ttl_seconds: int = _READ_TTL_SECONDS,
) -> None:
self._prisma_client_source = prisma_client_source
self._sink = sink
self._read_ttl_seconds = read_ttl_seconds
self._readings: LocalTTLCache = InMemoryCache(max_size_in_memory=10_000, default_ttl=read_ttl_seconds)
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
cached: Final = self._readings.get_cache(key=spend_key)
if isinstance(cached, _Reading):
return cached.window
window: Final = await self._read_row(spend_key=spend_key, budget_duration=budget_duration)
self._readings.set_cache(key=spend_key, value=_Reading(window=window), ttl=self._read_ttl_seconds)
return window
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None:
"""When the window every replica shares started, opening one at `now` if none is open.
Without a database every window is this replica's own, so it starts at `now`. None means the database did
not answer, so this replica cannot know which window the others are in yet.
"""
prisma_client: Final = self._prisma_client_source()
if prisma_client is None:
return now
try:
seen: Final = await _fetch_row(
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
)
if seen is not None and seen.is_open(budget_duration=budget_duration, now=now):
return seen.started_at
await open_window_spend_row(
prisma_client=prisma_client,
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
entity_id=spend_key,
window_duration=budget_duration,
window_start=datetime.fromtimestamp(now, tz=timezone.utc),
replaces_window_start=None
if seen is None
else datetime.fromtimestamp(seen.started_at, tz=timezone.utc),
)
opened: Final = await _fetch_row(
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
)
except Exception: # noqa: BLE001 # an unreachable database leaves this replica on its own window
verbose_proxy_logger.exception("Unable to open the shared model budget window")
return None
if opened is None or not opened.is_open(budget_duration=budget_duration, now=now):
return None
return opened.started_at
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None:
if self._prisma_client_source() is None:
return
try:
await self._sink(
build_window_spend_transaction(
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
entity_id=spend_key,
window_duration=budget_duration,
window_start=datetime.fromtimestamp(started_at, tz=timezone.utc),
spend=spend,
)
)
except Exception: # noqa: BLE001 # spend tracking must never fail the success callback
verbose_proxy_logger.debug("Unable to queue shared model budget spend", exc_info=True)
async def _read_row(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
prisma_client: Final = self._prisma_client_source()
if prisma_client is None:
return None
try:
return await _fetch_row(
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=False
)
except Exception: # noqa: BLE001 # an unreadable row leaves enforcement on this replica's own counter
verbose_proxy_logger.exception("Unable to read shared model budget spend")
return None
async def _fetch_row(
prisma_client: "PrismaClient", spend_key: str, budget_duration: str, use_writer: bool
) -> SharedWindow | None:
row: Final = await BudgetWindowSpendRepository(prisma_client, use_writer=use_writer).table.find_unique(
where={
"entity_type_entity_id_window_duration": {
"entity_type": MODEL_BUDGET_WINDOW_ENTITY_TYPE,
"entity_id": spend_key,
"window_duration": budget_duration,
}
}
)
if row is None:
return None
window_start: Final = row.window_start
started_at: Final = (
window_start if window_start.tzinfo is not None else window_start.replace(tzinfo=timezone.utc)
).timestamp()
return SharedWindow(started_at=started_at, spend=float(row.spend or 0.0))
PROXY_SHARED_WINDOWS: Final = DatabaseSharedWindows()

File diff suppressed because it is too large Load diff

View file

@ -118,6 +118,9 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
_SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$")
SCHEDULER_PRIORITY_HEADER: Final = "x-litellm-priority"
ROUTER_SCHEDULED_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion"})
# W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/
# e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
_TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE)

View file

@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass, replace
from datetime import date, datetime, timedelta
from itertools import groupby
from types import MappingProxyType
from typing import Final, Literal, NoReturn, Protocol
@ -28,6 +29,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import (
DailySpendData,
DailySpendMetadata,
GroupedData,
GuardrailKeyMetric,
GuardrailMetadata,
GuardrailMetric,
GuardrailSpendMetrics,
KeyMetadata,
KeyMetricWithMetadata,
MetricWithMetadata,
@ -39,6 +44,7 @@ from litellm.types.repositories.daily_activity import (
DailyActivityTable,
EntityRollupRow,
GroupingSetsRow,
GuardrailSpendRow,
KeyMetadataRow,
RollupMetricsRow,
SpendLogsWindow,
@ -1013,17 +1019,68 @@ def _fold_entity_rollups_sync(
)
def _guardrail_metric(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> GuardrailMetric:
return GuardrailMetric(
metrics=GuardrailSpendMetrics(
spend=sum(row.spend for row in rows),
untracked_units=sum(row.untracked_units for row in rows),
),
metadata=GuardrailMetadata(guardrail_name=rows[0].guardrail_name),
api_key_breakdown={
row.api_key: GuardrailKeyMetric(
metrics=GuardrailSpendMetrics(spend=row.spend, untracked_units=row.untracked_units),
metadata=day.breakdown.api_keys[row.api_key].metadata,
)
for row in rows
if row.api_key in day.breakdown.api_keys
},
)
def _day_with_guardrails(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> DailySpendData:
by_guardrail: Final = groupby(sorted(rows, key=lambda row: row.guardrail_id), key=lambda row: row.guardrail_id)
guardrails: Final = {guardrail_id: _guardrail_metric(day, tuple(group)) for guardrail_id, group in by_guardrail}
return day.model_copy(update={"breakdown": day.breakdown.model_copy(update={"guardrails": guardrails})})
def _with_guardrail_breakdown(
results: Sequence[DailySpendData], rows: Sequence[GuardrailSpendRow]
) -> tuple[DailySpendData, ...]:
rows_by_date: Final = MappingProxyType(
{day: tuple(group) for day, group in groupby(sorted(rows, key=lambda row: row.date), key=lambda row: row.date)}
)
return tuple(_day_with_guardrails(day, rows_by_date.get(day.date.strftime("%Y-%m-%d"), ())) for day in results)
async def _guardrail_spend_or_empty(
repository: DailyActivityRepository, scope: DailyActivityScope
) -> tuple[GuardrailSpendRow, ...]:
from prisma.errors import PrismaError
try:
return await repository.guardrail_spend(scope)
except PrismaError as exc:
verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc)
return ()
async def get_daily_activity_aggregated(
repository: DailyActivityRepository,
scope: DailyActivityScope,
*,
entity_metadata_field: Mapping[str, dict[str, object]] | None = None,
include_entity_breakdown: bool = False,
include_guardrail_breakdown: bool = False,
api_key_limit: int = constants.USAGE_TOP_API_KEYS_DEFAULT,
) -> SpendAnalyticsPaginatedResponse:
try:
aggregated_rows: Final = await repository.aggregated(
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
aggregated_rows, guardrail_rows = await asyncio.gather(
repository.aggregated(
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
),
_guardrail_spend_or_empty(repository, scope)
if include_guardrail_breakdown and not scope.model
else asyncio.sleep(0, result=()),
)
records: Final = aggregated_rows.grouping_rows
aggregated: Final = await _aggregate_grouping_sets_records(
@ -1058,8 +1115,13 @@ async def get_daily_activity_aggregated(
api_key_metadata=entity_key_metadata,
entity_metadata_field=entity_metadata_field,
)
results: Final = (
list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows))
if guardrail_rows
else aggregated["results"]
)
return SpendAnalyticsPaginatedResponse(
results=aggregated["results"],
results=results,
metadata=DailySpendMetadata(
total_spend=aggregated["totals"].spend,
total_flat_cost=aggregated["totals"].flat_cost,

View file

@ -300,6 +300,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
resolved.scope,
entity_metadata_field=resolved.entity_metadata,
include_entity_breakdown=resolver.include_entity_breakdown,
include_guardrail_breakdown=resolver.include_guardrail_breakdown,
api_key_limit=api_key_limit,
)
if resolver.entity != "team":

View file

@ -56,6 +56,7 @@ class EntityScopeResolver:
resolve: EntityScopeResolve
include_entity_breakdown: bool
operation_names: OperationNames
include_guardrail_breakdown: bool = False
def _query_ids(value: str | None) -> tuple[str, ...] | None:
@ -364,6 +365,7 @@ USER_RESOLVER = EntityScopeResolver(
query=_user_query,
resolve=_resolve_user,
include_entity_breakdown=False,
include_guardrail_breakdown=True,
operation_names=MappingProxyType(
{
"aggregated": "get_user_daily_activity_aggregated",
@ -385,6 +387,7 @@ TEAM_RESOLVER = EntityScopeResolver(
query=_team_query,
resolve=_resolve_team,
include_entity_breakdown=True,
include_guardrail_breakdown=True,
operation_names=MappingProxyType(
{
"aggregated": "get_team_daily_activity_aggregated",
@ -423,6 +426,7 @@ ORGANIZATION_RESOLVER = EntityScopeResolver(
query=_organization_query,
resolve=_resolve_organization,
include_entity_breakdown=True,
include_guardrail_breakdown=True,
operation_names=MappingProxyType(
{
"aggregated": "get_organization_daily_activity_aggregated",

View file

@ -499,7 +499,7 @@ async def new_user(
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
@ -1702,7 +1702,7 @@ async def user_update(
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.

View file

@ -2328,7 +2328,7 @@ async def generate_key_fn(
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
@ -2538,7 +2538,7 @@ async def generate_service_account_key_fn(
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
- guardrails: Optional[List[str]] - List of active guardrails for the key
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
@ -3855,7 +3855,7 @@ async def update_key_fn(
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
- spend: Optional[float] - Amount spent by key
- max_budget: Optional[float] - Max budget for key
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
@ -6303,7 +6303,7 @@ async def regenerate_key_fn(
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
- spend: Optional[float] - Amount spent by key
- max_budget: Optional[float] - Max budget for key
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.

View file

@ -11,6 +11,7 @@ All /team management endpoints
import asyncio
import copy
import itertools
import json
import math
import traceback
@ -139,7 +140,12 @@ from litellm.proxy.common_utils.user_api_key_cache import (
from litellm.proxy.db.db_span import db_span
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import (
RoutedModels,
build_model_max_budget_usage,
entries_outranked_by,
entry_names_model,
is_enforced_budget_entry,
proxy_routed_models,
resolve_model_budget,
)
from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied
@ -1467,8 +1473,10 @@ def _check_team_model_budget_update_authority(
data: UpdateTeamRequest,
user_api_key_dict: UserAPIKeyAuth,
existing_model_max_budget: Mapping[str, object] | None,
routed_models_source: Callable[[str | None], RoutedModels] = proxy_routed_models,
) -> None:
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap."""
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap, or
add an entry that takes precedence over one."""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget:
@ -1487,29 +1495,97 @@ def _check_team_model_budget_update_authority(
)
},
)
for model_name, proposed in requested.items():
governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget)
if governing is None:
continue
cap = governing.budget_config
if cap.max_budget is None:
continue
if (
proposed.max_budget is None
or proposed.max_budget > cap.max_budget
or proposed.budget_duration != cap.budget_duration
):
raise HTTPException(
status_code=403,
detail={
"error": (
f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its "
f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} "
f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per "
f"{proposed.budget_duration}."
)
},
routed_models: Final = routed_models_source(data.team_id)
for entry, proposed in requested.items():
existing = _existing_model_cap(existing_model_max_budget.get(entry))
if existing is not None:
_check_existing_model_cap_kept(entry=entry, cap_entry=entry, proposed=proposed, existing=existing)
if is_enforced_budget_entry(existing) or not is_enforced_budget_entry(proposed):
continue
governing = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=data.team_id)
if governing is not None and entry_names_model(entry=governing.budget_model, model=entry):
_check_existing_model_cap_kept(
entry=entry, cap_entry=governing.budget_model, proposed=proposed, existing=governing.budget_config
)
_check_new_entry_outranks_no_cap(
entry=entry,
existing_model_max_budget=existing_model_max_budget,
routed_models=routed_models,
team_id=data.team_id,
)
def _check_existing_model_cap_kept(entry: str, cap_entry: str, proposed: BudgetConfig, existing: BudgetConfig) -> None:
if existing.max_budget is None or not _looser_model_cap(proposed=proposed, cap=existing):
return
raise HTTPException(
status_code=403,
detail={
"error": (
f"Only a proxy admin can raise a team's model_max_budget for {entry!r} or change its "
f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration} "
f"(entry {cap_entry!r}), requested={proposed.max_budget} per {proposed.budget_duration}."
)
},
)
def _check_new_entry_outranks_no_cap(
entry: str,
existing_model_max_budget: Mapping[str, object],
routed_models: RoutedModels,
team_id: str | None,
) -> None:
"""A new entry gets its own counter, so the models it takes from a capped entry stop drawing from that cap."""
outranked: Final = next(
_capped_entries_outranked_by(
entry=entry,
existing_model_max_budget=existing_model_max_budget,
routed_models=routed_models,
team_id=team_id,
),
None,
)
if outranked is None:
return
outranked_name, cap = outranked
raise HTTPException(
status_code=403,
detail={
"error": (
f"Only a proxy admin can add a team model_max_budget entry {entry!r} that takes precedence over "
f"{outranked_name!r}, since the models they share would stop drawing from its max_budget="
f"{cap.max_budget} per {cap.budget_duration}."
)
},
)
def _capped_entries_outranked_by(
entry: str,
existing_model_max_budget: Mapping[str, object],
routed_models: RoutedModels,
team_id: str | None,
) -> Iterator[tuple[str, BudgetConfig]]:
governing: Final = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=team_id)
outranked_names: Final = itertools.chain(
()
if governing is None or entry_names_model(entry=governing.budget_model, model=entry)
else (governing.budget_model,),
entries_outranked_by(entry=entry, model_max_budget=existing_model_max_budget, routed_models=routed_models),
)
outranked_caps: Final = ((name, _existing_model_cap(existing_model_max_budget[name])) for name in outranked_names)
yield from (
(name, cap) for name, cap in outranked_caps if name != entry and cap is not None and cap.max_budget is not None
)
def _looser_model_cap(proposed: BudgetConfig, cap: BudgetConfig) -> bool:
return (
proposed.max_budget is None
or (cap.max_budget is not None and proposed.max_budget > cap.max_budget)
or (is_enforced_budget_entry(cap) and proposed.budget_duration != cap.budget_duration)
)
def _should_auto_add_team_creator(
@ -1572,7 +1648,7 @@ async def new_team(
- prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
@ -2350,7 +2426,7 @@ async def update_team(
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.

View file

@ -6,9 +6,12 @@
######################################################################
import asyncio
import functools
import itertools
import traceback
from collections.abc import Mapping, Sequence
from typing import Any, BinaryIO, Final, TypedDict, cast, get_args
from collections.abc import Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, BinaryIO, Final, NoReturn, TypedDict, cast, get_args
import httpx
from fastapi import (
@ -28,12 +31,13 @@ from typing_extensions import ReadOnly
import litellm
from litellm import CreateFileRequest, get_secret_str
from litellm._logging import verbose_proxy_logger
from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT
from litellm.litellm_core_utils.cloud_storage_security import (
is_managed_cloud_storage_uri,
)
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -112,6 +116,9 @@ from litellm.types.llms.openai import (
OpenAIFileObject,
OpenAIFilesPurpose,
)
from litellm.types.proxy.auth.user_api_key_auth import user_api_key_has_admin_view
from litellm.types.router import Deployment
from litellm.types.utils import LlmProviders
router: Final = APIRouter()
@ -1673,6 +1680,185 @@ def _as_file_list_page(response: object) -> object:
return FileListPage.model_validate(page)
async def _list_owned_managed_files(
managed_files_obj: BaseFileEndpoints,
user_api_key_dict: UserAPIKeyAuth,
purpose: str | None,
limit: int | None,
after: str | None,
before: str | None,
model_scope: FileListModelScope | None,
) -> FileListPage:
return await managed_files_obj.afile_list(
purpose=purpose,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
user_api_key_dict=user_api_key_dict,
limit=limit,
after=after,
before=before,
model_scope=model_scope,
)
@dataclass(frozen=True, slots=True)
class _UnscopedFileList:
status_code: int
detail: str
@dataclass(frozen=True, slots=True)
class _ListedFiles:
model_scope: FileListModelScope | None
custom_llm_provider: str | None
def _raise_unscoped_file_list(failure: _UnscopedFileList) -> NoReturn:
raise HTTPException(status_code=failure.status_code, detail=failure.detail)
def _single_target_model_name(target_model_names: str) -> str | _UnscopedFileList:
target_model_names_list: Final = target_model_names.split(",")
if len(target_model_names_list) != 1:
return _UnscopedFileList(
status_code=400,
detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']",
)
return target_model_names_list[0]
async def _requested_custom_llm_provider(request: Request, provider: str | None) -> str | None:
return (
provider
or get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or await get_custom_llm_provider_from_request_body(request=request)
)
def _default_file_list_limit(custom_llm_provider: str | None) -> int | None:
return ANTHROPIC_DEFAULT_FILE_LIST_LIMIT if custom_llm_provider == LlmProviders.ANTHROPIC.value else None
def _deployment_model_mapping_keys(deployment: Deployment) -> frozenset[str]:
return frozenset(key for key in (deployment.model_name, deployment.model_info.id) if key is not None)
def _model_mapping_keys(deployments: Iterable[Deployment]) -> frozenset[str]:
return frozenset(itertools.chain.from_iterable(map(_deployment_model_mapping_keys, deployments)))
def _deployments_serving_model(llm_router: Router, model: str) -> tuple[Deployment, ...]:
deployment: Final = llm_router.get_deployment(model_id=model) if llm_router.has_model_id(model) else None
if deployment is not None:
return (deployment,)
return tuple(Deployment.model_validate(listed) for listed in llm_router.get_model_list(model_name=model) or ())
def _resolved_deployment_provider(deployment: Deployment) -> str | None:
try:
return litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)[1]
except litellm.BadRequestError:
return None
def _deployments_on_provider(llm_router: Router | None, custom_llm_provider: str) -> Iterator[Deployment]:
if llm_router is None:
return iter(())
return (
deployment
for deployment_id in llm_router.get_model_ids()
if (deployment := llm_router.get_deployment(model_id=deployment_id)) is not None
and _resolved_deployment_provider(deployment) == custom_llm_provider
)
def _routed_model_name_on_provider(
llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str
) -> bool:
if model_name in llm_router.model_names:
return model_name in provider_model_names
return llm_router.pattern_router.routes_to_provider(model_name, custom_llm_provider)
def _stored_model_name_routes_to_provider(
llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str
) -> bool:
routing_group: Final = llm_router.get_routing_group(model_name)
routed_model_names: Final = (
tuple(routing_group.models) if routing_group is not None else (llm_router.routable_model_group(model_name),)
)
return any(
_routed_model_name_on_provider(llm_router, custom_llm_provider, provider_model_names, routed_model_name)
for routed_model_name in routed_model_names
)
def _files_on_provider(llm_router: Router | None, custom_llm_provider: str) -> _ListedFiles | _UnscopedFileList:
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import passthrough_file_model_id
if custom_llm_provider not in litellm.provider_list:
return _UnscopedFileList(status_code=400, detail=f"Unknown provider '{custom_llm_provider}' for listing files.")
passthrough_model_id: Final = passthrough_file_model_id(custom_llm_provider)
deployments: Final = tuple(_deployments_on_provider(llm_router, custom_llm_provider))
model_mapping_keys: Final = _model_mapping_keys(deployments) | frozenset(
(passthrough_model_id,) if passthrough_model_id is not None else ()
)
if llm_router is None:
return _ListedFiles(FileListModelScope(model_mapping_keys), custom_llm_provider)
return _ListedFiles(
FileListModelScope(
model_mapping_keys,
routes_stored_model_name=functools.partial(
_stored_model_name_routes_to_provider,
llm_router,
custom_llm_provider,
frozenset(deployment.model_name for deployment in deployments),
),
),
custom_llm_provider,
)
def _files_serving_model(llm_router: Router, model: str) -> _ListedFiles:
deployments: Final = _deployments_serving_model(llm_router, model)
providers: Final = frozenset(map(_resolved_deployment_provider, deployments))
return _ListedFiles(
FileListModelScope(_model_mapping_keys(deployments)), next(iter(providers)) if len(providers) == 1 else None
)
async def _listed_files(
llm_router: Router | None,
user_api_key_dict: UserAPIKeyAuth,
custom_llm_provider: str | None,
target_model_names: str | None,
routed_model: str | None,
) -> _ListedFiles | _UnscopedFileList:
if routed_model is not None and llm_router is not None:
return _files_serving_model(llm_router, routed_model)
if target_model_names:
target_model_name: Final = _single_target_model_name(target_model_names)
if isinstance(target_model_name, _UnscopedFileList):
return target_model_name
if llm_router is None:
return _UnscopedFileList(
status_code=500, detail="LLM Router not initialized. Ensure models added to proxy."
)
await get_authorized_credentials_for_model(
llm_router=llm_router,
model_id=target_model_name,
user_api_key_dict=user_api_key_dict,
operation_context="file list",
)
return _files_serving_model(llm_router, target_model_name)
if custom_llm_provider is None:
return _ListedFiles(None, None)
return _files_on_provider(llm_router, custom_llm_provider)
@router.get(
"/{provider}/v1/files",
dependencies=[Depends(user_api_key_auth)],
@ -1697,6 +1883,8 @@ async def list_files(
purpose: str | None = None,
limit: int | None = None,
after: str | None = None,
after_id: str | None = None,
before_id: str | None = None,
):
"""
Returns information about a specific file. that can be used across - Assistants API, Batch API
@ -1739,6 +1927,10 @@ async def list_files(
)
response: Any | None = None
owned_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files")
list_owned_files_only: Final = litellm.require_managed_files is True and not user_api_key_has_admin_view(
user_api_key_dict
)
# Check for model-based credential routing (no file_id encoding check for list)
should_route, model_used, _, credentials = await handle_model_based_routing(
@ -1750,7 +1942,36 @@ async def list_files(
check_file_id_encoding=False,
)
if should_route and credentials is not None:
if list_owned_files_only:
if after and after_id and after != after_id:
raise HTTPException(
status_code=400,
detail="Pass the list cursor as either 'after' or 'after_id', not both with different ids.",
)
if not isinstance(owned_files_hook, BaseFileEndpoints):
raise HTTPException(
status_code=500,
detail="Managed file listing is unavailable, so require_managed_files cannot scope the list to the caller.",
)
listed_files: Final = await _listed_files(
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=await _requested_custom_llm_provider(request=request, provider=provider),
target_model_names=target_model_names,
routed_model=model_used if should_route else None,
)
if isinstance(listed_files, _UnscopedFileList):
_raise_unscoped_file_list(listed_files)
response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response
managed_files_obj=owned_files_hook,
user_api_key_dict=user_api_key_dict,
purpose=purpose,
limit=limit if limit is not None else _default_file_list_limit(listed_files.custom_llm_provider),
after=after or after_id,
before=before_id,
model_scope=listed_files.model_scope,
)
elif should_route and credentials is not None:
# Use model-based routing with credentials from config
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
response = await litellm.afile_list(
@ -1762,12 +1983,9 @@ async def list_files(
verbose_proxy_logger.debug("Listed files using model: %s", model_used)
elif target_model_names and isinstance(target_model_names, str):
target_model_names_list: Final = target_model_names.split(",")
if len(target_model_names_list) != 1:
raise HTTPException(
status_code=400,
detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']",
)
target_model_name: Final = _single_target_model_name(target_model_names)
if isinstance(target_model_name, _UnscopedFileList):
_raise_unscoped_file_list(target_model_name)
if llm_router is None:
raise HTTPException(
status_code=500,
@ -1775,7 +1993,7 @@ async def list_files(
)
credentials = await get_authorized_credentials_for_model(
llm_router=llm_router,
model_id=target_model_names_list[0],
model_id=target_model_name,
user_api_key_dict=user_api_key_dict,
operation_context="file list",
)
@ -1786,20 +2004,17 @@ async def list_files(
**data,
)
else:
custom_llm_provider: Final = (
provider
or get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or await get_custom_llm_provider_from_request_body(request=request)
)
custom_llm_provider: Final = await _requested_custom_llm_provider(request=request, provider=provider)
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints):
response = await managed_files_obj.afile_list(
purpose=purpose,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response
managed_files_obj=managed_files_obj,
user_api_key_dict=user_api_key_dict,
purpose=purpose,
limit=limit,
after=after,
before=None,
model_scope=None,
)
else:
resolved_custom_llm_provider: Final = custom_llm_provider or "openai"

View file

@ -242,6 +242,12 @@ def _passthrough_sentinel_model_id(provider: str) -> str:
return f"_passthrough_{provider}"
def passthrough_file_model_id(provider: str) -> str | None:
if (provider, "/v1/files") not in _LIST_ROUTE_TABLE:
return None
return _passthrough_sentinel_model_id(provider)
# Key under which the provider marker is stored in a file row's model_mappings.
# Its value lands in flat_model_file_ids (built from model_mappings.values()),
# giving the file table a DB-queryable provider scope it otherwise lacks.

View file

@ -666,7 +666,11 @@ from litellm.proxy.list_api.common import (
problem_response,
request_validation_problem,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request
from litellm.proxy.litellm_pre_call_utils import (
SCHEDULER_PRIORITY_HEADER,
LiteLLMProxyRequestSetup,
add_litellm_data_to_request,
)
from litellm.proxy.logging_endpoints.callback_logs_endpoints import (
rust_control_plane_router,
)
@ -18500,6 +18504,11 @@ async def async_queue_request(
raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
router: Final = llm_router
data.update(
router.priority_header_update(
model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
)
)
queued_priority, _ = router.split_request_priority(model=data["model"], kwargs=data)
if "stream" in data and data["stream"] is True: # use generate_responses to stream responses

View file

@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
@@index([unified_file_id])
@@index([flat_model_file_ids], type: Gin)
@@index([team_id, created_at(sort: Desc)])
@@index([created_by, created_at(sort: Desc)])
}
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport {
@@index([window_end])
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
}

View file

@ -19,6 +19,7 @@ from litellm.repositories.daily_activity_sql import (
build_cache_leakage_keys_sql,
build_entity_rollup_sql,
build_export_sql,
build_guardrail_spend_sql,
build_key_page_sql,
build_key_search_sql,
build_model_top_keys_sql,
@ -36,6 +37,7 @@ from litellm.types.repositories.daily_activity import (
ExportRow,
ExportType,
GroupingSetsRow,
GuardrailSpendRow,
KeyMetadataRow,
KeyPage,
KeySpendRow,
@ -100,6 +102,7 @@ class DailyActivityDatabase(Protocol):
_GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...])
_ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...])
_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...])
_KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...])
_KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int)
_USER_SPEND_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
@ -191,6 +194,9 @@ class DailyActivityRepository:
distinct_api_keys=distinct_api_keys,
)
async def guardrail_spend(self, scope: DailyActivityScope) -> tuple[GuardrailSpendRow, ...]:
return _GUARDRAIL_SPEND_ADAPTER.validate_python(await self._query(build_guardrail_spend_sql(scope)))
async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]:
if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX:
raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}")

View file

@ -261,6 +261,25 @@ def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) ->
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit))
def build_guardrail_spend_sql(scope: DailyActivityScope) -> SqlQuery:
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
where_clause, where_params = build_where_clause(scope)
sql: Final = f"""
SELECT g.date, g.guardrail_id, COALESCE(t.guardrail_name, g.guardrail_id) AS guardrail_name, g.api_key,
COALESCE(SUM(g.cost), 0)::float AS spend,
COALESCE(SUM(CASE WHEN g.cost IS NULL THEN g.units ELSE g.untracked_units END), 0)::bigint AS untracked_units
FROM "LiteLLM_DailyGuardrailUsageUnits" g
LEFT JOIN "LiteLLM_GuardrailsTable" t ON t.guardrail_id = g.guardrail_id
WHERE g.date >= $1 AND g.date <= $2
AND EXISTS (
SELECT 1 FROM "{pg_table}"
WHERE {where_clause} AND "{pg_table}".date = g.date AND "{pg_table}".api_key = g.api_key
)
GROUP BY g.date, g.guardrail_id, t.guardrail_name, g.api_key
"""
return SqlQuery(sql=sql, params=where_params)
def _key_spend_select() -> str:
return """
COALESCE(SUM(spend), 0)::float AS spend,

View file

@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
from litellm.litellm_core_utils.internal_call_metadata import (
is_background_response,
requests_background,
without_budget_reservation,
)
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
@ -65,7 +69,7 @@ from litellm.types.llms.openai import (
ResponsesAPIStreamingResponse,
)
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
from litellm.types.utils import CallTypes
from litellm.types.utils import CallTypes, ModelResponse
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
@ -288,7 +292,7 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]:
return _status_code_for_error_fields(error_type, error_code), message
def _map_stream_error_to_exception(
def map_stream_error_to_exception(
error_obj: object,
model: str,
custom_llm_provider: str,
@ -578,6 +582,7 @@ class BaseResponsesAPIStreamingIterator:
self.request_data.get("input"),
self.request_data,
self._generated_content + self._generated_tool_arguments,
use_default_image_token_count=False,
)
if _estimate_wanted
else None
@ -731,7 +736,7 @@ class BaseResponsesAPIStreamingIterator:
self._earlier_rounds_spend = earlier
def _map_error_event_exception(self, error_obj: object) -> Exception:
return _map_stream_error_to_exception(
return map_stream_error_to_exception(
error_obj,
self.model or "",
self.custom_llm_provider or "",
@ -981,6 +986,39 @@ class BaseResponsesAPIStreamingIterator:
except Exception:
pass
def _record_partial_usage_for_failure(self) -> None:
generated_text: Final = self._generated_content + self._generated_tool_arguments
if (
not generated_text
or requests_background(self.request_data)
or is_background_response(self._created_response)
or self._completed_response_logged
or self.logging_obj.model_call_details.get("combined_usage_object") is not None
):
return
estimate: Final = _estimate_usage_safely(
self.model or "",
self.request_data.get("input"),
self.request_data,
generated_text,
use_default_image_token_count=True,
)
if estimate is None:
return
usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses
estimate
)
served_tier: Final[object] = getattr(self._created_response, "service_tier", None)
assembled: Final = ModelResponse(
model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None
)
if self._earlier_rounds_spend is None:
self.logging_obj.record_assembled_response_for_failure(assembled)
return
self.logging_obj.record_failure_after_earlier_rounds(
self._earlier_rounds_spend, usage, self.logging_obj.response_cost_calculator(result=assembled) or 0.0
)
def _handle_failure(self, exception: Exception):
"""
Trigger failure handlers before bubbling the exception.
@ -990,6 +1028,7 @@ class BaseResponsesAPIStreamingIterator:
if self._failure_handled:
return
self._failure_handled = True
self._record_partial_usage_for_failure()
traceback_exception: Final = traceback.format_exc()
end_time: Final = datetime.now()
@ -1060,15 +1099,20 @@ class BaseResponsesAPIStreamingIterator:
the same shape the success handler assembles for a natural end of stream.
Its response is the response.created payload (it carries the served
service_tier) plus the text and tool arguments relayed so far, with counted
usage since the provider never sent any. None when nothing was relayed yet
or the stream already reached its terminal event and billed itself.
usage since the provider never sent any. None when nothing was relayed yet,
the stream already reached its terminal event and billed itself, or its failure
log already billed the relayed output.
"""
created: Final = self._created_response
generated: Final = self._generated_content + self._generated_tool_arguments
if self.completed_response is not None or created is None or not generated:
if created is None or not generated or self._billed_without_disconnect():
return None
usage: Final = _estimate_usage_safely(
self.model or "", self.request_data.get("input"), self.request_data, generated
self.model or "",
self.request_data.get("input"),
self.request_data,
generated,
use_default_image_token_count=False,
)
if usage is None:
return None
@ -1082,6 +1126,9 @@ class BaseResponsesAPIStreamingIterator:
stamp_responses_usage_cost(partial, self.logging_obj)
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
def _billed_without_disconnect(self) -> bool:
return self.completed_response is not None or self._failure_handled
def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
self._yielded_first_chunk = True
if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:
@ -1218,6 +1265,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
raise StopAsyncIteration from e
self._raise_for_transport_error(e)
except httpx.TimeoutException as e:
if self.completed_response is not None:
self.finished = True
raise StopAsyncIteration from e
self._raise_for_transport_timeout(e)
except httpx.HTTPError as e:
# Handle HTTP errors
@ -1303,6 +1353,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
raise StopIteration from e
self._raise_for_transport_error(e)
except httpx.TimeoutException as e:
if self.completed_response is not None:
self.finished = True
raise StopIteration from e
self._raise_for_transport_timeout(e)
except httpx.HTTPError as e:
# Handle HTTP errors
@ -1727,6 +1780,8 @@ def _estimate_usage_from_text(
request_input: object,
responses_api_request: Mapping[str, object],
generated_text: str,
*,
use_default_image_token_count: bool = False,
) -> ResponseAPIUsage:
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
@ -1734,7 +1789,7 @@ def _estimate_usage_from_text(
only_since_compaction=True,
)
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
model=model, messages=messages
model=model, messages=messages, use_default_image_token_count=use_default_image_token_count
)
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
model=model, text=generated_text, count_response_tokens=True
@ -1751,6 +1806,8 @@ def _estimate_usage_safely(
request_input: object,
responses_api_request: Mapping[str, object],
generated_text: str,
*,
use_default_image_token_count: bool,
) -> ResponseAPIUsage | None:
try:
return _estimate_usage_from_text(
@ -1758,6 +1815,7 @@ def _estimate_usage_safely(
request_input=request_input,
responses_api_request=responses_api_request,
generated_text=generated_text,
use_default_image_token_count=use_default_image_token_count,
)
except Exception as e:
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
@ -2276,7 +2334,7 @@ class _ResponsesWsTurn:
)
if failed_event is None:
return None
return _map_stream_error_to_exception(
return map_stream_error_to_exception(
_ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj
)

View file

@ -359,6 +359,7 @@ from litellm.router_utils.request_priority import (
InvalidPriority,
parse_default_priority,
priority_error,
priority_header_fields,
request_drops_params,
resolve_request_priority,
unqueueable_priority_message,
@ -5411,14 +5412,22 @@ class Router:
request_priority: Final = resolve_request_priority(
requested=kwargs.get("priority"),
default_priority=self.default_priority,
drops_params=lambda: request_drops_params(
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
),
drops_params=lambda: self._drops_params(model=model, kwargs=kwargs),
)
if isinstance(request_priority, InvalidPriority):
raise priority_error(model=model, message=request_priority.message)
return request_priority, {key: value for key, value in kwargs.items() if key != "priority"}
def priority_header_update(
self, model: str, kwargs: Mapping[str, object], header: str | None
) -> Mapping[str, int | str]:
return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs))
def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool:
return request_drops_params(
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
)
async def _schedule_factory(
self,
model: str,
@ -6456,6 +6465,9 @@ class Router:
hop_kwargs: Final = mid_stream_fallback_hop_kwargs(
model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs
)
shared_logging: Final = kwargs.get("litellm_logging_obj")
if isinstance(shared_logging, LiteLLMLogging):
shared_logging.carry_failed_attempt_usage()
response: Final = await self._ageneric_api_call_with_fallbacks_helper(
model=model, original_generic_function=original_generic_function, **kwargs
)
@ -11847,6 +11859,11 @@ class Router:
indices: Final = self.model_name_to_deployment_indices.get(model_name)
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
def team_public_model_names_of(self, team_id: str) -> frozenset[str]:
"""The public names `team_id` serves its own deployments under."""
with self._deployment_index_lock:
return frozenset(name for owner, name in self.team_model_to_deployment_indices if owner == team_id)
def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None:
"""The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed."""
with self._deployment_index_lock:

View file

@ -13,7 +13,7 @@ from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._internal_context import service_target
@ -417,17 +417,35 @@ def _should_run_cooldown_logic(
return True
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...])
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...])
_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str])
def _fallback_target_name(target: str | Mapping[str, object]) -> str:
return target if isinstance(target, str) else str(target.get("model", ""))
def _fallback_target_names(
targets: Sequence[str | Mapping[str, object]] | str,
) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets]
def _fallback_chain_names(
entry: Mapping[str, Sequence[str | Mapping[str, object]]],
) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()}
entry: object,
) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
try:
validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry)
except ValidationError:
return None
return {key: _fallback_target_names(targets) for key, targets in validated.items()}
def _fallback_entries(configured: object) -> tuple[object, ...]:
try:
return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
except ValidationError:
return ()
def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool:
@ -437,9 +455,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_
if requested_model_group is None or not configured:
return False
chains: Final = tuple(
_fallback_chain_names(entry)
for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
if isinstance(entry, Mapping)
chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None
)
fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group)
return bool(fallback_group)

View file

@ -5,15 +5,19 @@ When enabled via router_settings.optional_pre_call_checks: ["enforce_model_rate_
- tpm/rpm: combined TPM + optional RPM (legacy)
- itpm/otpm: separate input/output tokens per minute
- tpd/rpd: tokens and requests per UTC calendar day
When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning
is logged the first time such a deployment is seen.
"""
import contextlib
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Callable
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from pydantic import BaseModel, ConfigDict
import litellm
from litellm._internal_context import with_service_target
@ -33,6 +37,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
io_token_refund_failure,
)
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.router import RouterErrors
from litellm.types.utils import StandardLoggingPayload
from litellm.utils import get_utc_datetime
@ -47,6 +52,62 @@ else:
class RoutingArgs:
ttl: int = 60 # 1min (RPM/TPM expire key)
daily_ttl: int = 86400
DAY_FORMAT: Final = "%Y-%m-%d"
def seconds_until_utc_midnight(now: datetime) -> int:
next_midnight: Final = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
return max(1, int((next_midnight - now).total_seconds()))
def _first_set(*values: int | None) -> int | None:
return next((value for value in values if value is not None), None)
class _DailyLimitFields(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
tpd: int | None = None
rpd: int | None = None
class _DailyLimitParams(_DailyLimitFields):
model: str | None = None
class _DailyLimitModelInfo(_DailyLimitFields):
id: str | None = None
class _DailyLimitDeployment(_DailyLimitFields):
model_name: str = ""
litellm_params: _DailyLimitParams = _DailyLimitParams()
model_info: _DailyLimitModelInfo = _DailyLimitModelInfo()
def limits(self) -> tuple[int | None, int | None]:
return (
_first_set(self.tpd, self.litellm_params.tpd, self.model_info.tpd),
_first_set(self.rpd, self.litellm_params.rpd, self.model_info.rpd),
)
def cache_keys(self, current_day: str) -> tuple[str, str]:
prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}"
return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}"
def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError:
return litellm.RateLimitError(
message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}",
llm_provider="",
model=self.litellm_params.model,
response=httpx.Response(
status_code=429,
content=f"{RouterErrors.user_defined_ratelimit_error.value} {kind} limit={limit}. current usage={usage}. id={self.model_info.id}, model_group={self.model_name}",
headers={"retry-after": str(seconds_until_utc_midnight(now))},
request=httpx.Request(method="model_rate_limit_check", url="https://github.com/BerriAI/litellm"),
),
)
class ModelRateLimitingCheck(CustomLogger):
@ -60,12 +121,60 @@ class ModelRateLimitingCheck(CustomLogger):
this check actively enforces those limits across ALL routing strategies.
"""
def __init__(self, dual_cache: DualCache):
def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime):
self.dual_cache = dual_cache
self.clock = clock
# model_ids already warned about itpm/otpm + tpm/rpm on the same deployment,
# so the warning is logged once per deployment rather than per request.
self._io_token_conflict_warned_ids: set[str] = set()
def _increment(self, key: str, value: int, ttl: int) -> int:
return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl)
async def _async_increment(
self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None
) -> int | None:
count: Final = await self.dual_cache.async_increment_cache(
key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span
)
return None if count is None else int(count)
@staticmethod
def _request_count_operations(
rpm_key: str | None, rpd_key: str | None, increment_value: int
) -> tuple[RedisPipelineIncrementOperation, ...]:
return tuple(
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl)
for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl))
if key is not None
)
async def _async_refund_request_counts(
self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None
) -> None:
operations: Final = self._request_count_operations(rpm_key, rpd_key, -1)
if operations:
await self.dual_cache.async_increment_cache_pipeline(
increment_list=list(operations), parent_otel_span=parent_otel_span
)
async def _async_increment_request_counts(
self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None
) -> tuple[int | None, int | None]:
operations: Final = self._request_count_operations(rpm_key, rpd_key, 1)
if not operations:
return None, None
counts: Final = await self.dual_cache.async_increment_cache_pipeline(
increment_list=list(operations), parent_otel_span=parent_otel_span
)
if counts is None or len(counts) != len(operations):
return None, None
count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)}
return (
None if rpm_key is None else count_by_key.get(rpm_key),
None if rpd_key is None else count_by_key.get(rpd_key),
)
def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None:
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
if tpm_limit is None and rpm_limit is None:
@ -129,6 +238,21 @@ class ModelRateLimitingCheck(CustomLogger):
return tpm, rpm
def _check_daily_limits(self, deployment: _DailyLimitDeployment, now: datetime) -> None:
tpd_limit, rpd_limit = deployment.limits()
if tpd_limit is None and rpd_limit is None:
return
tpd_key, rpd_key = deployment.cache_keys(now.strftime(DAY_FORMAT))
if tpd_limit is not None:
current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit)
if current_tpd is not None and current_tpd >= tpd_limit:
raise deployment.limit_error("tpd", tpd_limit, current_tpd, now)
if rpd_limit is not None:
current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl)
if current_rpd > rpd_limit:
self._increment(rpd_key, -1, RoutingArgs.daily_ttl)
raise deployment.limit_error("rpd", rpd_limit, current_rpd, now)
def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]:
"""Get the cache keys for TPM and RPM tracking."""
model_id: Final = deployment.get("model_info", {}).get("id")
@ -178,13 +302,10 @@ class ModelRateLimitingCheck(CustomLogger):
)
io_reservation_made = True
dt: Final = self.clock()
daily: Final = _DailyLimitDeployment.model_validate(deployment)
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
# If no limits are set, allow the request
if tpm_limit is None and rpm_limit is None:
return deployment
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute)
@ -213,8 +334,9 @@ class ModelRateLimitingCheck(CustomLogger):
# Check RPM limit (atomic increment-first to avoid race conditions)
if rpm_limit is not None:
current_rpm: Final = self.dual_cache.increment_cache(key=rpm_key, value=1, ttl=RoutingArgs.ttl)
current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl)
if current_rpm is not None and current_rpm > rpm_limit:
self._increment(rpm_key, -1, RoutingArgs.ttl)
raise litellm.RateLimitError(
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
llm_provider="",
@ -230,6 +352,12 @@ class ModelRateLimitingCheck(CustomLogger):
),
)
try:
self._check_daily_limits(daily, dt)
except litellm.RateLimitError:
if rpm_limit is not None:
self._increment(rpm_key, -1, RoutingArgs.ttl)
raise
return deployment
except litellm.RateLimitError:
@ -259,15 +387,14 @@ class ModelRateLimitingCheck(CustomLogger):
)
io_reservation_made = True
dt: Final = self.clock()
daily: Final = _DailyLimitDeployment.model_validate(deployment)
tpd_limit, rpd_limit = daily.limits()
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
# If no limits are set, allow the request
if tpm_limit is None and rpm_limit is None:
return deployment
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute)
tpd_key, rpd_key = daily.cache_keys(dt.strftime(DAY_FORMAT))
model_id: Final = deployment.get("model_info", {}).get("id")
model_name: Final = deployment.get("litellm_params", {}).get("model")
@ -293,30 +420,37 @@ class ModelRateLimitingCheck(CustomLogger):
num_retries=0, # Don't retry - return 429 immediately
)
# Check RPM limit (atomic increment-first to avoid race conditions)
if rpm_limit is not None:
current_rpm: Final = await self.dual_cache.async_increment_cache(
key=rpm_key,
value=1,
ttl=RoutingArgs.ttl,
parent_otel_span=parent_otel_span,
)
if current_rpm is not None and current_rpm > rpm_limit:
raise litellm.RateLimitError(
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
llm_provider="",
model=model_name,
response=httpx.Response(
status_code=429,
content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
headers={"retry-after": str(60)},
request=httpx.Request(
method="model_rate_limit_check",
url="https://github.com/BerriAI/litellm",
),
if tpd_limit is not None:
current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span)
if current_tpd is not None and current_tpd >= tpd_limit:
raise daily.limit_error("tpd", tpd_limit, current_tpd, dt)
charged_rpm_key: Final = rpm_key if rpm_limit is not None else None
charged_rpd_key: Final = rpd_key if rpd_limit is not None else None
current_rpm, current_rpd = await self._async_increment_request_counts(
charged_rpm_key, charged_rpd_key, parent_otel_span
)
if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit:
await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span)
raise litellm.RateLimitError(
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
llm_provider="",
model=model_name,
response=httpx.Response(
status_code=429,
content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
headers={"retry-after": str(60)},
request=httpx.Request(
method="model_rate_limit_check",
url="https://github.com/BerriAI/litellm",
),
num_retries=0, # Don't retry - return 429 immediately
)
),
num_retries=0, # Don't retry - return 429 immediately
)
if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit:
await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span)
raise daily.limit_error("rpd", rpd_limit, current_rpd, dt)
return deployment
@ -373,17 +507,16 @@ class ModelRateLimitingCheck(CustomLogger):
if not model or not total_tokens:
return
dt: Final = get_utc_datetime()
dt: Final = self.clock()
current_minute: Final = dt.strftime("%H-%M")
tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}"
tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}"
verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens)
await self.dual_cache.async_increment_cache(
key=tpm_key,
value=total_tokens,
ttl=RoutingArgs.ttl,
)
tokens: Final = int(total_tokens)
await self._async_increment(tpm_key, tokens, RoutingArgs.ttl)
await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl)
except Exception as e:
verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e)
@ -434,15 +567,14 @@ class ModelRateLimitingCheck(CustomLogger):
if not model or not total_tokens:
return
dt: Final = get_utc_datetime()
dt: Final = self.clock()
current_minute: Final = dt.strftime("%H-%M")
tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}"
tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}"
self.dual_cache.increment_cache(
key=tpm_key,
value=total_tokens,
ttl=RoutingArgs.ttl,
)
tokens: Final = int(total_tokens)
self._increment(tpm_key, tokens, RoutingArgs.ttl)
self._increment(tpd_key, tokens, RoutingArgs.daily_ttl)
except Exception as e:
verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e)

View file

@ -47,6 +47,15 @@ def resolve_request_priority(
return priority
def priority_header_fields(header: str | None, drops_params: Callable[[], bool]) -> Mapping[str, int | str]:
if not header:
return {}
priority: Final = _as_priority(header)
if isinstance(priority, InvalidPriority):
return {} if drops_params() else {"priority": header}
return {"priority": priority}
def parse_default_priority(value: object) -> int | None | InvalidPriority:
if value is None:
return None

View file

@ -189,6 +189,8 @@ class DeploymentOptions:
tpm: int | None = None
itpm: int | None = None
otpm: int | None = None
rpd: int | None = None
tpd: int | None = None
kubernetes_pod_discovery: bool | None = None
default_api_key_rpm_limit: int | None = None
default_api_key_tpm_limit: int | None = None

View file

@ -65,6 +65,34 @@ class MetricWithMetadata(MetricBase):
api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
class GuardrailSpendMetrics(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
spend: float = 0.0
untracked_units: int = 0
class GuardrailKeyMetric(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
metrics: GuardrailSpendMetrics
metadata: KeyMetadata
class GuardrailMetadata(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
guardrail_name: str
class GuardrailMetric(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
metrics: GuardrailSpendMetrics
metadata: GuardrailMetadata
api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict)
class BreakdownMetrics(LiteLLMBaseModel):
"""Breakdown of spend by different dimensions"""
@ -75,6 +103,7 @@ class BreakdownMetrics(LiteLLMBaseModel):
endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata}
api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata}
guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend
class DailySpendData(LiteLLMBaseModel):

View file

@ -168,6 +168,16 @@ class EntityRollupRow(RollupMetricsRow):
distinct_api_keys: int | None
@dataclass(frozen=True, slots=True)
class GuardrailSpendRow:
date: str
guardrail_id: str
guardrail_name: str
api_key: str
spend: float
untracked_units: int
@dataclass(frozen=True, slots=True)
class AggregatedRows:
grouping_rows: tuple[GroupingSetsRow, ...]

View file

@ -533,6 +533,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
rpm: int | None = None
itpm: int | None = None
otpm: int | None = None
tpd: int | None = None
rpd: int | None = None
timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/
# read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/
stream_timeout: float | str | None = None
@ -738,6 +740,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
rpm: int | None
itpm: int | None
otpm: int | None
tpd: ReadOnly[int | None]
rpd: ReadOnly[int | None]
order: int | None
weight: int | None
max_parallel_requests: int | None

View file

@ -10418,6 +10418,13 @@ class ProviderConfigManager:
return get_modelscope_image_generation_config(model)
elif LlmProviders.EDENAI == provider:
return litellm.EdenAIImageGenerationConfig()
elif LlmProviders.TENCENT == provider:
from litellm.llms.tencent.image_generation.transformation import (
TencentImageGenerationConfig,
is_tokenhub_image_model,
)
return TencentImageGenerationConfig() if is_tokenhub_image_model(model) else None
return None
@staticmethod
@ -10606,6 +10613,15 @@ class ProviderConfigManager:
)
return get_openrouter_image_edit_config(model)
elif LlmProviders.TENCENT == provider:
from litellm.llms.tencent.image_edit.transformation import (
TencentImageEditConfig,
)
from litellm.llms.tencent.image_generation.transformation import (
is_tokenhub_image_model,
)
return TencentImageEditConfig() if is_tokenhub_image_model(model) else None
return None
@staticmethod

View file

@ -21589,7 +21589,13 @@
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
"supports_video_input": true,
"input_cost_per_token_priority": 5.6252e-07,
"output_cost_per_token_priority": 3.37498e-06,
"cache_read_input_token_cost_priority": 5.628e-08,
"cache_creation_input_token_cost_priority": 5.6252e-07,
"input_cost_per_audio_token_priority": 1.12497e-06,
"cache_read_input_audio_token_cost_priority": 1.1249e-07
},
"databricks/databricks-gemini-3-1-flash-image": {
"cache_creation_input_token_cost": 6.2503e-07,
@ -21676,7 +21682,14 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.50002e-06,
"output_cost_per_token_priority": 2.699998e-05,
"cache_read_input_token_cost_priority": 4.5003e-07,
"cache_creation_input_token_cost_priority": 4.50002e-06,
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
},
"databricks/databricks-gemini-3-flash": {
"cache_creation_input_token_cost": 6.2503e-07,
@ -21710,7 +21723,13 @@
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
"supports_video_input": true,
"input_cost_per_token_priority": 1.12497e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.1249e-07,
"cache_creation_input_token_cost_priority": 1.12497e-06,
"input_cost_per_audio_token_priority": 2.25001e-06,
"cache_read_input_audio_token_cost_priority": 2.2498e-07
},
"databricks/databricks-gemini-3-pro": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -21757,7 +21776,11 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-7-flash": {
"cache_creation_input_token_cost": 7.4998e-07,
@ -21785,10 +21808,14 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-6-flash": {
"cache_creation_input_token_cost": 1.87502e-06,
"cache_creation_input_token_cost": 7.4998e-07,
"cache_read_input_token_cost": 7.497e-08,
"input_cost_per_token": 7.4998e-07,
"input_dbu_cost_per_token": 1.0714e-05,
@ -21816,7 +21843,11 @@
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.35002e-06,
"output_cost_per_token_priority": 6.75003e-06,
"cache_read_input_token_cost_priority": 1.3503e-07,
"cache_creation_input_token_cost_priority": 1.35002e-06
},
"databricks/databricks-gemini-3-5-flash": {
"cache_creation_input_token_cost": 1.87502e-06,
@ -21846,7 +21877,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 3.37498e-06,
"output_cost_per_token_priority": 2.025002e-05,
"cache_read_input_token_cost_priority": 3.3747e-07,
"cache_creation_input_token_cost_priority": 3.37498e-06
},
"databricks/databricks-gemini-3-5-flash-lite": {
"cache_creation_input_token_cost": 3.7499e-07,
@ -21878,7 +21913,11 @@
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 6.7501e-07,
"output_cost_per_token_priority": 5.62499e-06,
"cache_read_input_token_cost_priority": 6.748e-08,
"cache_creation_input_token_cost_priority": 6.7501e-07
},
"databricks/databricks-gemma-3-12b": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -21921,7 +21960,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
"supports_vision": false,
"input_cost_per_token_priority": 2.45e-06,
"output_cost_per_token_priority": 7.7e-06,
"cache_read_input_token_cost_priority": 4.55e-07,
"cache_creation_input_token_cost_priority": 2.45e-06
},
"databricks/databricks-glm-5-3": {
"cache_creation_input_token_cost": 1.4e-06,
@ -21951,7 +21994,11 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 2.45e-06,
"output_cost_per_token_priority": 7.7e-06,
"cache_read_input_token_cost_priority": 4.55e-07,
"cache_creation_input_token_cost_priority": 2.45e-06
},
"databricks/databricks-glm-5-3-flash": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -21980,7 +22027,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 2.625e-07,
"output_cost_per_token_priority": 8.75e-07,
"cache_read_input_token_cost_priority": 5.25e-08,
"cache_creation_input_token_cost_priority": 2.625e-07
},
"databricks/databricks-gpt-5": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -21999,7 +22050,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 2.49998e-06,
"output_cost_per_token_priority": 1.999998e-05,
"cache_read_input_token_cost_priority": 2.4997e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06
},
"databricks/databricks-gpt-5-1": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -22018,7 +22073,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 2.49998e-06,
"output_cost_per_token_priority": 1.999998e-05,
"cache_read_input_token_cost_priority": 2.4997e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06
},
"databricks/databricks-gpt-5-2": {
"cache_creation_input_token_cost": 1.75e-06,
@ -22037,7 +22096,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 3.5e-06,
"output_cost_per_token_priority": 2.8e-05,
"cache_read_input_token_cost_priority": 3.5e-07,
"cache_creation_input_token_cost_priority": 3.5e-06
},
"databricks/databricks-gpt-5-3-codex": {
"cache_creation_input_token_cost": 1.75e-06,
@ -22056,7 +22119,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 3.5e-06,
"output_cost_per_token_priority": 2.8e-05,
"cache_read_input_token_cost_priority": 3.5e-07,
"cache_creation_input_token_cost_priority": 3.5e-06
},
"databricks/databricks-gpt-5-4": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -22087,7 +22154,14 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.00003e-06,
"output_cost_per_token_priority": 2.999997e-05,
"cache_read_input_token_cost_priority": 5.0001e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
},
"databricks/databricks-gpt-5-4-mini": {
"cache_creation_input_token_cost": 7.4998e-07,
@ -22115,7 +22189,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.50003e-06,
"output_cost_per_token_priority": 8.99997e-06,
"cache_read_input_token_cost_priority": 1.5001e-07,
"cache_creation_input_token_cost_priority": 1.50003e-06
},
"databricks/databricks-gpt-5-4-nano": {
"cache_creation_input_token_cost": 1.9999e-07,
@ -22175,7 +22253,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 8.00002e-06,
"output_cost_per_token_priority": 4.000003e-05,
"cache_read_input_token_cost_priority": 8.0003e-07,
"cache_creation_input_token_cost_priority": 9.99999e-06,
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
},
"databricks/databricks-gpt-5-6-terra": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -22207,7 +22293,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 2.399999e-05,
"cache_read_input_token_cost_priority": 3.9998e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-5-6-luna": {
"cache_creation_input_token_cost": 2.4997e-07,
@ -22239,7 +22333,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 3.9998e-07,
"output_cost_per_token_priority": 2.40002e-06,
"cache_read_input_token_cost_priority": 3.997e-08,
"cache_creation_input_token_cost_priority": 5.0001e-07,
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
},
"databricks/databricks-gpt-5-5": {
"cache_creation_input_token_cost": 5.00003e-06,
@ -22270,7 +22372,14 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.249997e-05,
"output_cost_per_token_priority": 7.50001e-05,
"cache_read_input_token_cost_priority": 1.24999e-06,
"cache_creation_input_token_cost_priority": 1.249997e-05,
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
},
"databricks/databricks-gpt-5-5-pro": {
"cache_creation_input_token_cost": 2.999997e-05,
@ -22319,7 +22428,11 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"input_cost_per_token_priority": 4.5003e-07,
"output_cost_per_token_priority": 3.60003e-06,
"cache_read_input_token_cost_priority": 4.501e-08,
"cache_creation_input_token_cost_priority": 4.5003e-07
},
"databricks/databricks-gpt-5-nano": {
"cache_creation_input_token_cost": 4.998e-08,
@ -22469,7 +22582,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.25e-06,
"output_cost_per_token_priority": 2.625e-05,
"cache_read_input_token_cost_priority": 5.25e-07,
"cache_creation_input_token_cost_priority": 5.25e-06
},
"databricks/system.ai.kimi-k3": {
"litellm_provider": "databricks",
@ -22551,7 +22668,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false,
"thinking_always_on": true
"thinking_always_on": true,
"input_cost_per_token_priority": 3.85e-07,
"output_cost_per_token_priority": 3.85e-06,
"cache_read_input_token_cost_priority": 3.85e-07,
"cache_creation_input_token_cost_priority": 3.85e-07
},
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
"cache_creation_input_token_cost": 1.5001e-07,
@ -61865,6 +61986,39 @@
"supports_reasoning": true,
"supports_vision": false
},
"tencent/seedream-image-v5.0-flash": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.018,
"output_cost_per_image_token": 1.6e-06,
"source": "https://cloud.tencent.com/document/product/1823/130055",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"tencent/seedream-image-v5.0-lite": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.035,
"output_cost_per_image_token": 1.6e-06,
"source": "https://www.tencentcloud.com/document/product/1300/78937",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"tencent/seedream-image-v5.0-pro": {
"litellm_provider": "tencent",
"mode": "image_generation",
"output_cost_per_image": 0.09,
"output_cost_per_image_token": 1.6e-06,
"source": "https://www.tencentcloud.com/document/product/1300/78937",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
]
},
"cognition/swe-1.6": {
"input_cost_per_token": 5e-07,
"output_cost_per_token": 2.5e-06,
@ -77774,7 +77928,11 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 5.25e-07,
"output_cost_per_token_priority": 2.1e-06,
"cache_read_input_token_cost_priority": 5.25e-08,
"cache_creation_input_token_cost_priority": 5.25e-07
},
"databricks/databricks-deepseek-v4-1-flash-1017": {
"cache_creation_input_token_cost": 3.0002e-07,
@ -77897,7 +78055,11 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.50002e-06,
"output_cost_per_token_priority": 2.250003e-05,
"cache_read_input_token_cost_priority": 2.2498e-07,
"cache_creation_input_token_cost_priority": 4.50002e-06
},
"databricks/databricks-gpt-5-4-pro": {
"cache_creation_input_token_cost": 2.999997e-05,
@ -77959,7 +78121,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 1.99997e-05,
"cache_read_input_token_cost_priority": 1.9999e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-6-1-terra": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -77991,7 +78161,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.99997e-06,
"output_cost_per_token_priority": 9.99999e-06,
"cache_read_input_token_cost_priority": 1.0003e-07,
"cache_creation_input_token_cost_priority": 2.49998e-06,
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
},
"databricks/databricks-gpt-6-astra": {
"cache_creation_input_token_cost": 1.249997e-05,
@ -78023,7 +78201,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.999998e-05,
"output_cost_per_token_priority": 9.99999e-05,
"cache_read_input_token_cost_priority": 1.99997e-06,
"cache_creation_input_token_cost_priority": 2.500001e-05,
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
},
"databricks/databricks-gpt-6-luna": {
"cache_creation_input_token_cost": 1.2502e-07,
@ -78055,7 +78241,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 1.9999e-07,
"output_cost_per_token_priority": 1.0003e-06,
"cache_read_input_token_cost_priority": 2.002e-08,
"cache_creation_input_token_cost_priority": 2.4997e-07,
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
},
"databricks/databricks-gpt-6-sol": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -78087,7 +78281,15 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"input_cost_per_token_priority": 4.00001e-06,
"output_cost_per_token_priority": 1.99997e-05,
"cache_read_input_token_cost_priority": 3.9998e-07,
"cache_creation_input_token_cost_priority": 5.00003e-06,
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
},
"databricks/databricks-gpt-image-1": {
"cache_creation_input_token_cost": 5.00003e-06,

View file

@ -3793,8 +3793,8 @@
"responses": true,
"text_completion": true,
"embeddings": false,
"image_generations": false,
"image_edits": false,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": false,
"audio_speech": false,
"rerank": false,

View file

@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
@@index([unified_file_id])
@@index([flat_model_file_ids], type: Gin)
@@index([team_id, created_at(sort: Desc)])
@@index([created_by, created_at(sort: Desc)])
}
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport {
@@index([window_end])
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
}

View file

@ -0,0 +1,205 @@
---
name: litellm-setup
description: Set up a self-hosted LiteLLM proxy, the OpenAI-compatible LLM gateway, on this machine, from install through provider keys, config.yaml, a master key, optional Postgres and Admin UI, and a verified chat completion. Use when the user asks to set up LiteLLM, install the LiteLLM proxy, self-host an LLM gateway, or put one endpoint in front of OpenAI, Anthropic, Gemini, Azure, Bedrock, Ollama, or other providers. Not for using the LiteLLM Python SDK inside an app
---
# Set up a LiteLLM proxy
Take the user from nothing to a running LiteLLM proxy that has answered one chat completion. Run every command yourself unless a step says the user does it
Never print a secret. Provider keys, the master key, the salt key, the database URL, and virtual keys live only in `.env`, and `config.yaml` reads them through `os.environ/NAME`. Check that a value is there with `grep -c '^NAME=.' .env`, never with `cat .env` or `echo "$NAME"`. Load `.env` into a shell with this line, which exports each line as it is, so a `&` or `$` in a value stays literal, and first drops any database settings the shell already had and sets `LITELLM_MODE=PRODUCTION`, which stops `litellm` from also loading a `.env` it finds above its install folder, such as `~/.env`, so only this `.env` picks the database. The steps that need it start with it:
```bash
unset DATABASE_URL DIRECT_URL DATABASE_HOST; export LITELLM_MODE=PRODUCTION; while IFS= read -r line || [ -n "$line" ]; do case "$line" in [A-Za-z_]*=*) export "$line" ;; esac; done < .env
```
## 1. Ask once
Ask everything in one message, with defaults the user can accept by saying "go":
- Providers and models. Default is OpenAI with its newest flagship model, using the `OPENAI_API_KEY` in your environment when `[ -n "$OPENAI_API_KEY" ] && echo set` says it is there. Run that check before asking, and say in the question that the default uses that variable, naming it but never its value, since a key in your environment can belong to another account
- Whether they want virtual keys, teams, spend tracking, or the Admin UI. Those need Postgres, plain routing does not. Default is no
- uv (default) or Docker. With Postgres, Docker runs the repo's quickstart, which brings its own database, and uv needs a Postgres server they already have, with a new, empty database for LiteLLM, so someone with their own Postgres server picks uv
- Folder and port. Defaults are `./litellm-proxy` (`./litellm-gateway` for the quickstart) and 4000
Then follow the steps for their answers. Plain routing skips steps 3 and 6. The quickstart is steps 2, 3, 8, and 9. uv with their own Postgres skips step 3
In the same message, say that someone who picks uv and would rather answer prompts can run `litellm --setup` in a separate terminal window once LiteLLM is installed, not through Claude Code. Never run it yourself, since it needs a real terminal and without one prints "Setup cancelled." and writes nothing. Give them the command as `LITELLM_MODE=PRODUCTION <bin>/litellm --setup`, with `<bin>` replaced by the folder `uv tool dir --bin` prints, to run from the folder they want the config in. It writes `litellm_config.yaml` there with the provider keys and master key in plain text and offers to start the proxy, which then listens on every network interface and stops with Ctrl+C in that window. When the user is done, run this in that folder. It makes the file readable only by them, keeps it out of git (the wizard does neither), and prints only the model names. Never print that file or run `grep`, `sed`, or `cat` on it:
```bash
chmod 600 litellm_config.yaml; grep -qx 'litellm_config.yaml' .gitignore 2>/dev/null || echo 'litellm_config.yaml' >> .gitignore
"$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(*(m["model_name"] for m in yaml.safe_load(open("litellm_config.yaml"))["model_list"]), sep="\n")'
```
Then skip to step 8 and run it against their proxy, starting that shell command with this line instead of the loader line, which reads the master key without printing it:
```bash
export LITELLM_MASTER_KEY="$("$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(yaml.safe_load(open("litellm_config.yaml"))["general_settings"]["master_key"])')"
```
## 2. Check prerequisites
```bash
uv --version; docker info --format '{{.ServerVersion}}'; docker compose version; curl --version | head -1; openssl version; command -v python3
```
On native Windows, ask the user to run Claude Code inside WSL, since every command here is POSIX shell. Every path needs curl, openssl, and python3, so ask the user to install any that is missing. The uv path needs uv, which downloads a supported Python (LiteLLM supports 3.10 to 3.14) when the `--python` range in step 4 rules out the one on the machine. If uv is missing, offer to install it, and only after the user says yes run `curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh`, which leaves their shell profile alone (https://docs.astral.sh/uv/getting-started/installation/). It puts uv in `~/.local/bin`, so start every later command that runs `uv` or `litellm` with `export PATH="$HOME/.local/bin:$PATH";`. The Docker path needs a running Docker daemon, and the quickstart also needs Docker Compose v2. If `docker info` says `permission denied`, the user's account cannot use Docker without sudo. Offer the uv path, or ask them to add themselves to the `docker` group the way https://docs.docker.com/engine/install/linux-postinstall/ shows and start a new login session before you check again
Check the port is free with `curl -s -o /dev/null --max-time 2 http://127.0.0.1:4000; echo $?`, which prints `7` (could not connect) only when nothing listens there. Don't use `lsof` for this, since it misses ports that root or Docker hold. If it is taken, pick another port and use it everywhere below, because a proxy asked for a busy 4000 silently moves to a random port
## 3. Postgres with Docker: run the quickstart
The repo's quickstart starts the gateway, Postgres, and the Admin UI in Docker, listening on this machine only, and puts a `litellm-gateway` command in that folder that stops, starts, updates, or removes them. It asks nothing when it has no terminal. It writes the master key, salt key, and database password to its own `.env` and never prints them. Download it, then run it with the folder and port from step 1. It returns once the gateway is ready:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh -o quickstart.sh
LITELLM_DIR="$PWD/litellm-gateway" LITELLM_PORT=4000 sh quickstart.sh --yes
cd litellm-gateway
```
Models live in the database here, so add each one through the API instead of a `config.yaml`, naming it the way step 5 says. Pipe the request body in, so the provider key never shows up in output or on a command line. Unless the user agreed to the key in your environment, as step 5 says, ask them to add the line `OPENAI_API_KEY=<their key>` to this `.env` themselves. Run the loader line first:
```bash
python3 -c 'import json, os, sys; json.dump({"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": os.environ["OPENAI_API_KEY"]}}, sys.stdout)' \
| curl -sS -o /dev/null -w '%{http_code}\n' http://127.0.0.1:4000/model/new \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" -d @-
```
Run that in one shell command that starts with the loader line, since the master key is in `.env`. `200` means the model was added. Then go to step 8
## 4. Install
uv:
```bash
uv tool install --python '>=3.10,<3.15' 'litellm[proxy]'
export PATH="$(uv tool dir --bin):$PATH"; litellm --version
```
With their own Postgres, install `'litellm[proxy,extra_proxy]'` instead, with the same `--python` range, which adds the Prisma client the proxy uses for the database. Running `uv tool install` again with the new extras replaces the old install. Start every later command that runs `litellm` with `export PATH="$(uv tool dir --bin):$PATH";`, so this install runs even when the shell has no `litellm` or an older one from a venv or `pip install litellm`. `uv tool update-shell` puts it on the PATH of new shells by editing the user's shell profile, so run it only after the user says yes
Docker:
```bash
docker pull docker.litellm.ai/berriai/litellm:main-stable
```
## 5. Write `.env` and `config.yaml`
Run every command from here on in the folder from step 1. Create `.env` readable only by the user, keep it out of git, and add a master key unless it already has one. The proxy refuses to start without one, and `sk-` is the format of every key it issues:
```bash
mkdir -p litellm-proxy && cd litellm-proxy
touch .env && chmod 600 .env; grep -qx '.env' .gitignore 2>/dev/null || echo '.env' >> .gitignore
grep -q '^LITELLM_MASTER_KEY=.' .env || printf 'LITELLM_MASTER_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
```
Copy a provider key that is already in your environment only after the user agreed to that variable, in their step 1 answer or when you ask now, and copy it without printing it, as in `printf 'OPENAI_API_KEY=%s\n' "$OPENAI_API_KEY" >> .env`. Otherwise ask the user to add the line `OPENAI_API_KEY=<their key>` to `.env` themselves and tell you when it is there. If they paste a key into the chat anyway, write it to `.env` and never repeat it. Keep one `NAME=value` per line with no quotes, so `docker --env-file` and the loader line read the same value
Confirm each model's current LiteLLM name in the cost map instead of trusting memory, since model names go stale fast. This prints the chat models of the provider named last (`openai`, `anthropic`, `gemini`, `azure`, `bedrock`, `vertex_ai`, and so on) the way `config.yaml` takes them, with the provider in front. It matches by prefix, so `bedrock` also lists the newer Bedrock models and `vertex_ai` lists every Vertex family. For Ollama, use a model from `ollama list` instead:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json | python3 -c '
import json, sys
for name, row in json.load(sys.stdin).items():
provider = str(row.get("litellm_provider", "")) if isinstance(row, dict) else ""
if provider.startswith(sys.argv[1]) and row.get("mode") == "chat":
prefix = "bedrock" if provider == "bedrock_converse" else "vertex_ai" if provider.startswith("vertex_ai") else provider
print(name if name.startswith(prefix + "/") else prefix + "/" + name)
' openai
```
In `config.yaml`, `model` is a name from that list, as in `openai/gpt-6.1-sol` or `bedrock/us.anthropic.<name>`, and `model_name` is what clients will send. If the folder already has a `config.yaml`, show the user its `model_name` lines and ask before replacing it. This example uses the newest GPT model in the cost map when this skill was written, so replace it with what the user picked:
```yaml
model_list:
- model_name: gpt-6.1-sol
litellm_params:
model: openai/gpt-6.1-sol
api_key: os.environ/OPENAI_API_KEY
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
```
Add one entry per model. On Azure, ask the user for each deployment's name and the model it runs. Use `azure/<deployment name>` as `model`, since Azure routes by deployment and a name from the lookup works only when the deployment has that same name, and add `model_info:` with `base_model: azure/<model it runs>` to the entry, since LiteLLM prices each call by that name. Providers that need more than a key, like Azure or Bedrock, list their `litellm_params` on their page under https://docs.litellm.ai/docs/providers. Point every secret at `.env` through `os.environ/NAME` the same way
## 6. Their own Postgres, on the uv path
Ask the user to create a new, empty database for LiteLLM on their server (`CREATE DATABASE litellm;`) and add `DATABASE_URL=postgresql://...` for it to `.env` themselves. The proxy creates its tables there on first start and stops on a database that already holds other tables. If that URL goes through a transaction pooler such as PgBouncer, they also add `DIRECT_URL` for the same database without the pooler, which the migrations use. Then add a salt key. It encrypts the provider credentials the proxy stores in the database and falls back to the master key when unset, so set it once and never change it:
```bash
grep -q '^LITELLM_SALT_KEY=.' .env || printf 'LITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
```
Generate the Prisma client once inside the tool's environment. A pip or uv install doesn't ship it generated (Docker images do):
```bash
tool="$(uv tool dir)/litellm"
schema=$("$tool/bin/python" -I -c 'import importlib.util, os; print(os.path.join(importlib.util.find_spec("litellm").submodule_search_locations[0], "proxy", "schema.prisma"))')
PATH="$tool/bin:$PATH" prisma generate --schema "$schema"
```
## 7. Start the proxy
Bind to `127.0.0.1` unless the user wants other machines to reach it
uv, in one shell command that starts with the loader line, since the installed `litellm` does not read `.env` by itself:
```bash
export PATH="$HOME/.local/bin:$PATH"; export PATH="$(uv tool dir --bin):$PATH"
nohup litellm --config config.yaml --host 127.0.0.1 --port 4000 > litellm.log 2>&1 &
echo $! > litellm.pid
```
Docker:
```bash
docker run -d --name litellm --env-file .env -v "$PWD/config.yaml:/app/config.yaml:ro" -p 127.0.0.1:4000:4000 \
docker.litellm.ai/berriai/litellm:main-stable --config /app/config.yaml --port 4000
```
Wait on the health route, stopping early if the proxy exits. The first start with Postgres applies the database migrations and takes longer. Define `alive` for the path you took, `alive() { kill -0 "$(cat litellm.pid)" 2>/dev/null; }` for uv or `alive() { [ "$(docker inspect -f '{{.State.Running}}' litellm 2>/dev/null)" = true ]; }` for Docker, in the same shell command as the loop:
```bash
for i in $(seq 1 180); do curl -fsS http://127.0.0.1:4000/health/liveliness >/dev/null 2>&1 && break; alive || break; sleep 1; done
curl -sS http://127.0.0.1:4000/health/readiness
```
Readiness reports `"db":"connected"` with Postgres and `"db":"Not connected"` without it. If the proxy never comes up, read the end of `litellm.log` (or `docker logs litellm`) and check the troubleshooting list
## 8. Verify with a chat completion
Use a `model_name` from `config.yaml` (or step 3), in one shell command that starts with the loader line:
```bash
curl -sS http://127.0.0.1:4000/v1/chat/completions \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" \
-d '{"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "Say hello in five words."}]}'
```
Show the user the reply in `choices[0].message.content`. An error here means the setup is not done, so fix it before handing off
## 9. Hand off
Tell the user, without printing any secret:
- The base URL is `http://127.0.0.1:4000`. Any OpenAI-compatible client works with it as the base URL and the master key or a virtual key as the API key
- The master key is `LITELLM_MASTER_KEY` in the folder's `.env`, or `general_settings.master_key` in `litellm_config.yaml` when they used `litellm --setup`
- To stop it, `kill "$(cat litellm.pid)"` (uv), `docker stop litellm` (Docker), `<folder>/litellm-gateway stop` (quickstart), or Ctrl+C in its window (`litellm --setup`). To start it again, `docker start litellm`, `<folder>/litellm-gateway start`, on uv the whole step 7 command as you ran it, loader line and PATH export included, wrapped in `( ... )` so their shell keeps its own settings, shown as a code block and never paraphrased, since a new shell has neither, or for `litellm --setup` `LITELLM_MODE=PRODUCTION <bin>/litellm --config litellm_config.yaml --host 127.0.0.1 --port <port>` in a terminal window in that folder, with `<bin>` and `<port>` filled in, since running the wizard again writes a new config with a new master key. `update` and `uninstall` update or remove a quickstart install. Give the quickstart commands with the folder's full path, since a bare `litellm-gateway` is on PATH only when `~/.local/bin` is, and can belong to an earlier install
- With Postgres, the Admin UI is at `http://127.0.0.1:4000/ui` (user `admin`, password the master key), and `POST /key/generate` with the master key creates virtual keys with their own budgets and model lists. Without Postgres, both need step 3 or 6 first
- The Admin UI warns that the master-key login is a shared secret, which is fine while only this machine reaches the proxy. Before anyone else can reach it, invite an admin with their own password on the Internal Users page, then set `general_settings.disable_env_credential_login: true` in the proxy config
- To manage the running proxy from a coding agent (keys, teams, users, models, MCP servers), install `npx skills add BerriAI/litellm-skills -g`
## Troubleshooting
- `refused to start: no master key is set` means the loader line did not run in the shell command that started `litellm`, or `--env-file .env` is missing from `docker run`
- `a database URL is set but the prisma CLI is neither on PATH nor importable` means the install lacks `extra_proxy`, so rerun step 4 with it, or, on plain routing, that the command that started `litellm` did not start with the loader line
- `The Client hasn't been generated yet` or `Unable to find Prisma binaries` means the Prisma generate command in step 6 has not run
- `Incorrect API key provided` or another provider `AuthenticationError` means a provider key is wrong. Fix it in `.env` and restart, since the proxy reads the environment only at startup. On uv, stop it and rerun step 7. On Docker, `docker rm -f litellm` and rerun step 7, since `docker restart` keeps the environment the container was created with. On the quickstart the key is saved with the model, so delete the model and add it again: `GET /model/info` lists each model's `model_info.id`, and `POST /model/delete` with `{"id": "<id>"}` deletes it
- `Cannot automatically baseline this database` means `DATABASE_URL` points at a database that already holds tables LiteLLM did not create. Point it at a new, empty database
- `LLM Provider NOT provided` means a `model` in `config.yaml` (or step 3) lacks its provider in front. Use the name as the step 5 lookup prints it
- `Invalid model name passed in model=...` means the request's `model` is not a `model_name` in `config.yaml` or step 3. `GET /v1/models` lists the ones the key can use
- `No connected db.` on a proxy without Postgres means the bearer token is not the master key
- `Conflict. The container name "/litellm" is already in use` means an earlier run left a container. `docker start litellm` reuses it, or with the user's OK `docker rm -f litellm` and rerun step 7
- Inside Docker, `localhost` in `config.yaml` (an Ollama `api_base`, say) is the container itself. Use `host.docker.internal`, and on Linux also pass `--add-host=host.docker.internal:host-gateway`

View file

@ -175,6 +175,8 @@ ignored_function_names = [
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
"priority_header_update", # Tested through the priority header tests in test_common_request_processing.py
"_drops_params", # Read by split_request_priority and priority_header_update, tested through both
"_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py
"_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py
"_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py

View file

@ -1,3 +1,4 @@
import hashlib
import os
import queue
import secrets
@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key
_insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
_insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
_insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
_insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000)
ledger_before: Final = _ledger(scratch_database)
with pytest.raises(RuntimeError) as refused:
@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})"
@requires_db
def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is(
scratch_database: str,
) -> None:
"""Postgres compresses a long value before it checks the entry size, so 3000 repeated
bytes index fine while 3000 random hex bytes do not. The check must draw the same line."""
create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")')
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)')
conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500)))
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
)
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
conn.execute(create_index)
conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,))
conn.execute(create_index)
stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone()
assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored
(finding,) = findings
assert isinstance(finding, migration_data_checks.BlockedChange)
assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other"
assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)"
def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool:
conn.execute('DROP TABLE IF EXISTS "one_key"')
conn.execute('CREATE TABLE "one_key" ("k" TEXT)')
conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,))
try:
conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")')
except psycopg.errors.ProgramLimitExceeded:
return True
return False
@requires_db
def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None:
"""Keys that compress to a stored size just under and just over the entry limit go through the
check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it."""
noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40))
keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500))
create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")')
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)')
conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),))
compressed: Final = conn.execute(
'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704'
).fetchone()
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
)
rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed"
assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}"
(finding,) = findings
assert isinstance(finding, migration_data_checks.BlockedChange)
assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other"
@requires_db
def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:

View file

@ -0,0 +1,128 @@
import os
import subprocess
import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import psycopg
import pytest
from psycopg import sql
from litellm.accounting import DeliveryState, DeliveryTarget, EventKind
from tests.integration._support.database import scratch_database
_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)

View file

@ -0,0 +1,122 @@
import json
import uuid
from collections.abc import Mapping
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
MODEL: Final = "claude-haiku-4-5"
def _sse(event: Mapping[str, JsonValue]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def _message(identity: str, stop_reason: str | None) -> dict[str, JsonValue]:
return {
"id": "msg_" + identity,
"type": "message",
"role": "assistant",
"model": MODEL,
"content": [{"type": "text", "text": "alpha bravo"}],
"stop_reason": stop_reason,
"stop_sequence": None,
"usage": {"input_tokens": 199878, "output_tokens": 12799},
}
def _stream(identity: str, stop_reason: str) -> tuple[bytes, ...]:
return (
_sse({"type": "message_start", "message": {**_message(identity, None), "content": []}}),
_sse({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
_sse({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "alpha bravo"}}),
_sse({"type": "content_block_stop", "index": 0}),
_sse(
{
"type": "message_delta",
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
"usage": {"output_tokens": 12799},
}
),
_sse({"type": "message_stop"}),
)
def _anthropic(identity: str, stop_reason: str):
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/messages", request.target
if json.loads(request.body).get("stream"):
return Reply(content_type="text/event-stream", chunks=_stream(identity, stop_reason))
return Reply(body=json.dumps(_message(identity, stop_reason)).encode())
return respond
def _events(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
with gateway.client.stream(
"POST", path, json={**body, "stream": True}, headers={"Authorization": f"Bearer {gateway.key}"}
) as response:
assert response.status_code == 200, response.read().decode()
return tuple(
json.loads(line.removeprefix("data: "))
for line in response.iter_lines()
if line.startswith("data: ") and line != "data: [DONE]"
)
def _chat_finish_reason(gateway: Gateway, model: str, *, stream: bool) -> JsonValue:
body: Final = {"model": model, "max_tokens": 64000, "messages": [{"role": "user", "content": "copy the list"}]}
if stream:
reasons: Final = tuple(
choice["finish_reason"]
for chunk in _events(gateway, "/v1/chat/completions", body)
for choice in chunk["choices"] # comprehension-ok: chunk -> choices flatten
if choice.get("finish_reason")
)
assert len(reasons) == 1, reasons
return reasons[0]
response: Final = gateway.request("POST", "/v1/chat/completions", body)
assert response.status_code == 200, response.text
return response.json()["choices"][0]["finish_reason"]
def _responses_status(gateway: Gateway, model: str, *, stream: bool) -> JsonValue:
body: Final = {"model": model, "max_output_tokens": 64000, "input": "copy the list"}
if stream:
terminal: Final = tuple(
event["response"]
for event in _events(gateway, "/v1/responses", body)
if event["type"] in ("response.completed", "response.incomplete", "response.failed")
)
assert len(terminal) == 1, terminal
return terminal[0]["status"]
response: Final = gateway.request("POST", "/v1/responses", body)
assert response.status_code == 200, response.text
return response.json()["status"]
@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"])
@pytest.mark.parametrize(
("stop_reason", "finish_reason"),
[("model_context_window_exceeded", "length"), ("max_tokens", "length"), ("end_turn", "stop")],
)
def test_chat_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, finish_reason: str, stream: bool) -> None:
identity: Final = uuid.uuid4().hex
with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key")
assert _chat_finish_reason(gateway, model, stream=stream) == finish_reason
@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"])
@pytest.mark.parametrize(
("stop_reason", "status"),
[("model_context_window_exceeded", "incomplete"), ("max_tokens", "incomplete"), ("end_turn", "completed")],
)
def test_responses_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, status: str, stream: bool) -> None:
identity: Final = uuid.uuid4().hex
with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key")
assert _responses_status(gateway, model, stream=stream) == status

View file

@ -32,6 +32,7 @@ GEMINI: Final = "gemini-3-pro-preview"
ZAI_GLM: Final = "glm-5"
ECHOED_KEY_MODEL: Final = f"{KIMI}-echoed-key"
BEDROCK_CLAUDE: Final = f"us.anthropic.{CLAUDE}"
DATABRICKS_CLAUDE: Final = f"databricks-{CLAUDE}"
AZURE_API_VERSION: Final = "2025-04-01-preview"
SYNTHETIC_KEY: Final = "synthetic-provider-key"
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
@ -215,6 +216,10 @@ def backend(kind: str, wire_url: str) -> Backend:
f"bedrock/invoke/{BEDROCK_CLAUDE}",
MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY, "aws_region_name": "us-east-1"}),
)
case "databricks_claude":
return Backend(
f"databricks/{DATABRICKS_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY})
)
case _:
raise AssertionError(f"unknown backend kind {kind}")

View file

@ -0,0 +1,96 @@
from __future__ import annotations
import uuid
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import httpx
import pytest
from integration._support.client import gateway_from_environment
from integration._support.openai_wire import chat_reply
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._cache_control_marks_support import owned_config
from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
GROUP: Final = "daily-limited-group"
LIMITED_ID: Final = "free-tier-two-a-day"
UNLIMITED_ID: Final = "paid-tier"
RPD: Final = 2
EXTRA_REQUESTS: Final = 2
@dataclass(frozen=True, slots=True)
class Rig:
owned: OwnedProxy
limited: Wire
unlimited: Wire
def _reply(request: Request) -> Reply:
if is_model_info_probe(request):
return model_list_reply()
return chat_reply("chatcmpl-daily", OPENAI_MODEL, "served", stream=False)
def _deployment(model_id: str, wire: Wire, **params: JsonValue) -> dict[str, JsonValue]:
return {
"model_name": GROUP,
"litellm_params": {
"model": OPENAI_MODEL,
"api_base": wire.url + "/v1",
"api_key": "synthetic-openai-key",
**params,
},
"model_info": {"id": model_id},
}
def _config(directory: Path, limited: Wire, unlimited: Wire) -> Path:
return owned_config(
directory,
[_deployment(LIMITED_ID, limited, rpd=RPD, order=1), _deployment(UNLIMITED_ID, unlimited, order=2)],
litellm_settings={"cache": False},
router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 1},
)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("daily-limits")
with (
gateway_from_environment() as gateway,
wire_server(_reply) as limited,
wire_server(_reply) as unlimited,
owned_proxy_process(gateway, directory, {}, config=_config(directory, limited, unlimited)) as owned,
):
yield Rig(owned, limited, unlimited)
def _chat(rig: Rig) -> httpx.Response:
return rig.owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": GROUP, "messages": [{"role": "user", "content": f"daily {uuid.uuid4().hex}"}]},
)
def _forwarded(wire: Wire) -> int:
return sum(1 for request in wire.drain() if not is_model_info_probe(request))
def test_requests_past_the_deployment_rpd_are_served_by_the_other_deployment_in_the_group(rig: Rig) -> None:
responses: Final = tuple(_chat(rig) for _ in range(RPD + EXTRA_REQUESTS))
assert [response.status_code for response in responses] == [200] * (RPD + EXTRA_REQUESTS), [
response.text for response in responses
]
assert [response.headers["x-litellm-model-id"] for response in responses] == [LIMITED_ID] * RPD + [
UNLIMITED_ID
] * EXTRA_REQUESTS
assert (_forwarded(rig.limited), _forwarded(rig.unlimited)) == (RPD, EXTRA_REQUESTS)

View file

@ -1,4 +1,5 @@
import asyncio
import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
@ -8,9 +9,11 @@ import anthropic
import httpx
import openai
import pytest
from integration._support.anthropic_sse import message_json
from integration._support.client import Gateway, eventually, list_value, object_value, string_value
from integration._support.openai_wire import openai_error
from integration._support.wire import Request, wire_server
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, wire_server
from integration.routing._reasoning_effort_levels_support import (
AZURE_API_VERSION,
CLAUDE,
@ -34,6 +37,7 @@ from integration.routing._reasoning_effort_levels_support import (
marker_of,
messages_request,
new_marker,
path_of,
prompt,
provider_reply,
responses_request,
@ -41,6 +45,8 @@ from integration.routing._reasoning_effort_levels_support import (
)
from pydantic import JsonValue
from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS
CHAT: Final = "/v1/chat/completions"
MESSAGES: Final = "/v1/messages"
RESPONSES: Final = "/v1/responses"
@ -49,6 +55,7 @@ _ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"}
_STRING_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": "2048"}
_XHIGH_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 9000}
_BOOL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": True}
_FALSE_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": False}
_NULL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": None}
_LIST_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": [2048]}
_EMPTY_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": ""}
@ -67,6 +74,8 @@ _KIMI_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort
_CLAUDE_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(CLAUDE_LEVELS)})
_GPT_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(GPT_LEVELS)})
_UNDECLARED: Final[Mapping[str, JsonValue]] = MappingProxyType({})
_EMPTY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": []})
_NONE_ONLY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": ["none"]})
_ECHOED_MAP_KEY: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"key": "azure/chat-latest", "litellm_provider": "azure"}
)
@ -134,7 +143,7 @@ def _invoke_chat(effort: str, thinking: Mapping[str, JsonValue], betas: tuple[st
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt(marker)}]}],
"thinking": dict(thinking),
"output_config": {"effort": effort},
"max_tokens": 4096,
"max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS,
"anthropic_version": "bedrock-2023-05-31",
"anthropic_beta": list(betas),
}
@ -144,6 +153,14 @@ def _with_output_config(expected: Expected, output_config: JsonValue) -> Expecte
return lambda marker: {**expected(marker), "output_config": output_config}
def _with_fields(expected: Expected, fields: Mapping[str, JsonValue]) -> Expected:
return lambda marker: {**expected(marker), **fields}
def _without(expected: Expected, field: str) -> Expected:
return lambda marker: {key: value for key, value in expected(marker).items() if key != field}
def _converse_chat(effort: str, thinking: Mapping[str, JsonValue]) -> Expected:
return lambda marker: {
"messages": [{"role": "user", "content": [{"text": prompt(marker)}]}],
@ -320,6 +337,7 @@ _OPENAI_RESPONSES: Final = "/v1/responses"
_AZURE_RESPONSES: Final = f"/openai/responses?api-version={AZURE_API_VERSION}"
_UNSUPPORTED: Final = "does not support parameters"
_GEMINI: Final = f"/models/{GEMINI}:generateContent"
_DATABRICKS: Final = "/chat/completions"
CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
{
@ -1750,6 +1768,139 @@ CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
_anthropic_chat_without_effort(_ADAPTIVE),
_ANTHROPIC,
),
"chat-anthropic-false-budget-400": Cell(
CHAT, "anthropic", _CLAUDE_INFO, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
),
"chat-anthropic-undeclared-null-budget-400": Cell(
CHAT, "anthropic", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
),
"chat-anthropic-undeclared-false-budget-drop-params-sends-adaptive-without-effort": Cell(
CHAT,
"anthropic",
_UNDECLARED,
{"thinking": _FALSE_BUDGET, "drop_params": True},
_anthropic_chat_without_effort(_ADAPTIVE),
_ANTHROPIC,
),
"chat-anthropic-undeclared-missing-budget-sends-low": Cell(
CHAT,
"anthropic",
_UNDECLARED,
{"thinking": _MISSING_BUDGET},
_anthropic_chat("low", _ADAPTIVE),
_ANTHROPIC,
),
"messages-anthropic-undeclared-false-budget-400": Cell(
MESSAGES, "anthropic", _UNDECLARED, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
),
"messages-anthropic-empty-string-budget-400": Cell(
MESSAGES, "anthropic", _CLAUDE_INFO, {"thinking": _EMPTY_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
),
"messages-anthropic-undeclared-null-budget-drop-params-sends-adaptive-without-effort": Cell(
MESSAGES,
"anthropic",
_UNDECLARED,
{"thinking": _NULL_BUDGET, "drop_params": True},
_native_messages(None, _ADAPTIVE),
_ANTHROPIC,
),
"responses-anthropic-false-budget-400": Cell(
RESPONSES,
"anthropic",
_CLAUDE_INFO,
{"thinking": _FALSE_BUDGET},
None,
_ANTHROPIC,
400,
_NON_NUMERIC_BUDGET,
),
"chat-converse-undeclared-false-budget-400": Cell(
CHAT,
"bedrock_converse",
_UNDECLARED,
{"thinking": _FALSE_BUDGET},
None,
_CONVERSE,
400,
_NON_NUMERIC_BUDGET,
),
"messages-invoke-undeclared-null-budget-400": Cell(
MESSAGES, "bedrock_invoke", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _INVOKE, 400, _NON_NUMERIC_BUDGET
),
"chat-databricks-claude-undeclared-false-budget-400": Cell(
CHAT,
"databricks_claude",
_UNDECLARED,
{"thinking": _FALSE_BUDGET},
None,
_DATABRICKS,
400,
_NON_NUMERIC_BUDGET,
),
"chat-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
CHAT,
"anthropic",
_EMPTY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_anthropic_chat_without_effort(_ADAPTIVE),
_ANTHROPIC,
),
"chat-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
CHAT,
"anthropic",
_NONE_ONLY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_anthropic_chat_without_effort(_ADAPTIVE),
_ANTHROPIC,
),
"chat-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
CHAT,
"anthropic",
_EMPTY_DECLARATION,
{"thinking": _BUDGET},
_anthropic_chat("medium", _ADAPTIVE),
_ANTHROPIC,
),
"chat-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
CHAT,
"bedrock_converse",
_EMPTY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_converse_chat_without_effort(_ADAPTIVE),
_CONVERSE,
),
"messages-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
MESSAGES,
"anthropic",
_EMPTY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_native_messages(None, _ADAPTIVE),
_ANTHROPIC,
),
"messages-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
MESSAGES,
"anthropic",
_NONE_ONLY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_native_messages(None, _ADAPTIVE),
_ANTHROPIC,
),
"messages-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
MESSAGES,
"anthropic",
_EMPTY_DECLARATION,
{"thinking": _BUDGET},
_native_messages({"effort": "medium"}, _ADAPTIVE),
_ANTHROPIC,
),
"messages-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
MESSAGES,
"bedrock_converse",
_EMPTY_DECLARATION,
{"thinking": _BUDGET, "drop_params": True},
_converse_messages(None),
_CONVERSE,
),
}
)
@ -1815,20 +1966,65 @@ def _bedrock_native_schema(body: Mapping[str, JsonValue], *, invoke: bool) -> Js
return JSON_OBJECT.validate_json(string_value(json_schema["schema"]))
_CONVERSE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
"textFormat": {
"type": "json_schema",
"structure": {"jsonSchema": {"name": "answer", "schema": json.dumps(_NATIVE_OUTPUT_SCHEMA)}},
}
}
_INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
"effort": "high",
"format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA},
}
_JSON_TOOL: Final = "json_tool_call"
_CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {
key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties"
}
_CONVERSE_WITH_NATIVE_OUTPUT: Final = _with_fields(
_converse_chat("high", _ADAPTIVE), {"outputConfig": _CONVERSE_NATIVE_OUTPUT}
)
_INVOKE_WITH_NATIVE_OUTPUT: Final = _with_output_config(
_invoke_chat("high", _ADAPTIVE, ("structured-outputs-2025-11-13",)), _INVOKE_NATIVE_OUTPUT
)
_CONVERSE_WITH_JSON_TOOL: Final = _with_fields(
_converse_chat("high", _ADAPTIVE),
{
"toolConfig": {
"tools": [
{
"toolSpec": {
"name": _JSON_TOOL,
"description": _JSON_TOOL,
"inputSchema": {"json": _CONVERSE_TOOL_SCHEMA},
}
}
]
}
},
)
_INVOKE_WITH_JSON_TOOL: Final = _with_fields(
_without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"),
{"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]},
)
@pytest.mark.parametrize(
("kind", "target", "fields"),
("kind", "target", "fields", "expected"),
[
pytest.param("bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, id="converse-budget"),
pytest.param(
"bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, _CONVERSE_WITH_NATIVE_OUTPUT, id="converse-budget"
),
pytest.param(
"bedrock_invoke",
_INVOKE,
{"thinking": _BUDGET, "output_config": {"effort": "medium"}},
_INVOKE_WITH_NATIVE_OUTPUT,
id="invoke-output-config-beside-budget",
),
],
)
def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_declared_level(
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue]
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue], expected: Expected
) -> None:
marker: Final = new_marker()
response_format: Final[dict[str, JsonValue]] = {
@ -1844,10 +2040,116 @@ def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_decla
assert response.status_code == 200, response.text[:2000]
assert answer(marker) in response.text, response.text[:2000]
assert [(request.method, request.target) for request in received] == [("POST", target)]
body: Final = body_of(received[0])
assert _bedrock_native_schema(body, invoke=target == _INVOKE) == _NATIVE_OUTPUT_SCHEMA, body
assert {"toolConfig", "tools"}.isdisjoint(body), body
assert sent_effort(received[0]) == "high", body
assert body_of(received[0]) == expected(marker)
_NATIVE_OUTPUT_REJECTION: Final = (
"The model returned the following errors: output_config.format: Extra inputs are not permitted"
)
def _rejecting_native_output(request: Request) -> Reply:
native: Final = _bedrock_native_schema(body_of(request), invoke=path_of(request).endswith("/invoke"))
if native is None:
return provider_reply(request)
return Reply(status=400, body=json.dumps({"message": _NATIVE_OUTPUT_REJECTION}).encode())
@pytest.mark.parametrize(
("kind", "target", "rejected", "retried"),
[
pytest.param(
"bedrock_converse", _CONVERSE, _CONVERSE_WITH_NATIVE_OUTPUT, _CONVERSE_WITH_JSON_TOOL, id="converse"
),
pytest.param("bedrock_invoke", _INVOKE, _INVOKE_WITH_NATIVE_OUTPUT, _INVOKE_WITH_JSON_TOOL, id="invoke"),
],
)
def test_the_retry_after_bedrock_rejects_native_output_keeps_the_declared_level(
gateway: Gateway, kind: str, target: str, rejected: Expected, retried: Expected
) -> None:
marker: Final = new_marker()
response_format: Final[dict[str, JsonValue]] = {
"type": "json_schema",
"json_schema": {"name": "answer", "schema": _NATIVE_OUTPUT_SCHEMA},
}
with wire_server(_rejecting_native_output) as wire, gateway.scenario() as scenario:
model: Final = deployment(scenario, kind, wire, _CLAUDE_INFO_WITH_NATIVE_OUTPUT)
response: Final = gateway.request(
"POST", CHAT, chat_request(model, marker, response_format=response_format, thinking=_BUDGET)
)
received: Final = wire.drain()
assert response.status_code == 200, response.text[:2000]
assert [(request.method, request.target) for request in received] == [("POST", target), ("POST", target)]
assert [body_of(request) for request in received] == [rejected(marker), retried(marker)]
_VERTEX_PROJECT: Final = "scripted-project"
_VERTEX_LOCATION: Final = "us-east5"
_VERTEX_TOKEN: Final = "/_oauth/token"
_VERTEX_RAW_PREDICT: Final = (
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{CLAUDE}:rawPredict"
)
_VERTEX_ACCESS_TOKEN: Final = json.dumps({"access_token": "scripted", "token_type": "Bearer", "expires_in": 3600})
def _vertex(expected: Expected) -> Expected:
return _with_fields(_without(expected, "model"), {"anthropic_version": "vertex-2023-10-16"})
def _vertex_reply(request: Request) -> Reply:
if request.target == _VERTEX_TOKEN:
return Reply(body=_VERTEX_ACCESS_TOKEN.encode())
marker: Final = marker_of(request)
return Reply(body=message_json(f"msg_{marker}", CLAUDE, answer(marker)))
@pytest.mark.parametrize(
("route", "fields", "expected", "status"),
[
pytest.param(
CHAT,
{"reasoning_effort": "medium"},
_vertex(_anthropic_chat("high", {"type": "adaptive", "display": "summarized"})),
200,
id="chat-medium-to-high",
),
pytest.param(
MESSAGES,
{"thinking": _BUDGET},
_vertex(_native_messages({"effort": "high"}, _ADAPTIVE)),
200,
id="messages-budget-to-high",
),
pytest.param(CHAT, {"thinking": _FALSE_BUDGET}, None, 400, id="chat-false-budget-400"),
],
)
def test_a_vertex_ai_claude_deployment_gets_one_of_its_declared_levels(
gateway: Gateway, route: str, fields: dict[str, JsonValue], expected: Expected | None, status: int
) -> None:
marker: Final = new_marker()
with wire_server(_vertex_reply) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"vertex_ai/{CLAUDE}",
api_base=wire.url,
api_key=None,
vertex_project=_VERTEX_PROJECT,
vertex_location=_VERTEX_LOCATION,
vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url),
model_info=dict(_CLAUDE_INFO),
)
body: Final = (
chat_request(model, marker, **fields) if route == CHAT else messages_request(model, marker, **fields)
)
response: Final = gateway.request("POST", route, body)
received: Final = tuple(request for request in wire.drain() if request.target != _VERTEX_TOKEN)
assert response.status_code == status, response.text[:2000]
if expected is None:
assert _NON_NUMERIC_BUDGET in response.text, response.text[:2000]
assert received == (), [request.body[:300] for request in received]
return
assert answer(marker) in response.text, response.text[:2000]
assert [(request.method, request.target) for request in received] == [("POST", _VERTEX_RAW_PREDICT)]
assert body_of(received[0]) == expected(marker)
def _single(received: tuple[Request, ...], marker: str) -> Request:

View file

@ -0,0 +1,67 @@
"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose
``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the
caller's dict is left as the caller wrote it (LIT-10239)"""
import json
from collections.abc import Mapping
from typing import Final
import pytest
from integration._support.client import eventually
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
import litellm
_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0"
_PROMPT_TOKENS: Final = 5
_COMPLETION_TOKENS: Final = 1
_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1}
_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two")
def _bedrock(request: Request) -> Reply:
"""Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named."""
guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"]
body: Final[dict[str, JsonValue]] = {
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6},
"metrics": {"latencyMs": 3},
"trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}},
}
return Reply(body=json.dumps(body).encode())
def _guardrail_cost() -> float:
rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
return sum(units * rates[counter] for counter, units in _USAGE.items())
def _llm_cost() -> float:
prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")]
return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"]
def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None:
billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed
litellm.success_callback.append(
lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"])
)
shared_metadata: Final = {"tag": "x"}
with wire_server(_bedrock) as wire:
for guardrail_id in _GUARDRAIL_IDS:
litellm.completion(
model=_MODEL,
messages=[{"role": "user", "content": "hello"}],
api_base=wire.url,
aws_access_key_id="AKIASCRIPTED",
aws_secret_access_key="scripted",
aws_region_name="us-east-1",
guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"},
metadata=shared_metadata,
)
costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30)
assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS))
assert shared_metadata == {"tag": "x"}

View file

@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm import Router
from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments
from litellm.types.router import DeploymentTypedDict
_MODEL: Final = "gpt-5.4"
@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No
"/primary/chat/completions",
"/backup/chat/completions",
)
@pytest.mark.parametrize(
"router_kwargs",
[{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}],
ids=["string_fallback_target", "string_default_fallbacks"],
)
@pytest.mark.asyncio
async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None:
with wire_server(_peer) as wire:
router: Final = Router(
model_list=[
_deployment("primary", f"{wire.url}/primary", "primary-deployment"),
_deployment("backup", f"{wire.url}/backup", "backup-deployment"),
],
allowed_fails=0,
num_retries=0,
**router_kwargs,
)
with pytest.raises(litellm.AuthenticationError):
await router.acompletion(
model="primary",
messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}],
)
cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None)
assert tuple(cooled) == ("primary-deployment",)

View file

@ -413,3 +413,71 @@ def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) ->
second: Final = timed(lambda: single_worker.proxy.request("POST", path, following))
assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text)
assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5)
@pytest.mark.parametrize("placement", ("guardrail", "global"))
def test_skip_tool_message_in_guardrail_forwards_tool_results(tmp_path: Path, placement: str) -> None:
def configure(configuration: dict[str, JsonValue]) -> None:
guardrails: Final = configuration["guardrails"]
settings: Final = configuration["litellm_settings"]
assert isinstance(guardrails, list) and isinstance(guardrails[0], dict) and isinstance(settings, dict)
params: Final = guardrails[0]["litellm_params"]
assert isinstance(params, dict)
(params if placement == "guardrail" else settings)["skip_tool_message_in_guardrail"] = True
with tool_rig(tmp_path, configure=configure) as rig:
rig.wire.drain()
replies: Final = (
rig.proxy.request(
"POST",
"/v1/messages",
{
"model": "claude",
"max_tokens": 32,
"cache": {"no-cache": True},
"messages": [
{"role": "user", "content": TEXT},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"env": TEXT}}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": TEXT}],
},
],
},
),
rig.proxy.request(
"POST",
"/v1/chat/completions",
{
"model": "gpt",
"cache": {"no-cache": True},
"messages": [
{"role": "user", "content": TEXT},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "read_file", "arguments": json.dumps({"env": TEXT})},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": TEXT},
],
},
),
)
assert all(reply.status_code == 200 for reply in replies), [reply.text for reply in replies]
sent: Final = rig.wire.drain()
assert len(sent) == 2
for request in sent:
messages = json.loads(request.body)["messages"]
assert SECRET not in json.dumps(messages[:2])
assert SECRET in json.dumps(messages[2])

View file

@ -0,0 +1,263 @@
import json
import uuid
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from itertools import count, islice
from pathlib import Path
from types import MappingProxyType
from typing import Final
from urllib.parse import parse_qs, urlsplit
import pytest
import yaml
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from anthropic import Anthropic
from openai import OpenAI
from pydantic import BaseModel, ConfigDict, JsonValue
pytestmark: Final = pytest.mark.shares_worker_state
UNMANAGED_FILE: Final = "file-uploaded-straight-to-the-provider"
UPLOAD: Final = b"isolation check\n"
@dataclass(frozen=True, slots=True)
class Deployment:
provider: str
model: str
litellm_model: str
api_key: str
api_path: str
OPENAI: Final = Deployment("openai", "gpt-4.1-managed-listing", "openai/gpt-4.1", "sk-managed-listing", "/v1")
ANTHROPIC: Final = Deployment(
"anthropic", "claude-managed-listing", "anthropic/claude-haiku-5-5", "sk-ant-managed-listing", ""
)
DEPLOYMENTS: Final = (pytest.param(OPENAI, id="openai"), pytest.param(ANTHROPIC, id="anthropic"))
@dataclass(frozen=True, slots=True)
class ListCall:
path: str
params: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
headers: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
LIST_FORMS: Final = (
pytest.param(lambda d: ListCall(f"/{d.provider}/v1/files"), id="provider-path"),
pytest.param(lambda d: ListCall("/v1/files", params={"provider": d.provider}), id="provider-query"),
pytest.param(lambda d: ListCall("/v1/files", headers={"custom-llm-provider": d.provider}), id="provider-header"),
pytest.param(lambda d: ListCall("/v1/files", params={"target_model_names": d.model}), id="target-model-names"),
pytest.param(lambda d: ListCall("/v1/files", headers={"x-litellm-model": d.model}), id="model-header"),
)
class FilePage(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
data: list[dict[str, JsonValue]]
has_more: bool
@property
def ids(self) -> list[str]:
return [str(entry["id"]) for entry in self.data]
def _openai_file(file_id: str) -> dict[str, JsonValue]:
return {
"id": file_id,
"object": "file",
"bytes": len(UPLOAD),
"created_at": 1700000000,
"filename": "notes.txt",
"purpose": "user_data",
"status": "processed",
}
def _anthropic_file(file_id: str) -> dict[str, JsonValue]:
return {
"id": file_id,
"type": "file",
"filename": "notes.txt",
"mime_type": "text/plain",
"size_bytes": len(UPLOAD),
"created_at": "2026-01-01T00:00:00Z",
"downloadable": False,
}
def _reply(body: Mapping[str, JsonValue]) -> Reply:
return Reply(body=json.dumps(body).encode())
def _list_reply(provider: str, file_ids: tuple[str, ...]) -> Reply:
if provider == ANTHROPIC.provider:
return _reply(
{
"data": [_anthropic_file(file_id) for file_id in file_ids],
"has_more": False,
"first_id": file_ids[0],
"last_id": file_ids[-1],
}
)
return _reply({"object": "list", "data": [_openai_file(file_id) for file_id in file_ids]})
def _responder(provider_files: Mapping[str, tuple[str, ...]]) -> Callable[[Request], Reply]:
uploads: Final = {provider: count() for provider in provider_files}
def respond(request: Request) -> Reply:
target: Final = urlsplit(request.target)
path: Final = target.path
provider: Final = ANTHROPIC.provider if "x-api-key" in request.headers else OPENAI.provider
files: Final = provider_files[provider]
if (request.method, path) == ("GET", "/v1/models"):
return _reply({"object": "list", "data": []})
if (request.method, path) == ("POST", "/v1/files"):
uploaded: Final = files[next(uploads[provider])]
return _reply(_anthropic_file(uploaded) if provider == ANTHROPIC.provider else _openai_file(uploaded))
if (request.method, path) == ("GET", "/v1/files"):
limit: Final = int(parse_qs(target.query).get("limit", ["10000"])[0])
return _list_reply(provider, (UNMANAGED_FILE, *reversed(files))[:limit])
return Reply(status=404, body=json.dumps({"error": {"message": f"unscripted {request.target}"}}).encode())
return respond
def _config(wire: Wire) -> dict[str, JsonValue]:
return {
"model_list": [
{
"model_name": deployment.model,
"litellm_params": {
"model": deployment.litellm_model,
"api_key": deployment.api_key,
"api_base": f"{wire.url}{deployment.api_path}",
},
}
for deployment in (OPENAI, ANTHROPIC)
],
"litellm_settings": {"require_managed_files": True},
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"},
}
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
owner: str
intruder: str
files_by_provider: Mapping[str, tuple[str, ...]]
every_file: tuple[str, ...]
@property
def base_url(self) -> str:
return str(self.proxy.client.base_url).rstrip("/")
def list(self, call: ListCall, key: str | None = None) -> FilePage:
response: Final = self.proxy.request("GET", call.path, key=key, params=call.params, headers=call.headers)
assert response.status_code == 200, response.text
return FilePage.model_validate_json(response.text)
def _upload(proxy: Gateway, key: str, deployments: tuple[Deployment, ...]) -> str:
with OpenAI(base_url=f"{str(proxy.client.base_url).rstrip('/')}/openai/v1", api_key=key, max_retries=0) as client:
uploaded: Final = client.files.with_raw_response.create(
file=("notes.txt", UPLOAD, "text/plain"),
purpose="user_data",
extra_body={"target_model_names": ",".join(deployment.model for deployment in deployments)},
)
assert uploaded.status_code == 200, uploaded.text
return uploaded.parse().id
@contextmanager
def _rig(directory: Path) -> Iterator[Rig]:
build: Final = uuid.uuid4().hex
provider_files: Final = {
OPENAI.provider: (f"file-older-{build}", f"file-newer-{build}"),
ANTHROPIC.provider: (f"file_older_{build}", f"file_newer_{build}"),
}
with gateway_from_environment() as shared, wire_server(_responder(provider_files)) as wire:
config: Final = directory / "config.yaml"
config.write_text(yaml.safe_dump(_config(wire)))
with owned_proxy(shared, directory, {}, config=config) as proxy, proxy.scenario() as scenario:
owner: Final = scenario.key(user_id=scenario.user())
intruder: Final = scenario.key(user_id=scenario.user())
on_both: Final = _upload(proxy, owner, (OPENAI, ANTHROPIC))
on_openai: Final = _upload(proxy, owner, (OPENAI,))
on_anthropic: Final = _upload(proxy, owner, (ANTHROPIC,))
yield Rig(
proxy,
owner,
intruder,
MappingProxyType(
{OPENAI.provider: (on_openai, on_both), ANTHROPIC.provider: (on_anthropic, on_both)}
),
(on_anthropic, on_openai, on_both),
)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
with _rig(tmp_path_factory.mktemp("managed-file-listing")) as built:
yield built
@pytest.mark.parametrize("deployment", DEPLOYMENTS)
@pytest.mark.parametrize("form", LIST_FORMS)
def test_another_key_lists_none_of_the_files_on_the_shared_key(
rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall]
) -> None:
page: Final = rig.list(form(deployment), key=rig.intruder)
assert (page.ids, page.has_more) == ([], False)
@pytest.mark.parametrize("deployment", DEPLOYMENTS)
@pytest.mark.parametrize("form", LIST_FORMS)
def test_owner_lists_their_own_managed_files_on_that_provider_or_model_newest_first(
rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall]
) -> None:
page: Final = rig.list(form(deployment), key=rig.owner)
assert (page.ids, page.has_more) == (list(rig.files_by_provider[deployment.provider]), False)
def test_owner_lists_every_managed_file_when_no_provider_or_model_is_named(rig: Rig) -> None:
page: Final = rig.list(ListCall("/v1/files"), key=rig.owner)
assert (page.ids, page.has_more) == (list(rig.every_file), False)
@pytest.mark.parametrize("cursor_param", ["after", "after_id"])
def test_owner_pages_through_a_provider_list_with_limit_and_a_cursor(rig: Rig, cursor_param: str) -> None:
first: Final = rig.list(ListCall("/openai/v1/files", params={"limit": "1"}), key=rig.owner)
second: Final = rig.list(
ListCall("/openai/v1/files", params={"limit": "1", cursor_param: first.ids[0]}), key=rig.owner
)
expected: Final = rig.files_by_provider[OPENAI.provider]
assert [(page.ids, page.has_more) for page in (first, second)] == [
([expected[0]], True),
([expected[1]], False),
]
def test_anthropic_sdk_auto_pages_the_owners_files_to_the_end(rig: Rig) -> None:
with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client:
listed: Final = [file.id for file in islice(client.beta.files.list(limit=1), 5)]
assert listed == list(rig.files_by_provider[ANTHROPIC.provider])
def test_anthropic_sdk_pages_backwards_to_the_newest_file_and_stops(rig: Rig) -> None:
newest_first: Final = rig.files_by_provider[ANTHROPIC.provider]
with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client:
listed: Final = [file.id for file in islice(client.beta.files.list(limit=1, before_id=newest_first[-1]), 5)]
assert listed == list(reversed(newest_first[:-1]))
def test_master_key_still_lists_the_files_on_the_shared_provider_key(rig: Rig) -> None:
listed: Final = rig.list(ListCall(f"/{ANTHROPIC.provider}/v1/files")).ids
assert UNMANAGED_FILE in listed

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -26,7 +26,7 @@
},
"results": [
{
"breakdown": {
"breakdown": {"guardrails": {},
"api_keys": {
"key-a": {
"metadata": {

View file

@ -26,7 +26,7 @@
},
"results": [
{
"breakdown": {
"breakdown": {"guardrails": {},
"api_keys": {
"key-a": {
"metadata": {

View file

@ -26,7 +26,7 @@
},
"results": [
{
"breakdown": {
"breakdown": {"guardrails": {},
"api_keys": {
"key-a": {
"metadata": {

View file

@ -26,7 +26,7 @@
},
"results": [
{
"breakdown": {
"breakdown": {"guardrails": {},
"api_keys": {
"key-a": {
"metadata": {

View file

@ -1,6 +1,7 @@
import asyncio
import base64
import json
import math
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
@ -10,9 +11,18 @@ from typing import Final
import httpx
import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value
from integration._support.client import (
Gateway,
Scenario,
delete_key_if_present,
eventually,
held,
list_value,
object_value,
string_value,
)
from integration._support.database import read_rows, write_rows
from integration._support.otlp_sink import decode_spans
from integration._support.otlp_sink import Span, decode_spans
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Request, wire_server
from integration.caching.response_cache_case import entry_key, redis_store
@ -859,7 +869,7 @@ LANGFUSE_COST: Final = "langfuse.observation.cost_details"
_LANGFUSE_COST_DETAILS: Final = TypeAdapter(dict[str, float])
def _langfuse_callback(host: str) -> dict[str, JsonValue]:
def _langfuse_callback(host: str, **callback_vars: JsonValue) -> dict[str, JsonValue]:
return {
"callback_name": "langfuse",
"callback_type": "success",
@ -867,12 +877,17 @@ def _langfuse_callback(host: str) -> dict[str, JsonValue]:
"langfuse_public_key": LANGFUSE_PUBLIC_KEY,
"langfuse_secret_key": LANGFUSE_SECRET_KEY,
"langfuse_host": host,
**callback_vars,
},
}
def _logging_to(langfuse_url: str, **callback_vars: JsonValue) -> dict[str, JsonValue]:
return {"logging": [_langfuse_callback(langfuse_url, **callback_vars)]}
def _key_with_its_own_logging(scenario: Scenario, langfuse_url: str) -> str:
return scenario.key(metadata={"logging": [_langfuse_callback(langfuse_url)]})
return scenario.key(metadata=_logging_to(langfuse_url))
def _key_in_a_team_with_logging(scenario: Scenario, langfuse_url: str) -> str:
@ -892,17 +907,33 @@ def _trace_exports(langfuse: Recorder) -> tuple[Request, ...]:
return tuple(request for request in langfuse.requests() if request.target.startswith(LANGFUSE_TRACES_PATH))
def _exported_costs(langfuse: Recorder) -> tuple[float, ...]:
spans: Final = tuple(
def _exported_spans(langfuse: Recorder) -> tuple[Span, ...]:
return tuple(
chain.from_iterable(
decode_spans(export.body, export.headers.get("content-type", "")) for export in _trace_exports(langfuse)
)
)
return tuple(
_LANGFUSE_COST_DETAILS.validate_json(string_value(span["attributes"][LANGFUSE_COST]))["total"]
for span in spans
if LANGFUSE_COST in span["attributes"]
)
def _span_cost(span: Span) -> float | None:
cost_details: Final = span["attributes"].get(LANGFUSE_COST)
return None if cost_details is None else _LANGFUSE_COST_DETAILS.validate_json(string_value(cost_details))["total"]
def _exported_costs(langfuse: Recorder) -> tuple[float, ...]:
return tuple(cost for cost in map(_span_cost, _exported_spans(langfuse)) if cost is not None)
def _priced_spans(langfuse: Recorder) -> tuple[Span, ...]:
return tuple(span for span in _exported_spans(langfuse) if (_span_cost(span) or 0.0) > 0)
def _billed_once(langfuse: Recorder, seconds: float) -> Span:
priced: Final = eventually(lambda: _priced_spans(langfuse), lambda spans: len(spans) > 0, seconds=seconds)
costs: Final = tuple(_span_cost(span) or 0.0 for span in priced)
assert len(costs) == 1, costs
assert math.isclose(costs[0], PRICE), costs
return priced[0]
@pytest.mark.timeout(240)
@ -925,3 +956,133 @@ def test_the_billed_cost_reaches_the_langfuse_project_set_on_the_creating_key_or
assert [cost for cost in exported if cost > 0] == [pytest.approx(PRICE)]
credentials: Final = base64.b64encode(f"{LANGFUSE_PUBLIC_KEY}:{LANGFUSE_SECRET_KEY}".encode()).decode()
assert {export.headers.get("authorization") for export in _trace_exports(langfuse)} == {f"Basic {credentials}"}
@pytest.mark.timeout(240)
def test_the_billed_cost_reaches_the_creating_keys_langfuse_project_and_never_the_readers(gateway: Gateway) -> None:
with (
wire_server(langfuse_sink) as creator_wire,
wire_server(langfuse_sink) as reader_wire,
gateway.scenario() as scenario,
):
creator_langfuse: Final = Recorder(creator_wire)
reader_langfuse: Final = Recorder(reader_wire)
team_id: Final = scenario.team()
user_id: Final = scenario.member(team_id)
creator: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(creator_wire.url))
reader: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(reader_wire.url))
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(scenario, creator, upstream.model)
reads: Final = tuple(read_response(gateway, reader, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
_billed_once(creator_langfuse, SPEND_LANDING_SECONDS)
eventually(lambda: _trace_exports(reader_langfuse), lambda exports: len(exports) > 0)
assert _priced_spans(reader_langfuse) == ()
MESSAGE_LOGGING: Final[Mapping[str, Mapping[str, JsonValue]]] = {
"logged": {},
"redacted": {"turn_off_message_logging": "true"},
}
@pytest.mark.timeout(240)
@pytest.mark.parametrize("message_logging", MESSAGE_LOGGING.keys())
def test_the_billed_observation_carries_the_response_text_only_when_the_creating_key_keeps_message_logging(
gateway: Gateway, message_logging: str
) -> None:
with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario:
langfuse: Final = Recorder(wire)
key: Final = scenario.key(metadata=_logging_to(wire.url, **MESSAGE_LOGGING[message_logging]))
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(scenario, key, upstream.model)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
billed: Final = _billed_once(langfuse, SPEND_LANDING_SECONDS)
assert ("pong" in json.dumps(billed["attributes"])) is (message_logging == "logged"), billed["attributes"]
def _generated_key(scenario: Scenario, **fields: JsonValue) -> str:
key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"])
scenario.cleanups.callback(delete_key_if_present, scenario.gateway, key)
return key
def _delete_team_if_present(gateway: Gateway, team_id: str) -> None:
if gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)):
gateway.post("/team/delete", {"team_ids": [team_id]})
Retirement = Callable[[Scenario, str, str], tuple[str, Callable[[], object]]]
def _rotated_with_grace_then_relogged(
scenario: Scenario, billed_url: str, stale_url: str
) -> tuple[str, Callable[[], object]]:
key: Final = _generated_key(scenario, metadata=_logging_to(stale_url))
def rotate() -> None:
rotated: Final = string_value(
scenario.gateway.post("/key/regenerate", {"key": key, "grace_period": "1h"})["key"]
)
scenario.cleanups.callback(delete_key_if_present, scenario.gateway, rotated)
scenario.gateway.post("/key/update", {"key": rotated, "metadata": _logging_to(billed_url)})
return key, rotate
def _deleted_key(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]:
key: Final = _generated_key(scenario, metadata=_logging_to(billed_url))
return key, lambda: scenario.gateway.post("/key/delete", {"keys": [key]})
def _deleted_team(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]:
team_id: Final = string_value(
scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"]
)
scenario.cleanups.callback(_delete_team_if_present, scenario.gateway, team_id)
added: Final = scenario.gateway.request("POST", f"/team/{team_id}/callback", _langfuse_callback(billed_url))
assert added.status_code == 200, added.text
key: Final = _generated_key(scenario, team_id=team_id)
return key, lambda: scenario.gateway.post("/team/delete", {"team_ids": [team_id]})
RETIREMENTS: Final[Mapping[str, Retirement]] = {
"key_rotated_with_grace": _rotated_with_grace_then_relogged,
"key_deleted": _deleted_key,
"team_deleted": _deleted_team,
}
@pytest.mark.timeout(240)
@pytest.mark.parametrize("retired", RETIREMENTS.keys())
def test_a_retired_creators_billed_cost_reaches_its_own_langfuse_project(gateway: Gateway, retired: str) -> None:
with (
wire_server(langfuse_sink) as billed_wire,
wire_server(langfuse_sink) as stale_wire,
gateway.scenario() as scenario,
):
billed_langfuse: Final = Recorder(billed_wire)
stale_langfuse: Final = Recorder(stale_wire)
key, retire = RETIREMENTS[retired](scenario, billed_wire.url, stale_wire.url)
upstream: Final = scripted_upstream(scenario, background_routes("queued"))
response_id: Final = create_background_response(scenario, key, upstream.model)
retire()
upstream.script(background_routes("completed"))
_billed_once(billed_langfuse, ONE_POLL_SECONDS + SPEND_LANDING_SECONDS)
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
assert _priced_spans(stale_langfuse) == ()
@pytest.mark.timeout(240)
def test_a_creator_whose_langfuse_keys_resolve_to_nothing_is_still_billed_once(gateway: Gateway) -> None:
unset: Final = f"os.environ/INTEGRATION_UNSET_LANGFUSE_{uuid.uuid4().hex.upper()}"
with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario:
key: Final = scenario.key(metadata=_logging_to(wire.url, langfuse_public_key=unset, langfuse_secret_key=unset))
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(scenario, key, upstream.model)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert math.isclose(number(one_charged_row(upstream.model_id)["spend"]), PRICE)
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]

View file

@ -243,3 +243,183 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo
assert object_value(object_value(models["gpt-5"])["metrics"])["spend"] == 10.0
finally:
_clean(day)
def _seed_guardrail_units(day: str, rows: Sequence[tuple[str, str, str, float | None, int, int]]) -> None:
for guardrail_id, api_key, usage_unit, cost, units, untracked_units in rows:
write_rows(
'INSERT INTO "LiteLLM_DailyGuardrailUsageUnits" (guardrail_id, date, team_id, api_key, usage_unit,'
" units, cost, untracked_units, updated_at) VALUES (%s, %s, '', %s, %s, %s, %s, %s, now())",
(
guardrail_id,
day,
api_key,
usage_unit,
str(units),
None if cost is None else str(cost),
str(untracked_units),
),
)
def _guardrail_spend(breakdown: dict[str, JsonValue]) -> dict[str, tuple[float, dict[str, float]]]:
return {
guardrail_id: (
float(object_value(object_value(entry)["metrics"])["spend"]),
{
api_key: float(object_value(object_value(key_entry)["metrics"])["spend"])
for api_key, key_entry in object_value(object_value(entry)["api_key_breakdown"]).items()
},
)
for guardrail_id, entry in object_value(breakdown["guardrails"]).items()
}
def _day_breakdown(gateway: Gateway, day: str, **params: str) -> dict[str, JsonValue]:
results: Final = _RESULTS.validate_python(_activity(gateway, day, **params)["results"])
assert len(results) == 1
return object_value(object_value(results[0])["breakdown"])
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_breaks_down_guardrail_spend(gateway: Gateway) -> None:
day: Final = _unique_day()
_seed(
day,
[
(_row_id(), "user-a", day, "key-a", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 1.0, 1, 1),
(_row_id(), "user-b", day, "key-b", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 2.0, 1, 1),
],
)
_seed_guardrail_units(
day,
[
("pii-filter", "key-a", "contentPolicyUnits", 0.10, 5, 2),
("pii-filter", "key-a", "sensitiveInformationPolicyUnits", 0.05, 1, 0),
("pii-filter", "key-b", "contentPolicyUnits", 0.20, 1, 0),
("pii-filter", "key-b", "legacyUnits", None, 7, 0),
("topic-filter", "key-a", "topicPolicyUnits", 0.01, 1, 0),
("topic-filter", "key-b", "topicPolicyUnits", None, 1, 0),
("pii-filter", "orphan-key", "contentPolicyUnits", 9.0, 1, 0),
],
)
try:
everyone: Final = _day_breakdown(gateway, day)
assert _guardrail_spend(everyone) == {
"pii-filter": (pytest.approx(0.35), {"key-a": pytest.approx(0.15), "key-b": pytest.approx(0.20)}),
"topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01), "key-b": 0.0}),
}
api_keys: Final = object_value(everyone["api_keys"])
pii_filter: Final = object_value(object_value(everyone["guardrails"])["pii-filter"])
assert object_value(pii_filter["metrics"]) == {"spend": pytest.approx(0.35), "untracked_units": 9}
pii_api_keys: Final = object_value(pii_filter["api_key_breakdown"])
assert object_value(object_value(pii_api_keys["key-a"])["metrics"]) == {
"spend": pytest.approx(0.15),
"untracked_units": 2,
}
assert object_value(object_value(pii_api_keys["key-b"])["metrics"]) == {
"spend": pytest.approx(0.20),
"untracked_units": 7,
}
assert object_value(everyone["guardrails"])["topic-filter"] == {
"metrics": {"spend": pytest.approx(0.01), "untracked_units": 1},
"metadata": {"guardrail_name": "topic-filter"},
"api_key_breakdown": {
"key-a": {
"metrics": {"spend": pytest.approx(0.01), "untracked_units": 0},
"metadata": object_value(api_keys["key-a"])["metadata"],
},
"key-b": {
"metrics": {"spend": 0.0, "untracked_units": 1},
"metadata": object_value(api_keys["key-b"])["metadata"],
},
},
}
assert object_value(object_value(object_value(everyone["models"])["gpt-5"])["metrics"])["spend"] == 3.0
assert _guardrail_spend(_day_breakdown(gateway, day, user_id="user-a")) == {
"pii-filter": (pytest.approx(0.15), {"key-a": pytest.approx(0.15)}),
"topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01)}),
}
assert _guardrail_spend(_day_breakdown(gateway, day, api_key_limit="1")) == {
"pii-filter": (pytest.approx(0.35), {"key-b": pytest.approx(0.20)}),
"topic-filter": (pytest.approx(0.01), {"key-b": 0.0}),
}
limited_pii_filter: Final = object_value(
object_value(_day_breakdown(gateway, day, api_key_limit="1")["guardrails"])["pii-filter"]
)
assert object_value(limited_pii_filter["metrics"])["untracked_units"] == 9
assert set(object_value(limited_pii_filter["api_key_breakdown"])) == {"key-b"}
assert _guardrail_spend(_day_breakdown(gateway, day, model="gpt-5")) == {}
assert _guardrail_spend(_day_breakdown(gateway, day, model="")) == _guardrail_spend(everyone)
finally:
_clean(day)
write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,))
def _entity_breakdown(gateway: Gateway, url: str, day: str, **params: str) -> dict[str, JsonValue]:
response: Final = gateway.request("GET", url, params={"start_date": day, "end_date": day, **params})
assert response.status_code == 200, response.text
results: Final = _RESULTS.validate_python(object_value(response.json())["results"])
assert len(results) == 1
return object_value(object_value(results[0])["breakdown"])
@pytest.mark.asyncio
@pytest.mark.parametrize(
("url", "table", "entity_column", "filter_param"),
[
("/team/daily/activity/aggregated", "LiteLLM_DailyTeamSpend", "team_id", "team_ids"),
(
"/organization/daily/activity/aggregated",
"LiteLLM_DailyOrganizationSpend",
"organization_id",
"organization_ids",
),
],
)
async def test_team_and_organization_activity_break_down_guardrail_spend_by_name(
gateway: Gateway, url: str, table: str, entity_column: str, filter_param: str
) -> None:
day: Final = _unique_day()
db_guardrail_id: Final = str(uuid.uuid4())
db_guardrail_name: Final = f"pii-filter-{db_guardrail_id[:8]}"
for entity_id, api_key, spend in (("entity-a", "key-a", 1.0), ("entity-b", "key-b", 2.0)):
write_rows(
f'INSERT INTO "{table}" (id, {entity_column}, date, api_key, model, spend, api_requests,'
" successful_requests, updated_at) VALUES (%s, %s, %s, %s, 'gpt-5', %s, 1, 1, now())",
(_row_id(), entity_id, day, api_key, str(spend)),
)
write_rows(
'INSERT INTO "LiteLLM_GuardrailsTable" (guardrail_id, guardrail_name, litellm_params, updated_at)'
" VALUES (%s, %s, '{}', now())",
(db_guardrail_id, db_guardrail_name),
)
_seed_guardrail_units(
day,
[
(db_guardrail_id, "key-a", "contentPolicyUnits", 0.10, 1, 0),
(db_guardrail_id, "key-b", "contentPolicyUnits", 0.20, 1, 0),
("config-topic-filter", "key-a", "topicPolicyUnits", 0.05, 1, 0),
(db_guardrail_id, "orphan-key", "contentPolicyUnits", 9.0, 1, 0),
],
)
try:
everyone: Final = _entity_breakdown(gateway, url, day)
assert _guardrail_spend(everyone) == {
db_guardrail_id: (pytest.approx(0.30), {"key-a": pytest.approx(0.10), "key-b": pytest.approx(0.20)}),
"config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}),
}
guardrails: Final = object_value(everyone["guardrails"])
assert object_value(guardrails[db_guardrail_id])["metadata"] == {"guardrail_name": db_guardrail_name}
assert object_value(guardrails["config-topic-filter"])["metadata"] == {"guardrail_name": "config-topic-filter"}
assert _guardrail_spend(_entity_breakdown(gateway, url, day, **{filter_param: "entity-a"})) == {
db_guardrail_id: (pytest.approx(0.10), {"key-a": pytest.approx(0.10)}),
"config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}),
}
finally:
write_rows(f'DELETE FROM "{table}" WHERE date = %s', (day,))
write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,))
write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id = %s', (db_guardrail_id,))

View file

View file

@ -0,0 +1,235 @@
from datetime import datetime, timezone
from decimal import Decimal
from typing import Final
import pytest
from litellm.accounting import (
CostEvent,
Dimension,
DimensionName,
InboundUsageEvent,
Measurement,
Metric,
MetricName,
Outcome,
OutboundUsageEvent,
PricingSnapshot,
measurements_for,
)
from litellm.types.utils import CallTypes
_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc)
_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc)
_DIMENSIONS: Final = (
Dimension(name=DimensionName.TEAM, value="team"),
Dimension(name=DimensionName.TAG, value="alpha"),
Dimension(name=DimensionName.TAG, value="beta"),
)
def _snapshot(amount: Decimal) -> PricingSnapshot:
return PricingSnapshot(
calculator_version="test",
model="test-model",
provider=None,
pricing_source="upstream_header",
priced_model_key="test-model",
upstream_reported_cost=amount,
total_cost=amount,
)
@pytest.mark.parametrize(
("outcome", "outcome_metric"),
[
(Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS),
(Outcome.FAILED, MetricName.FAILED_REQUESTS),
(Outcome.CANCELLED, MetricName.FAILED_REQUESTS),
],
)
def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric):
event = InboundUsageEvent(
event_id="caller-event",
request_id="caller",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=outcome,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))),
)
assert measurements_for(event) == (
Measurement(
source_event_id="caller-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)),
),
Measurement(
source_event_id="caller-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=outcome_metric, value=Decimal(1)),
),
)
def test_usage_measurements_keep_every_metric_in_the_start_time_bucket():
event = OutboundUsageEvent(
event_id="usage-event",
attempt_id="attempt",
request_id="caller",
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(
Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)),
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)),
Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),
Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)),
),
)
assert measurements_for(event) == (
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)),
),
)
@pytest.mark.parametrize("metrics", [None, ()])
def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics):
event = OutboundUsageEvent(
event_id="usage-event",
attempt_id="attempt",
request_id=None,
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=(),
metrics=metrics,
)
assert measurements_for(event) == ()
@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")])
def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount):
event = CostEvent(
event_id="cost-event",
usage_event_id="usage-event",
occurred_at=_END,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),),
snapshot=_snapshot(amount),
)
assert measurements_for(event) == (
Measurement(
source_event_id="cost-event",
occurred_at=_END,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"),
),
)
def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges():
caller = InboundUsageEvent(
event_id="caller-event",
request_id="caller",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
)
retry = OutboundUsageEvent(
event_id="retry-usage",
attempt_id="retry",
request_id="caller",
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.FAILED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),),
)
success = OutboundUsageEvent(
event_id="success-usage",
attempt_id="success",
request_id="caller",
retry_of_attempt_id="retry",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),),
)
retry_cost = CostEvent(
event_id="retry-cost",
usage_event_id="retry-usage",
occurred_at=_START,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),),
snapshot=_snapshot(Decimal("0.025")),
)
success_cost = CostEvent(
event_id="success-cost",
usage_event_id="success-usage",
occurred_at=_START,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),),
snapshot=_snapshot(Decimal("0.125")),
)
measurements = (
*measurements_for(caller),
*measurements_for(retry),
*measurements_for(success),
*measurements_for(retry_cost),
*measurements_for(success_cost),
)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150")
assert {m.source_event_id for m in measurements} == {
"caller-event",
"retry-usage",
"success-usage",
"retry-cost",
"success-cost",
}

View file

@ -0,0 +1,484 @@
import re
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from enum import Enum
from pathlib import Path
from typing import Final
import pytest
from pydantic import ValidationError
from litellm.accounting import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
DeliveryState,
DeliveryTarget,
Dimension,
DimensionFilter,
DimensionName,
EventKind,
Granularity,
InboundUsageEvent,
Metric,
MetricName,
Multiplier,
Outcome,
OutboundUsageEvent,
PricedQuantity,
PricingSnapshot,
Rate,
validate_aggregate_rows,
)
from litellm.accounting.types import PricingUnit
from litellm.types.utils import CallTypes
_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc)
_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1))
_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10))
_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD")
_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team")
def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED):
return InboundUsageEvent(
event_id="inbound:request",
request_id="request",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_START,
outcome=outcome,
dimensions=(_TEAM,),
metrics=metrics,
)
def _outbound(metrics=(_TOKEN,), **changes):
return OutboundUsageEvent.model_validate(
{
"event_id": "outbound:attempt",
"attempt_id": "attempt",
"request_id": "request",
"retry_of_attempt_id": None,
"call_type": CallTypes.acompletion,
"started_at": _START,
"ended_at": _END,
"outcome": Outcome.SUCCEEDED,
"dimensions": (_TEAM,),
"metrics": metrics,
**changes,
}
)
def _query(**changes):
return AggregateQuery.model_validate(
{
"metric": MetricName.SPEND,
"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),),
"group_by": (DimensionName.MODEL,),
"start": _START,
"end": _END,
"granularity": Granularity.DAY,
**changes,
}
)
def test_wire_values_preserve_decimal_precision_and_normalize_aware_times():
offset = timezone(timedelta(hours=-8))
event = CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset),
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(),
)
wire = event.model_dump(mode="json")
assert wire["occurred_at"] == "2026-01-01T00:00:00Z"
assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}]
assert event.occurred_at == _START
assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value
@pytest.mark.parametrize(
("name", "value", "unit", "message"),
[
(MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"),
(MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"),
(MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"),
(MetricName.SPEND, Decimal("NaN"), "USD", "finite"),
(MetricName.SPEND, Decimal("Infinity"), "USD", "finite"),
(MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"),
(MetricName.SPEND, 0.1, "USD", "instance of Decimal"),
],
)
def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message):
with pytest.raises(ValidationError, match=message):
Metric(name=name, value=value, unit=unit)
@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()])
def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics):
with pytest.raises(ValidationError, match="Inbound events"):
_inbound(metrics)
@pytest.mark.parametrize(
("outcome", "metric"),
[
(Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))),
(Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))),
],
)
def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric):
with pytest.raises(ValidationError, match="count must match outcome"):
_inbound((_REQUEST, metric), outcome)
@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)])
def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics):
with pytest.raises(ValidationError, match="Outbound events own only"):
_outbound(metrics)
@pytest.mark.parametrize(
("changes", "message"),
[
({"started_at": datetime(2026, 1, 1)}, "timezone"),
({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"),
({"retry_of_attempt_id": "attempt"}, "cannot retry itself"),
({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"),
({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"),
(
{
"metrics": (
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)),
Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)),
)
},
"reasoning_tokens cannot exceed completion_tokens",
),
(
{
"metrics": (
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)),
)
},
"cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens",
),
({"event_id": " "}, "pattern"),
({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"),
],
)
def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message):
with pytest.raises(ValidationError, match=message):
_outbound(**changes)
def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable():
unknown = _outbound(None)
zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),))
assert unknown.metrics is None
assert zero.metrics is not None and zero.metrics[0].value == 0
assert _inbound().started_at == _inbound().ended_at
with pytest.raises(ValidationError, match="frozen"):
zero.attempt_id = "different"
with pytest.raises(ValidationError, match="frozen"):
zero.metrics[0].value = Decimal(10)
@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)])
def test_cost_events_require_one_authoritative_spend(charges):
with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"):
CostEvent(
event_id="cost",
usage_event_id="usage",
occurred_at=_START,
dimensions=(),
charges=charges,
snapshot=_snapshot(),
)
def test_cost_event_cannot_reuse_its_usage_identity():
with pytest.raises(ValidationError, match="distinct identities"):
CostEvent(
event_id="usage",
usage_event_id="usage",
occurred_at=_START,
dimensions=(),
charges=(_SPEND,),
snapshot=_snapshot(),
)
@pytest.mark.parametrize("granularity", [Granularity.DAY, None])
def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity):
query = _query(granularity=granularity)
assert query.granularity is granularity
assert query.filters[0].values == ("one", "two")
assert query.group_by == (DimensionName.MODEL,)
assert query.start == _START and query.end == _END
@pytest.mark.parametrize(
("changes", "message"),
[
({"end": _START}, "End must follow start"),
({"start": _START + timedelta(hours=1)}, "UTC day boundaries"),
({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"),
({"granularity": "hour"}, "day"),
({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"),
({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"),
({"filters": ({"name": "team", "values": ()},)}, "at least 1"),
],
)
def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message):
with pytest.raises(ValidationError, match=message):
_query(**changes)
def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups():
first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1))
second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2))
other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3))
rows = (second, other, first)
assert validate_aggregate_rows(rows) is rows
assert validate_aggregate_rows(()) == ()
def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order():
model = Dimension(name=DimensionName.MODEL, value="model")
first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1))
overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2))
with pytest.raises(ValueError, match="overlap within a dimension group"):
validate_aggregate_rows((first, overlapping))
def test_rows_reject_empty_intervals():
with pytest.raises(ValidationError, match="End must follow start"):
AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0))
def test_admission_decision_preserves_a_denial_reason():
decision = AdmissionDecision(allowed=False, reason="Team budget exhausted")
assert not decision.allowed
assert decision.reason == "Team budget exhausted"
def _snapshot(**changes: object) -> PricingSnapshot:
return PricingSnapshot.model_validate(
{
"calculator_version": "1.80.0",
"model": "gpt-4o",
"provider": "openai",
"pricing_source": "cost_map",
"priced_model_key": "openai/gpt-4o",
"multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),),
"rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),),
"usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),),
"input_cost": Decimal("0.000025"),
"total_cost": _SPEND.value,
**changes,
}
)
def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend():
event = CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=_START,
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(),
)
assert event.snapshot is not None
assert event.snapshot.total_cost == _SPEND.value
assert event.snapshot.multipliers[0].factor == Decimal("1.1")
with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"):
CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=_START,
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(total_cost=Decimal("0.5")),
)
@pytest.mark.parametrize(
("changes", "message"),
(
(
{
"rates": (
Rate(measure="input", unit="token", usd_per_unit=Decimal(1)),
Rate(measure="input", unit="token", usd_per_unit=Decimal(2)),
)
},
"Duplicate rate measures and units",
),
({"usage": (_SPEND,)}, "PricedQuantity"),
({"usage": (_REQUEST,)}, "PricedQuantity"),
({"pricing_source": "guess"}, "pricing_source"),
({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"),
({"schema_version": 2}, "schema_version"),
(
{"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})},
"Duplicate additional cost names",
),
),
)
def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str):
with pytest.raises(ValidationError, match=message):
_snapshot(**changes)
@pytest.mark.parametrize(
("unit", "value"),
[("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")],
)
def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None:
quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value))
rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003"))
amount: Final = quantity.value * rate.usd_per_unit
snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount)
restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json())
assert restored == snapshot
assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount
@pytest.mark.parametrize(
("usage", "message"),
[
((), "Every rate requires"),
(({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"),
(({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"),
(
({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2,
"Duplicate priced quantity",
),
],
)
def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities(
usage: tuple[Mapping[str, object], ...], message: str
) -> None:
with pytest.raises(ValidationError, match=message):
_snapshot(usage=usage)
@pytest.mark.parametrize("unit", ["token", "character", "image", "request"])
def test_priced_counts_remain_integral(unit: PricingUnit) -> None:
with pytest.raises(ValidationError, match="integral"):
PricedQuantity(measure="input", unit=unit, value=Decimal("1.25"))
@pytest.mark.parametrize(
("value", "message"),
[
(Decimal(-1), "greater than or equal"),
(Decimal("NaN"), "finite"),
(Decimal("Infinity"), "finite"),
(1.25, "instance of Decimal"),
],
)
def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None:
with pytest.raises(ValidationError, match=message):
PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value})
with pytest.raises(ValidationError, match=message):
Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value})
@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}])
def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None:
with pytest.raises(ValidationError, match="snapshot"):
CostEvent.model_validate(
{
"event_id": "cost",
"usage_event_id": "usage",
"occurred_at": _START,
"dimensions": (),
"charges": (_SPEND,),
**snapshot_fields,
}
)
def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None:
snapshot: Final = _snapshot(
pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value
)
restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json())
assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value
assert restored.rates == restored.usage == ()
def test_subset_token_counts_round_trip_without_changing_totals() -> None:
metrics: Final = (
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)),
Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)),
)
event: Final = _outbound(metrics=metrics)
assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics
@pytest.mark.parametrize(
"schema_path",
[
"schema.prisma",
"litellm/proxy/schema.prisma",
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
],
)
@pytest.mark.parametrize(
("enum_name", "python_enum"),
[
("AccountingEventKind", EventKind),
("AccountingDeliveryState", DeliveryState),
("AccountingDeliveryTarget", DeliveryTarget),
],
)
def test_prisma_accepts_exactly_the_public_accounting_enum_values(
schema_path: str, enum_name: str, python_enum: type[Enum]
) -> None:
schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text()
declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema)
assert declaration is not None, f"Missing persisted enum {enum_name}"
assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum)
def test_one_event_carries_every_matched_model_access_group():
groups = (
Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"),
Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"),
)
assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups)
with pytest.raises(ValidationError, match="Duplicate dimension values"):
_outbound(dimensions=(groups[0], groups[0]))
def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts():
multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1"))
with pytest.raises(ValidationError, match="Duplicate multiplier names"):
_snapshot(multipliers=(multiplier, multiplier))
with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"):
_snapshot(
usage=(
PricedQuantity(measure="input", unit="token", value=_TOKEN.value),
PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)),
PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)),
)
)
def test_agent_live_spend_is_a_delivery_target():
assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable"

View file

@ -235,13 +235,9 @@ def test_in_memory_cache_eviction_order():
# Add items with different TTLs
now = time.time()
in_memory_cache.set_cache(
key="early_expire", value="value_1", ttl=100
) # expires in 100 seconds
in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100) # expires in 100 seconds
time.sleep(0.01)
in_memory_cache.set_cache(
key="late_expire", value="value_2", ttl=200
) # expires in 200 seconds
in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200) # expires in 200 seconds
# Verify TTL order
early_ttl = in_memory_cache.ttl_dict["early_expire"]
@ -379,3 +375,20 @@ def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None:
assert cache.get_cache("second") is None
assert cache.get_cache("third") == "replacement"
assert cache.get_cache("fourth") == "new"
@pytest.mark.asyncio
async def test_async_increment_pipeline_keeps_each_key_for_its_own_ttl() -> None:
now: Final = {"t": 1_000.0}
cache: Final = InMemoryCache(default_ttl=600, clock=lambda: now["t"])
await cache.async_increment_pipeline(
increment_list=[
{"key": "minute", "increment_value": 1, "ttl": 60},
{"key": "day", "increment_value": 1, "ttl": 86400},
]
)
now["t"] += 601
assert cache.get_cache("minute") is None
assert cache.get_cache("day") == 1

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