scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge branch 'main' into litellm_tts_true_streaming

Browse source
This commit is contained in:
Gabriel Creti 2026-10-29 17:54:07 +01:00 • committed by GitHub
commit 027145b1ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
108 changed files with 7626 additions and 1026 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
codecov.yaml
View file

@ -15,9 +15,10 @@ ignore:
flag_management:
default_rules:
carryforward: true
# Dead flags no CI job uploads anymore: their carried-forward sessions were
# measured against old revisions, and the stale line maps mark comment lines
# of since-edited files as missed, sinking patch coverage on unrelated PRs.
# Dead flags no CI job uploads anymore, and flags whose job is path-gated on
# PRs: their carried-forward sessions were measured against old revisions, and
# the stale line maps mark comment lines of since-edited files as missed,
# sinking patch coverage on unrelated PRs.
individual_flags:
- name: proxy-mgmt-behavior
carryforward: false
@ -99,6 +100,10 @@ flag_management:
carryforward: false
- name: proxy-db-endpoints-and-responses
carryforward: false
- name: roi-postgres
carryforward: false
- name: redis-compat
carryforward: false
component_management:
individual_components:

6
helm/litellm-helm/README.md
View file

@ -259,3 +259,9 @@ itself.
`lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
## Shutdown accounting
`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged
Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml`

252
litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py
View file

@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it.
"""
import os
import re
from collections.abc import Callable, Iterator
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
REFUSAL: Final = "Existing rows would make a pending schema change fail:"
Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
Kind = Literal["too_long"]
@dataclass(frozen=True, slots=True)
@ -54,9 +52,6 @@ class Check:
statement: str
source: str
index: str | None = None
referenced_table: str | None = None
referenced_columns: tuple[str, ...] = ()
expression: str | None = None
@dataclass(frozen=True, slots=True)
@ -100,33 +95,6 @@ class ColumnType:
ColumnTypeLookup = Callable[[str, str], ColumnType | None]
_IDENT: Final = r'"?(\w+)"?'
_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
_CREATE_INDEX: Final = re.compile(
rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
re.IGNORECASE | re.DOTALL,
)
_ALTER_TABLE: Final = re.compile(
rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
)
_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
_ADD_COLUMN: Final = re.compile(
rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
re.IGNORECASE | re.DOTALL,
)
_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
_ADD_FOREIGN_KEY: Final = re.compile(
rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
re.IGNORECASE | re.DOTALL,
)
_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
_COMMENT: Final = re.compile(r"--[^\n]*")
_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
def migration_data_check_timeout_seconds() -> float:
"""Seconds one count query may run before the preflight gives up on verifying it."""
@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float:
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
def statements(migration_sql: str) -> tuple[str, ...]:
"""The statements of one migration file, comments stripped, empty ones dropped."""
parts: Final = _COMMENT.sub("", migration_sql).split(";")
return tuple(part.strip() for part in parts if part.strip())
def _split_top_level(text: str) -> tuple[str, ...]:
"""Split on the commas outside parentheses."""
depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
cuts: Final = tuple(
position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
)
starts: Final = (0, *(cut + 1 for cut in cuts))
stops: Final = (*cuts, len(text))
pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
return tuple(piece for piece in pieces if piece)
def _plain_columns(column_list: str) -> tuple[str, ...] | None:
"""The column names of an index or constraint column list, None when one entry is an
expression the derivation cannot turn into a count."""
matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
if any(match is None for match in matches):
return None
return tuple(match.group(1) for match in matches if match is not None)
def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
unique, name, table, method, column_list = match.groups()
columns: Final = _plain_columns(column_list)
if columns is None:
return
label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
if method is None or method.lower() == "btree":
yield Check("too_long", table, columns, label, source, index=name)
if unique:
yield Check("duplicates", table, columns, label, source, index=name)
def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
return
if (add_column := _ADD_COLUMN.match(action)) is not None:
column, body = add_column.groups()
if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
yield Check(
"any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
)
return
if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
own: Final = _plain_columns(columns)
theirs: Final = _plain_columns(referenced_columns)
if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
return
yield Check(
"orphans",
table,
own,
_constraint_label(table, constraint, "FOREIGN KEY"),
source,
referenced_table=referenced_table,
referenced_columns=theirs,
)
return
if (unique := _ADD_UNIQUE.match(action)) is not None:
constraint, columns = unique.groups()
unique_columns: Final = _plain_columns(columns)
if unique_columns is not None:
yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
return
if (check := _ADD_CHECK.match(action)) is not None:
constraint, expression, not_valid = check.groups()
if not_valid is None:
yield Check(
"check_violations",
table,
(),
_constraint_label(table, constraint, "CHECK"),
source,
expression=expression,
)
def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
return f'ALTER TABLE "{table}" ADD {named}{kind}'
def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
"""The count queries one SQL statement needs before it can be applied, empty for a
statement that cannot fail on existing rows or that the derivation cannot read."""
if (index := _CREATE_INDEX.match(statement_text)) is not None:
return tuple(_index_checks(index, statement_text, source))
if (alter := _ALTER_TABLE.match(statement_text)) is not None:
table, actions = alter.groups()
return tuple(
check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
) # comprehension-ok: flattening the per-action checks of one statement
return ()
def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
return tuple(
check
for statement_text in statements(migration_sql)
for check in checks_for_statement(statement_text, f"migration {name}")
) # comprehension-ok: flattening the per-statement checks of one migration
def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
"""The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
@ -424,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...],
def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
from psycopg import sql
table: Final = sql.Identifier(schema, check.table)
columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
match check.kind:
case "too_long":
return _entry_size_query(check, schema, types, limit_bytes)
case "duplicates":
not_null: Final = sql.SQL(" AND ").join(
sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
).format(t=table, nn=not_null, cols=columns)
case "nulls":
return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
t=table, col=sql.Identifier(check.columns[0])
)
case "any_row":
return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
case "orphans":
pairs: Final = sql.SQL(" AND ").join(
sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
)
present: Final = sql.SQL(" AND ").join(
sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
case "check_violations":
return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
t=table,
expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file
)
return _entry_size_query(check, schema, types, limit_bytes)
def _applies(
@ -474,15 +296,7 @@ def _applies(
return False
if check.index is not None and _index_is_valid(connection, schema, check.index):
return False
if check.kind == "any_row":
return types is None
if types is None:
return False
if check.kind == "too_long":
return any(column.can_overflow_btree for column in types)
if check.kind == "orphans":
return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
return True
return types is not None and any(column.can_overflow_btree for column in types)
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
@ -515,21 +329,7 @@ def _run_check(
rows: Final = _int(values[0], 0) if values else 0
if rows == 0:
return None
match check.kind:
case "too_long":
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
case "duplicates":
return BlockedChange(check, rows, "duplicated value groups")
case "nulls":
return BlockedChange(check, rows, "NULL")
case "any_row":
return BlockedChange(
check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
)
case "orphans":
return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
case "check_violations":
return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
def find_blockers(
@ -571,19 +371,7 @@ def find_blockers(
def _remedy(check: Check) -> str:
joined: Final = ", ".join(check.columns)
match check.kind:
case "too_long":
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
case "duplicates":
return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
case "nulls":
return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
case "any_row":
return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
case "orphans":
return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
case "check_violations":
return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
def _describe(finding: Finding) -> str:
@ -607,30 +395,4 @@ def _describe(finding: Finding) -> str:
def report(findings: tuple[Finding, ...]) -> str:
"""The operator-facing refusal: one block per finding, then what to do next."""
lines: Final = "\n".join(_describe(finding) for finding in findings)
return (
f"{REFUSAL}\n{lines}\n"
"No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
)
def pending_change_checks(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
migrations_dir: Path | None,
indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
) -> tuple[Check, ...]:
"""Every check the pending migrations and the request-log indexes need. `migrations_dir`
is None when the schema is pushed rather than migrated, so only the indexes are checked."""
migration_checks: Final = (
tuple(
check
for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
for check in checks_for_migration(name, migration_sql)
) # comprehension-ok: flattening the per-migration checks
if migrations_dir is not None
else ()
)
index_checks: Final = tuple(
check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
) # comprehension-ok: flattening the per-index checks
return migration_checks + index_checks
return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up."

184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py Normal file
View file

@ -0,0 +1,184 @@
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each
inside its own transaction that is always rolled back, so Postgres itself says whether this
database's rows and schema take them. Each migration holds the locks the real upgrade will,
and releases them before the next one runs."""
from __future__ import annotations
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from pathlib import Path
from typing import TYPE_CHECKING, Final
from litellm_proxy_extras.migration_data_checks import pending_migrations
if TYPE_CHECKING:
import psycopg
_CONCURRENTLY: Final = re.compile(
r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE
)
UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"})
Migration = tuple[str, str]
@dataclass(frozen=True, slots=True)
class Applied:
name: str
seconds: float
replayed: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class Rejected:
name: str
error: str
seconds: float
@dataclass(frozen=True, slots=True)
class Baseline:
"""A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying
history, so the dry run asks the same question the boot does. `refusal` is None when the boot
will baseline, otherwise its reason (a custom migration folder, or the schema diff)."""
refusal: str | None
@property
def matches(self) -> bool:
return self.refusal is None
@dataclass(frozen=True, slots=True)
class DryRun:
applied: tuple[Applied, ...]
rejected: Rejected | None
baseline: Baseline | None = None
@property
def ok(self) -> bool:
return self.rejected is None and (self.baseline is None or self.baseline.matches)
def transactional(migration_sql: str) -> str:
"""The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones."""
return _CONCURRENTLY.sub(r"\1", migration_sql)
def _attempt(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
chain: tuple[Migration, ...],
clock: Callable[[], float],
) -> tuple[float, psycopg.Error | None]:
import psycopg
from psycopg import sql
started: Final = clock()
try:
with connection.transaction(force_rollback=True):
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
for _, migration_sql in chain:
connection.execute(transactional(migration_sql))
except psycopg.Error as exc:
return clock() - started, exc
return clock() - started, None
def _run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
pending: tuple[Migration, ...],
position: int,
done: DryRun,
clock: Callable[[], float],
) -> DryRun:
"""One migration in its own rolled-back transaction. A migration that fails because something
an earlier pending migration creates does not exist yet is rerun with those replayed in front."""
if done.rejected is not None:
return done
name: Final = pending[position][0]
seconds, error = _attempt(connection, schema, (pending[position],), clock)
if error is None:
return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline)
if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0:
return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline)
earlier: Final = pending[:position]
chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock)
if chained_error is None:
replayed: Final = tuple(earlier_name for earlier_name, _ in earlier)
return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline)
return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline)
def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
from psycopg import sql
row: Final = connection.execute(
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),)
).fetchone()
return row is not None and row[0] is not None
def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
row: Final = connection.execute(
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1",
(schema,),
).fetchone()
return row is not None
def dry_run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
migrations_dir: Path,
baseline_check: Callable[[], str | None] | None = None,
clock: Callable[[], float] = time.monotonic,
) -> DryRun:
"""Apply every pending migration in order, each in its own transaction that is rolled back
before the next one starts. The first migration Postgres rejects ends the run with its error.
`baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for
the replay on a database that has tables but no migration history, which the boot baselines."""
if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
return DryRun((), None, Baseline(baseline_check()))
pending: Final = pending_migrations(connection, schema, migrations_dir)
return reduce(
lambda done, position: _run(connection, schema, pending, position, done, clock),
range(len(pending)),
DryRun((), None),
)
def _timing(applied: Applied) -> str:
replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else ""
return f"\n {applied.name}: {applied.seconds:.1f} s{replay}"
def report(result: DryRun) -> str:
"""What the operator reads: the per-migration time the real upgrade will hold its locks,
or the migration Postgres rejected with Postgres's own words. Nothing was applied either way."""
if result.baseline is not None:
if result.baseline.matches:
return (
"No migration history and the schema matches this build, "
"the proxy baselines it at boot, nothing to dry-run."
)
return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}"
timings: Final = "".join(_timing(applied) for applied in result.applied)
if result.rejected is None:
count: Final = len(result.applied)
if count == 0:
return "No pending migrations, this database is on this version's schema."
return (
f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next "
f"so its locks are released, nothing changed:{timings}"
)
return (
f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, "
f"nothing was applied.\n{result.rejected.error}"
+ (f"\nApplied before it in the dry run:{timings}" if timings else "")
)

21
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py
View file

@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
)
def has_packaged_history(migrations_dir: Path) -> bool:
"""Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
the only history the boot baselines an existing database onto."""
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
return (
bool(migrations)
and migrations == migration_files(packaged_dir)
and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
)
def baseline_current_schema(
coordinator: MigrationCoordinator,
schema: str,
@ -95,14 +107,9 @@ def baseline_current_schema(
) -> None:
from psycopg import sql
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
if (
not migrations
or migrations != migration_files(packaged_dir)
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
):
if not has_packaged_history(migrations_dir):
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
migrations: Final = migration_files(migrations_dir)
coordinator.acquire_prisma_lock()
existing: Final = coordinator.connection.execute(

107
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -11,18 +11,13 @@ from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from typing import TYPE_CHECKING, Final, Protocol
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_data_checks import (
Finding,
find_blockers,
migration_data_check_timeout_seconds,
pending_change_checks,
report,
)
from litellm_proxy_extras.migration_dry_run import dry_run
from litellm_proxy_extras.migration_dry_run import report as dry_run_report
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import (
apply_replica_identity_full,
)
from litellm_proxy_extras.request_log_indexes import (
REQUEST_LOG_INDEXES,
ensure_request_log_indexes,
filter_request_log_index_diff,
)
@ -52,7 +46,7 @@ class LensCheckConnect(Protocol):
) -> "psycopg.Connection[tuple[object, ...]]": ...
def str_to_bool(value: Optional[str]) -> bool:
def str_to_bool(value: str | None) -> bool:
if value is None:
return False
return value.lower() in ("true", "1", "t", "y", "yes")
@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str:
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
def _without_spend_logs_pk_clauses(statement: str) -> str | None:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int:
return max(_migration_timestamp(n) for n in names)
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]:
if isinstance(command, tuple):
return tuple(_redact_credentials(str(argument)) for argument in command)
if isinstance(command, list):
@ -404,7 +398,7 @@ class ProxyExtrasDBManager:
return tuple(row) if row is not None else ()
@staticmethod
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
def _failed_migration_logs(migration_name: str, started_at: str) -> str | None:
row: Final = ProxyExtrasDBManager._read_migration_ledger(
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
"AND finished_at IS NULL AND rolled_back_at IS NULL",
@ -721,10 +715,10 @@ class ProxyExtrasDBManager:
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
def _prisma_schema_param(url: str) -> str | None:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
from urllib.parse import parse_qsl, urlparse
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
if not migrated:
return False
@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager:
return True
@staticmethod
def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
"""What in the existing rows would make a pending migration or request-log index
fail, read with count queries only: no write, no table lock. Empty when nothing
blocks, when no database is configured, or when it cannot be reached, so the
migration step that follows reports the connection problem the way it does today."""
def _baseline_refusal(prisma_dir: Path) -> str | None:
"""Why `baseline_current_schema` would refuse this database at boot, None when it would baseline:
the migrations folder must be the packaged one, then `prisma migrate diff` must find the live
database and this build's schema.prisma identical."""
from litellm_proxy_extras.migration_lock import migration_environment
from litellm_proxy_extras.migration_recovery import has_packaged_history
if not has_packaged_history(prisma_dir):
return (
f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the "
"migrations this build ships, and the proxy only baselines onto the packaged history."
)
try:
prisma_toolchain.run_prisma(
(
_get_prisma_command(),
"migrate",
"diff",
"--from-schema-datasource",
str(prisma_dir / "schema.prisma"),
"--to-schema-datamodel",
str(prisma_dir / "schema.prisma"),
"--exit-code",
),
timeout=prisma_toolchain.prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except subprocess.CalledProcessError as exc:
diff: Final = str(exc.stdout or exc.stderr or exc).strip()
return f"the schema differs from this build. prisma migrate diff:\n{diff}"
except subprocess.TimeoutExpired as exc:
return f"prisma migrate diff did not finish within {exc.timeout} s"
return None
@staticmethod
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
"""`--check_migrations_against_existing_data`: the pending migrations applied for real on
DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports
what the upgrade would hit. (True, the per-migration timings) when they all apply,
(False, Postgres's own error) when one does not."""
import psycopg
if not use_migrate:
return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.")
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return ()
return (False, "DATABASE_URL is not set.")
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
os.environ.get("DIRECT_URL") or database_url
)
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir())
try:
with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
result: Final = dry_run(
connection,
schema,
prisma_dir / "migrations",
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
)
except psycopg.OperationalError as exc:
logger.warning(
"Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
)
return ()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
"""The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
return report(findings) if findings else None
@staticmethod
def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
"""Refuse to run migrations that the existing rows would make fail, before any is
applied. The RuntimeError carries the report and reaches the operator through the
proxy's "Database migration cannot proceed" exit."""
refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
if refusal is not None:
raise RuntimeError(refusal)
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")
return (result.ok, dry_run_report(result))
@staticmethod
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager:
time.sleep(random.randrange(5, 15))
finally:
os.chdir(original_dir)
pass
return False

11
litellm/exceptions.py
View file

@ -323,6 +323,17 @@ class UnsupportedEndpointError(BadRequestError):
pass
class EndpointNotSupportedByModelError(BadRequestError):
def __init__(self, message: str, model: str, llm_provider: str = "") -> None:
super().__init__( # pyright: ignore[reportUnknownMemberType] # inherited public exception constructor is untyped
message=message, model=model, llm_provider=llm_provider
)
self.message = message
self.type = "invalid_request_error"
self.param = "model"
self.code = "400"
class UnprocessableEntityError(openai.UnprocessableEntityError):
def __init__(
self,

12
litellm/images/main.py
View file

@ -68,6 +68,11 @@ from litellm.utils import (
)
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
_PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE: Final = (
LlmProviders.DASHSCOPE,
LlmProviders.QWENCLOUD,
LlmProviders.QWEN_AI_PLATFORM,
)
# Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
_ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None
@ -271,6 +276,7 @@ def image_generation(
"""
try:
args: Final = locals()
requested_api_base: Final = api_base
aimg_generation: Final = kwargs.get("aimg_generation", False)
litellm_call_id: Final = kwargs.get("litellm_call_id", None)
logger_fn: Final = kwargs.get("logger_fn", None)
@ -468,7 +474,11 @@ def image_generation(
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
# Resolve api_base from litellm.api_base if not explicitly provided
_api_base: Final = api_base or litellm.api_base
_api_base: Final = (
requested_api_base
if custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE
else api_base or litellm.api_base
)
litellm_params_dict["api_base"] = _api_base
return llm_http_handler.image_generation_handler(

125
litellm/litellm_core_utils/logging_worker.py
View file

@ -6,7 +6,7 @@ import atexit
import contextvars
import inspect
import logging
from collections.abc import Coroutine, Iterator
from collections.abc import Awaitable, Callable, Coroutine, Iterator
from typing import Final
from typing_extensions import TypedDict
@ -66,14 +66,19 @@ class LoggingWorker:
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
*,
retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self.timeout_summary_window = timeout_summary_window
self._queue: asyncio.Queue[LoggingTask] | None = None
self._worker_task: asyncio.Task | None = None
self._running_tasks: set[asyncio.Task] = set()
self._worker_task: asyncio.Task[None] | None = None
self._running_tasks: set[asyncio.Task[None]] = set()
self._overflow_tasks: set[asyncio.Task[None]] = set() # mutable-ok: done callbacks discard refs
self._pending_enqueues: dict[int, LoggingTask] = {} # mutable-ok: retries and shutdown share callback ownership
self._retry_sleep = retry_sleep
self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks
self._sem: asyncio.Semaphore | None = None
self._bound_loop: asyncio.AbstractEventLoop | None = None
@ -168,6 +173,7 @@ class LoggingWorker:
self._sem = None
self._worker_task = None
self._running_tasks.clear()
self._overflow_tasks.clear()
self._timeout_summary_task = None
self._timeout_burst_count = 0
self._queue = new_queue
@ -185,6 +191,10 @@ class LoggingWorker:
self._sem = asyncio.Semaphore(self.concurrency)
if self._worker_task is None or self._worker_task.done():
self._worker_task = asyncio.create_task(self._worker_loop())
if not self._overflow_tasks:
self._aggressive_clear_in_progress = False
for pending in tuple(self._pending_enqueues.values()):
self._schedule_delayed_enqueue_retry(pending)
async def _process_log_task(
self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
@ -278,12 +288,24 @@ class LoggingWorker:
# Capture the current context when enqueueing
task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
self._enqueue_task(task)
def _enqueue_task(self, task: LoggingTask) -> None:
if self._queue is None:
return
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
# Queue is full - handle it appropriately
verbose_logger.exception("LoggingWorker queue is full")
self._pending_enqueues[id(task)] = task
self._handle_queue_full(task)
else:
self._pending_enqueues.pop(id(task), None)
def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None:
task: Final = asyncio.create_task(coroutine)
self._overflow_tasks.add(task)
task.add_done_callback(self._overflow_tasks.discard)
def _should_start_aggressive_clear(self) -> bool:
"""
@ -326,7 +348,7 @@ class LoggingWorker:
if self._should_start_aggressive_clear():
self._mark_aggressive_clear_started()
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
asyncio.create_task(self._aggressively_clear_queue_async(task))
self._start_overflow_task(self._aggressively_clear_queue_async(task))
else:
# Cooldown active or clear in progress, schedule a delayed retry
self._schedule_delayed_enqueue_retry(task)
@ -363,7 +385,7 @@ class LoggingWorker:
delay: Final = self._calculate_retry_delay()
# Schedule the retry as a background task
asyncio.create_task(self._retry_enqueue_task(task, delay))
self._start_overflow_task(self._retry_enqueue_task(task, delay))
except RuntimeError:
# No event loop, drop the task as we can't schedule a retry
pass
@ -373,17 +395,8 @@ class LoggingWorker:
Retry enqueueing the task after delay, preserving original context.
This is called as a background task from _schedule_delayed_enqueue_retry.
"""
await asyncio.sleep(delay)
# Try to enqueue the task directly, preserving its original context
if self._queue is None:
return
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
# Still full - handle it appropriately (clear or retry again)
self._handle_queue_full(task)
await self._retry_sleep(delay)
self._enqueue_task(task)
def _extract_tasks_from_queue(self) -> list[LoggingTask]:
"""
@ -423,15 +436,10 @@ class LoggingWorker:
return
extracted_tasks: Final = self._extract_tasks_from_queue()
# Add new task to extracted tasks to process directly
if new_task is not None:
self._track_dequeued(new_task)
extracted_tasks.append(new_task)
# Process extracted tasks directly
self._enqueue_task(new_task)
if extracted_tasks:
await self._process_extracted_tasks(extracted_tasks)
await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks))
except Exception as e:
verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e)
finally:
@ -439,7 +447,7 @@ class LoggingWorker:
self._aggressive_clear_in_progress = False
async def _process_single_task(self, task: LoggingTask) -> None:
"""Process a single task and mark it done."""
"""Process a task extracted from the queue and mark it done."""
if self._queue is None:
return
@ -455,17 +463,6 @@ class LoggingWorker:
self._untrack_dequeued(task)
self._queue.task_done()
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
"""
Process tasks that were extracted from the queue to make room.
Processes them concurrently without semaphore limits for maximum speed.
"""
if not tasks or self._queue is None:
return
# Process all tasks concurrently for maximum speed
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
"""
Ensure the logging worker is initialized and enqueue the coroutine.
@ -473,32 +470,63 @@ class LoggingWorker:
self.start()
self.enqueue(async_coroutine)
async def stop(self) -> None:
@property
def pending_count(self) -> int:
return (
(self._queue.qsize() if self._queue is not None else 0)
+ len(self._dequeued_tasks)
+ len(self._pending_enqueues)
)
async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None:
"""Stop the logging worker and clean up resources."""
if discard_pending and self._queue is not None:
discarded: Final = self._drain_pending(self._queue)
for task in discarded:
task["coroutine"].close()
self._queue.task_done()
if discarded:
verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded))
if self._timeout_summary_task is not None:
self._timeout_summary_task.cancel()
self._timeout_summary_task = None
self._emit_timeout_summary()
if self._worker_task is None and not self._running_tasks:
# No worker launched and no in-flight tasks to drain.
return
if discard_pending:
rejected: Final = tuple(self._pending_enqueues.values())
self._pending_enqueues.clear()
for task in rejected:
task["coroutine"].close()
if rejected:
verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected))
tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
if self._worker_task:
# Include the main worker loop so it stops fetching work.
tasks_to_cancel.append(self._worker_task)
tasks_to_cancel: Final = (
*self._running_tasks,
*self._overflow_tasks,
*((self._worker_task,) if self._worker_task is not None else ()),
)
if not tasks_to_cancel:
return
for task in tasks_to_cancel:
# Propagate cancellation to every pending task.
task.cancel()
# Wait for cancellation to settle; ignore errors raised during shutdown.
await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout)
await asyncio.gather(*done, return_exceptions=True)
if pending:
verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending))
self._worker_task = None
if discard_pending and self._queue is not None:
for task in self._unstarted_dequeued_tasks():
task["coroutine"].close()
self._untrack_dequeued(task)
self._queue.task_done()
if self._worker_task is not None and self._worker_task.done():
self._worker_task = None
# Drop references to completed tasks so we can restart cleanly.
self._running_tasks.clear()
self._running_tasks.difference_update(done)
async def flush(self) -> None:
"""Flush the logging queue.
@ -516,6 +544,9 @@ class LoggingWorker:
return
self.start()
await self._queue.join()
while self._overflow_tasks:
await asyncio.wait(tuple(self._overflow_tasks))
await self._queue.join()
async def clear_queue(self):
"""

60
litellm/litellm_core_utils/served_output_texts.py
View file

@ -4,12 +4,13 @@ post-call guardrail rewrote rather than the provider response the proxy assemble
from __future__ import annotations
from collections.abc import Sequence
from itertools import chain
from typing import Final, Literal
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import ModelResponse, ModelResponseStream
from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse
SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts"
@ -44,6 +45,19 @@ class _ChatChoices(LiteLLMBaseModel):
choices: list[object]
class _TextCompletionChoice(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
index: int | None = None
text: str | None = None
class _TextCompletionChunk(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
choices: tuple[_TextCompletionChoice, ...]
def _as_json_object(response: object) -> dict[str, object] | None:
candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response
try:
@ -108,8 +122,13 @@ def served_output_texts(response: object) -> ServedTexts | None:
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
chat_chunks: Final = tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream))
if chunks and len(chat_chunks) == len(chunks):
return _indexed_texts_by_choice(tuple(chain.from_iterable(_chat_chunk_texts(chunk) for chunk in chat_chunks)))
if chunks and all(isinstance(chunk, TextCompletionResponse) for chunk in chunks):
return _indexed_texts_by_choice(
tuple(chain.from_iterable(_text_completion_chunk_texts(chunk) for chunk in chunks))
)
from litellm.llms.anthropic.messages.stream_assembly import (
assemble_anthropic_sse_stream,
is_anthropic_sse_stream,
@ -121,19 +140,26 @@ def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
return _chat_texts(assembled) if assembled is not None else None
def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None:
contents: Final = tuple(
content
for chunk in chunks
def _chat_chunk_texts(chunk: ModelResponseStream) -> tuple[tuple[int, str | None], ...]:
return tuple(
(choice.index or 0, content if isinstance(content := choice.delta.content, str) else None)
for choice in chunk.choices
if (choice.index or 0) == index and isinstance(content := choice.delta.content, str)
)
return "".join(contents) if contents else None
def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None:
choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0)
texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count))
def _text_completion_chunk_texts(chunk: object) -> tuple[tuple[int, str | None], ...]:
parsed: Final = _TextCompletionChunk.model_validate(_as_json_object(chunk))
return tuple((choice.index or 0, choice.text) for choice in parsed.choices)
def _joined_choice_text(indexed_texts: Sequence[tuple[int, str | None]], index: int) -> str | None:
texts: Final = tuple(text for text_index, text in indexed_texts if text_index == index and text is not None)
return "".join(texts) if texts else None
def _indexed_texts_by_choice(indexed_texts: Sequence[tuple[int, str | None]]) -> ServedTexts | None:
choice_count: Final = max((index + 1 for index, _ in indexed_texts), default=0)
texts: Final = tuple(_joined_choice_text(indexed_texts, index) for index in range(choice_count))
return texts if any(text is not None for text in texts) else None
@ -170,6 +196,8 @@ def _choice_with_text(choice: object, text: str | None) -> object:
if choice_obj is None or text is None:
return choice
message: Final = _as_json_object(choice_obj.get("message"))
if message is None or message.get("content") == text:
return choice
return {**choice_obj, "message": {**message, "content": text}}
if message is not None:
return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}}
if isinstance(choice_obj.get("text"), str):
return {**choice_obj, "text": text}
return choice

9
litellm/llms/dashscope/chat/transformation.py
View file

@ -10,6 +10,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
from ..common_utils import resolve_dashscope_openai_api_base
class DashScopeChatConfig(OpenAIGPTConfig):
@ -55,11 +56,9 @@ class DashScopeChatConfig(OpenAIGPTConfig):
def _get_openai_compatible_provider_info(
self, api_base: str | None, api_key: str | None
) -> tuple[str | None, str | None]:
api_base = (
api_base or get_secret_str("DASHSCOPE_API_BASE") or "https://dashscope.aliyuncs.com/compatible-mode/v1"
)
resolved_api_base: Final = resolve_dashscope_openai_api_base(api_base)
dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
return api_base, dynamic_api_key
return resolved_api_base, dynamic_api_key
def get_openai_compatible_provider_info(
self,
@ -69,7 +68,7 @@ class DashScopeChatConfig(OpenAIGPTConfig):
return self._get_openai_compatible_provider_info(api_base, api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1"
return resolve_dashscope_openai_api_base(api_base)
def get_complete_url(
self,

136
litellm/llms/dashscope/common_utils.py
View file

@ -2,6 +2,7 @@
Common utilities for the DashScope LLM provider.
"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from urllib.parse import urlparse
@ -19,10 +20,123 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1"
DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks"
DASHSCOPE_ANTHROPIC_PATH: Final = "/apps/anthropic"
DASHSCOPE_NATIVE_V1_PATH: Final = "/api/v1"
DASHSCOPE_COMPATIBLE_API_PATH: Final = "/compatible-api/v1"
DASHSCOPE_RERANK_PATH: Final = f"{DASHSCOPE_COMPATIBLE_API_PATH}/reranks"
DASHSCOPE_IMAGE_PATH: Final = "/api/v1/services/aigc/multimodal-generation/generation"
DASHSCOPE_DEFAULT_ORIGIN: Final = "https://dashscope.aliyuncs.com"
_STANDARD_DASHSCOPE_HOSTS: Final = frozenset(
{
"dashscope.aliyuncs.com",
"dashscope-intl.aliyuncs.com",
"dashscope-us.aliyuncs.com",
}
)
_WORKSPACE_HOST_SUFFIX: Final = ".maas.aliyuncs.com"
_EXCLUDED_WORKSPACE_PREFIXES: Final = ("trial.", "token-plan.")
_PROTOCOL_PATH_PREFIXES: Final = (
DASHSCOPE_CHAT_COMPATIBLE_PATH,
DASHSCOPE_ANTHROPIC_PATH,
DASHSCOPE_NATIVE_V1_PATH,
DASHSCOPE_COMPATIBLE_API_PATH,
)
def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
@dataclass(frozen=True, slots=True)
class DashScopeEndpoints:
origin: str
openai: str
anthropic: str
native: str
def _is_dashscope_host(host: str) -> bool:
if host in _STANDARD_DASHSCOPE_HOSTS:
return True
return host.endswith(_WORKSPACE_HOST_SUFFIX) and not host.startswith(_EXCLUDED_WORKSPACE_PREFIXES)
def _dashscope_origin(api_base: str | None) -> str | None:
if not api_base:
return None
parsed: Final = urlparse(api_base)
if parsed.scheme not in ("http", "https") or not parsed.netloc:
return None
if parsed.query or parsed.fragment:
return None
host: Final = (parsed.hostname or "").lower()
if not _is_dashscope_host(host):
return None
path: Final = parsed.path.rstrip("/")
if path and not any(path == prefix or path.startswith(f"{prefix}/") for prefix in _PROTOCOL_PATH_PREFIXES):
return None
return f"{parsed.scheme}://{parsed.netloc}"
def resolve_dashscope_endpoints(api_base: str | None) -> DashScopeEndpoints | None:
"""
Derive DashScope's protocol-specific bases from any standard DashScope base.
Accepts the domain root, an OpenAI-compatible base, an Anthropic Messages
base, or a native `/api/v1` base for `dashscope.aliyuncs.com`,
`dashscope-intl.aliyuncs.com`, `dashscope-us.aliyuncs.com`, and
`{workspace}.{region}.maas.aliyuncs.com`. Returns None for any other host
or path so callers keep the caller-supplied base untouched.
"""
origin: Final = _dashscope_origin(api_base)
if origin is None:
return None
return DashScopeEndpoints(
origin=origin,
openai=f"{origin}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
anthropic=f"{origin}{DASHSCOPE_ANTHROPIC_PATH}",
native=f"{origin}{DASHSCOPE_NATIVE_V1_PATH}",
)
def is_dashscope_native_endpoint_url(api_base: str) -> bool:
return urlparse(api_base).path.rstrip("/").startswith(f"{DASHSCOPE_NATIVE_V1_PATH}/")
def resolve_dashscope_openai_api_base(
api_base: str | None,
env_var: str = "DASHSCOPE_API_BASE",
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
) -> str:
candidate: Final = api_base or get_secret_str(env_var)
if candidate and is_dashscope_native_endpoint_url(candidate):
return candidate
endpoints: Final = resolve_dashscope_endpoints(candidate)
if endpoints is not None:
return endpoints.openai
return candidate or default
def resolve_dashscope_anthropic_api_base(
api_base: str | None,
env_var: str = "DASHSCOPE_API_BASE",
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_ANTHROPIC_PATH}",
) -> str:
candidate: Final = api_base or get_secret_str(env_var)
endpoints: Final = resolve_dashscope_endpoints(candidate)
if endpoints is not None:
return endpoints.anthropic
return candidate or default
def dashscope_base_selects_anthropic_messages(api_base: str | None) -> bool:
candidate: Final = api_base or get_secret_str("DASHSCOPE_API_BASE")
if not candidate:
return False
path: Final = urlparse(candidate).path.rstrip("/")
if path.endswith(DASHSCOPE_ANTHROPIC_PATH) or f"{DASHSCOPE_ANTHROPIC_PATH}/" in path:
return True
return not path and resolve_dashscope_endpoints(candidate) is not None
def _chat_shaped_aliyun_origin(api_base: str | None) -> str | None:
if api_base is None:
return None
parsed: Final = urlparse(api_base)
@ -30,14 +144,24 @@ def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com")
if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH:
return None
return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}"
return f"{parsed.scheme}://{parsed.netloc}"
def _is_compatible_api_base(api_base: str) -> bool:
path: Final = urlparse(api_base).path.rstrip("/")
return path == DASHSCOPE_COMPATIBLE_API_PATH or path.startswith(f"{DASHSCOPE_COMPATIBLE_API_PATH}/")
def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str:
remapped: Final = _rerank_base_for_chat_shaped_base(api_base)
if api_base is not None and remapped is None:
if api_base is None:
return get_secret_str(env_var) or default_rerank_base
if _is_compatible_api_base(api_base):
return api_base
return get_secret_str(env_var) or remapped or default_rerank_base
chat_origin: Final = _chat_shaped_aliyun_origin(api_base)
if chat_origin is not None:
return get_secret_str(env_var) or f"{chat_origin}{DASHSCOPE_RERANK_PATH}"
endpoints: Final = resolve_dashscope_endpoints(api_base)
return f"{endpoints.origin}{DASHSCOPE_RERANK_PATH}" if endpoints is not None else api_base
def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig":

4
litellm/llms/dashscope/embed/transformation.py
View file

@ -25,7 +25,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import DashScopeError
from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base
DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1"
@ -91,7 +91,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
return resolved_api_key
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
return resolve_dashscope_openai_api_base(api_base, default=DEFAULT_API_BASE)
def validate_environment(
self,

35
litellm/llms/dashscope/image_generation/transformation.py
View file

@ -39,6 +39,13 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import ImageObject, ImageResponse
from ..common_utils import (
DASHSCOPE_CHAT_COMPATIBLE_PATH,
DASHSCOPE_IMAGE_PATH,
is_dashscope_native_endpoint_url,
resolve_dashscope_endpoints,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
@ -53,8 +60,6 @@ _DECODED_JSON: Final = TypeAdapter(object)
DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
CHAT_COMPATIBLE_MODE_PATH: Final = "/compatible-mode/v1"
# Maps OpenAI size strings (WxH) to DashScope size strings (W*H)
OPENAI_TO_DASHSCOPE_SIZE: Final[dict] = {
"256x256": "256*256",
@ -102,8 +107,14 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
raise ValueError("DASHSCOPE_API_KEY is not set")
return resolved_api_key
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("DASHSCOPE_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return DEFAULT_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("DASHSCOPE_API_BASE")
def get_complete_url(
self,
@ -114,10 +125,18 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
litellm_params: dict,
stream: bool | None = None,
) -> str:
image_api_base: Final = (
api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None
)
return self._resolve_image_api_base(image_api_base)
image_env_api_base: Final = self._image_api_base_env()
configured_api_base: Final = api_base or self._configured_chat_api_base()
if not configured_api_base or configured_api_base.rstrip("/").endswith(DASHSCOPE_CHAT_COMPATIBLE_PATH):
return image_env_api_base or self._default_image_api_base()
if is_dashscope_native_endpoint_url(configured_api_base):
return configured_api_base
endpoints: Final = resolve_dashscope_endpoints(configured_api_base)
if endpoints is None:
return configured_api_base
if not api_base and image_env_api_base:
return image_env_api_base
return f"{endpoints.origin}{DASHSCOPE_IMAGE_PATH}"
def validate_environment(
self,

7
litellm/llms/dashscope/messages/__init__.py Normal file
View file

@ -0,0 +1,7 @@
"""
DashScope Anthropic-compatible Messages Module
"""
from .transformation import DashScopeAnthropicMessagesConfig
__all__ = ["DashScopeAnthropicMessagesConfig"]

96
litellm/llms/dashscope/messages/transformation.py Normal file
View file

@ -0,0 +1,96 @@
"""
DashScope Anthropic-compatible Messages transformation config.
"""
from typing import ClassVar, Final
import litellm
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.url_utils import ensure_api_route
from litellm.llms.anthropic.messages.native.transformation import (
AnthropicMessagesConfig,
)
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
from ..common_utils import (
dashscope_base_selects_anthropic_messages,
resolve_dashscope_anthropic_api_base,
)
class DashScopeAnthropicMessagesConfig(AnthropicMessagesConfig):
"""
DashScope exposes an Anthropic-compatible Messages API at
``{origin}/apps/anthropic`` for every DashScope domain, including
``dashscope.aliyuncs.com`` and ``{workspace}.{region}.maas.aliyuncs.com``.
"""
_workload_identity_eligible: ClassVar[bool] = False
@property
def custom_llm_provider(self) -> str | None:
return "dashscope"
def should_strip_billing_metadata(self) -> bool:
return True
def handles_web_search_natively(self) -> bool:
return False
def serves_api_base(self, model: str, api_base: str | None) -> bool:
return dashscope_base_selects_anthropic_messages(api_base)
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
return (
normalize_nonempty_secret_str(api_key)
or normalize_nonempty_secret_str(get_secret_str("DASHSCOPE_API_KEY"))
or normalize_nonempty_secret_str(litellm.api_key)
)
@staticmethod
def get_api_base(api_base: str | None = None) -> str:
return resolve_dashscope_anthropic_api_base(api_base)
def validate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors the legacy base override signature
model: str,
messages: list[object], # mutable-ok: mirrors the legacy base override signature
optional_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
litellm_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature
resolved_api_key: Final = self.get_api_key(api_key=api_key)
if resolved_api_key is None:
present_headers: Final = {name.lower() for name in headers}
if "x-api-key" not in present_headers and "authorization" not in present_headers:
raise AuthenticationError(
message=(
"Missing DashScope API key. Set DASHSCOPE_API_KEY or pass api_key explicitly. "
"Anthropic credentials are never sent to a DashScope endpoint."
),
llm_provider="dashscope",
model=model,
)
return super().validate_anthropic_messages_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportReturnType] # base signature uses an unparameterized dict
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=resolved_api_key,
api_base=api_base,
)
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
litellm_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
stream: bool | None = None,
) -> str:
return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages")

16
litellm/llms/dashscope/qwen_ai_platform.py
View file

@ -3,7 +3,7 @@ from typing import Final
from litellm.secret_managers.main import get_secret_str
from .chat.transformation import DashScopeChatConfig
from .common_utils import resolve_dashscope_family_rerank_api_base
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
from .embed.transformation import DashScopeEmbeddingConfig
from .image_generation.transformation import DashScopeImageGenerationConfig
from .rerank.transformation import DashScopeRerankConfig
@ -38,7 +38,7 @@ class QwenAIPlatformChatConfig(DashScopeChatConfig):
return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
@ -46,7 +46,7 @@ class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
return _require_qwen_ai_platform_api_key(api_key)
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
class QwenAIPlatformRerankConfig(DashScopeRerankConfig):
@ -63,8 +63,14 @@ class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig):
def _resolve_api_key(self, api_key: str | None) -> str:
return _require_qwen_ai_platform_api_key(api_key)
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return QWEN_AI_PLATFORM_IMAGE_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("QWEN_AI_PLATFORM_API_BASE")
class QwenAIPlatformVideoConfig(DashScopeVideoConfig):

16
litellm/llms/dashscope/qwencloud.py
View file

@ -3,7 +3,7 @@ from typing import Final
from litellm.secret_managers.main import get_secret_str
from .chat.transformation import DashScopeChatConfig
from .common_utils import resolve_dashscope_family_rerank_api_base
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
from .embed.transformation import DashScopeEmbeddingConfig
from .image_generation.transformation import DashScopeImageGenerationConfig
from .rerank.transformation import DashScopeRerankConfig
@ -38,7 +38,7 @@ class QwenCloudChatConfig(DashScopeChatConfig):
return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
@ -46,7 +46,7 @@ class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
return _require_qwencloud_api_key(api_key)
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
class QwenCloudRerankConfig(DashScopeRerankConfig):
@ -63,8 +63,14 @@ class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig):
def _resolve_api_key(self, api_key: str | None) -> str:
return _require_qwencloud_api_key(api_key)
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("QWENCLOUD_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return QWENCLOUD_IMAGE_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("QWENCLOUD_API_BASE")
class QwenCloudVideoConfig(DashScopeVideoConfig):

5
litellm/llms/dashscope/videos/transformation.py
View file

@ -38,6 +38,8 @@ from litellm.types.videos.utils import (
)
from litellm.utils import get_model_info
from ..common_utils import resolve_dashscope_endpoints
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -206,6 +208,9 @@ def _parse_task_response(raw_response: httpx.Response) -> _TaskResponse:
def _normalized_api_base(api_base: str) -> str:
endpoints: Final = resolve_dashscope_endpoints(api_base)
if endpoints is not None:
return endpoints.origin
trimmed: Final = api_base.rstrip("/")
if trimmed.endswith(DASHSCOPE_COMPATIBLE_MODE_PATH):
return trimmed[: -len(DASHSCOPE_COMPATIBLE_MODE_PATH)]

198
litellm/llms/openai/completion/guardrail_translation/handler.py
View file

@ -5,18 +5,30 @@ This module provides guardrail translation support for OpenAI's text completion
The handler processes the 'prompt' parameter for guardrails.
"""
from collections.abc import Mapping
import json
from collections.abc import Mapping, Sequence
from itertools import chain
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_chat_stream_usage,
stream_item_field,
stream_item_items,
)
from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import TextCompletionResponse
class OpenAITextCompletionHandler(BaseTranslation):
@ -32,6 +44,9 @@ class OpenAITextCompletionHandler(BaseTranslation):
- A list of strings (for batch completions)
"""
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
async def process_input_messages(
self,
data: dict,
@ -125,7 +140,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
request_data: dict[str, object] | None = None,
) -> "TextCompletionResponse":
"""
Process output response by applying guardrails to completion text.
@ -162,11 +177,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
if "response" not in request_data:
request_data["response"] = response
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
# Include model information from the response if available
@ -195,3 +206,168 @@ class OpenAITextCompletionHandler(BaseTranslation):
)
return response
async def process_output_streaming_response(
self,
responses_so_far: list[object], # mutable-ok: rewrites the caller's buffered chunks in place
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict[str, object] | None = None, # mutable-ok: the hooks share and mutate the request payload
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Sequence[object]:
chunks: Final = _text_completion_chunks(responses_so_far)
if not self._check_streaming_has_ended(chunks):
return responses_so_far
assembled: Final = self._assemble_ended_stream(chunks)
pre_guardrail_texts: Final = _choice_texts(assembled)
await self.process_output_response(
response=assembled,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if not deliver_ended_stream_rewrites:
return responses_so_far
rewrites: Final = tuple(
after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before
)
for choice_index, text in rewrites:
self._write_choice_text(chunks, choice_index, text)
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = _text_completion_chunks(responses_so_far)
if not self._check_streaming_has_ended(chunks):
return None
assembled: Final = self._assemble_ended_stream(chunks)
return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True)
def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]:
chunks: Final = _text_completion_chunks(responses_so_far)
streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks))
finished_indices: Final = frozenset(
_choice_index(choice) for choice in streamed_choices if choice.finish_reason
)
unfinished_indices: Final = tuple(
dict.fromkeys(
_choice_index(choice) for choice in streamed_choices if _choice_index(choice) not in finished_indices
)
)
if not unfinished_indices:
return tuple(responses_so_far)
terminator: Final = TextCompletionResponse(
id=chunks[0].id,
model=chunks[0].model,
object="text_completion",
choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices],
)
return (*responses_so_far, terminator)
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes]:
streamed: Final = _text_completion_chunks(responses_so_far or ())
source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model)
message_chunk: Final = TextCompletionResponse(
id=source.id,
created=source.created,
model=source.model or exc.model,
object="text_completion",
choices=[TextChoices(index=0, text=exc.message)],
).model_copy(update={"usage": None})
final_chunk: Final = TextCompletionResponse(
id=source.id,
created=source.created,
model=source.model or exc.model,
object="text_completion",
choices=[TextChoices(index=0, text="", finish_reason="content_filter")],
).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)})
return tuple(
f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk)
)
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes] | None:
from litellm.proxy.common_request_processing import sse_error_payload
_, error_obj = sse_error_payload(exc)
return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far)
return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices)
@staticmethod
def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse:
streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far))
choice_indices: Final = tuple(dict.fromkeys(_choice_index(choice) for choice in streamed_choices))
assembled_choices: Final = [
TextChoices(
index=index,
text="".join(_choice_text(choice) for choice in streamed_choices if _choice_index(choice) == index),
finish_reason=next(
(
choice.finish_reason
for choice in streamed_choices
if _choice_index(choice) == index and choice.finish_reason
),
None,
),
)
for index in choice_indices
]
first: Final = responses_so_far[0]
return TextCompletionResponse(
id=first.id,
created=first.created,
model=first.model,
object="text_completion",
choices=assembled_choices,
usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None),
)
@staticmethod
def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None:
streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far)
head, *tail = (choice for choice in streamed_choices if _choice_index(choice) == choice_index)
head.text = text
for choice in tail:
if isinstance(stream_item_field(choice, "text"), str):
choice.text = ""
def _blocked_stream_usage(original_response: object) -> Usage | None:
prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response)
if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0:
return None
return Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]:
return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse))
def _choice_text(choice: TextChoices) -> str:
text: Final = stream_item_field(choice, "text")
return text if isinstance(text, str) else ""
def _choice_index(choice: TextChoices) -> int:
return choice.index or 0
def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]:
return tuple((_choice_index(choice), _choice_text(choice)) for choice in response.choices)

8
litellm/proxy/batches_endpoints/litellm_executed_batches.py
View file

@ -35,7 +35,7 @@ from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import LITELLM_EXECUTED_BATCH_PROVIDERS, ExtractedFileData, LiteLLMBatch, LlmProviders
from litellm.utils import resolve_api_key_secret_reference
from litellm.utils import aresolve_api_key_secret_reference
if TYPE_CHECKING:
from prisma import types as prisma_types
@ -227,7 +227,7 @@ async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_clie
return response.status_code == httpx.codes.NOT_FOUND
def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None:
async def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None:
model: Final = credentials.get("model")
api_base: Final = credentials.get("api_base")
api_key: Final = credentials.get("api_key")
@ -244,7 +244,7 @@ def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str,
return None
if resolved_api_base is None:
return None
return resolved_api_base, resolve_api_key_secret_reference(
return resolved_api_base, await aresolve_api_key_secret_reference(
resolved_api_key, model=model, llm_provider=provider, upstream=credentials
)
@ -255,7 +255,7 @@ async def litellm_executed_provider_for(
provider: Final = litellm_executed_provider_of(credentials)
if provider is None:
return None
upstream: Final = _upstream_of(credentials, provider)
upstream: Final = await _upstream_of(credentials, provider)
if upstream is None:
return None
return provider if await lacks_files_api(*upstream) else None

50
litellm/proxy/common_request_processing.py
View file

@ -129,6 +129,7 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
from litellm.proxy.native_compaction import with_proxy_compaction_executor
from litellm.proxy.route_llm_request import (
endpoint_unsupported_by_model,
route_request,
)
from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id
@ -871,6 +872,17 @@ _DECODED_JSON: Final = TypeAdapter(object)
_ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _fallback_probe_data(
pristine: object, fallback_model: str, key_aliases: object, group_alias: object
) -> Mapping[str, object]:
return {
**_REQUEST_FIELDS.validate_python(pristine),
"model": resolve_served_model(
fallback_model, _DECODED_JSON.validate_python(litellm.model_alias_map), key_aliases, group_alias
),
}
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
try:
return _REQUEST_MODEL.validate_python(data.get("model"), strict=True)
@ -2676,6 +2688,17 @@ class ProxyBaseLLMRequestProcessing:
original_model,
fallback_models,
)
raw_router_settings: Final = _DECODED_JSON.validate_python(
_REQUEST_FIELDS.validate_python(rate_limited_data).get("router_settings_override")
)
group_alias: Final = _DECODED_JSON.validate_python(
_REQUEST_FIELDS.validate_python(raw_router_settings).get("model_group_alias")
if isinstance(raw_router_settings, Mapping)
else None
)
key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
object, user_api_key_dict.aliases
)
try:
for fallback_model in fallback_models:
if fallback_model == original_model:
@ -2691,6 +2714,13 @@ class ProxyBaseLLMRequestProcessing:
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:
if endpoint_unsupported_by_model(
route_type,
_fallback_probe_data(pristine, fallback_model, key_aliases, group_alias),
llm_router,
user_api_key_dict.team_id,
):
continue
return await self.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
@ -4029,13 +4059,13 @@ class ProxyBaseLLMRequestProcessing:
(assembled_response, cache_hit); the closure also runs
non-apply_guardrail post-call hooks via
_run_deferred_stream_guardrails.
- Every other anthropic_messages/aresponses stream gets a closure
that dispatches on the stored args shape, because the arming site
cannot tell the producers apart: native iterators store a single
ready-made logging coroutine to enqueue, while bridged streams
(LiteLLMCompletionStreamingIterator, and the plain SSE generator
AnthropicStreamWrapper returns for bridged /v1/messages) share
their inner CustomStreamWrapper's logging_obj and so store
- Every other anthropic_messages/aresponses/atext_completion stream
gets a closure that dispatches on the stored args shape, because the
arming site cannot tell the producers apart: native iterators store
a single ready-made logging coroutine to enqueue, while bridged
streams (LiteLLMCompletionStreamingIterator, TextCompletionStreamWrapper,
and the plain SSE generator AnthropicStreamWrapper returns for bridged
/v1/messages) share their inner CustomStreamWrapper's logging_obj and so store
(assembled_response, cache_hit); for those the closure only
dispatches success logging, matching the route's pre-existing
hook surface.
@ -4072,7 +4102,11 @@ class ProxyBaseLLMRequestProcessing:
logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete
return
if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response):
if route_type not in (
"anthropic_messages",
"aresponses",
"atext_completion",
) or not self._is_streaming_response(response):
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER

5
litellm/proxy/common_utils/callback_utils.py
View file

@ -436,10 +436,7 @@ def get_model_group_from_litellm_kwargs(kwargs: dict) -> str | None:
def get_model_group_from_request_data(data: dict) -> str | None:
_metadata: Final = data.get("metadata", None) or {}
_model_group: Final = _metadata.get("model_group", None)
if _model_group is not None:
return _model_group
return None
return _model_group if isinstance(_model_group, str) else None
def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, str]:

2
litellm/proxy/db/db_spend_update_writer.py
View file

@ -1900,8 +1900,6 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
# NOTE: Daily tag spend is committed by a separate scheduler job.
################## Daily End-User Spend Update Transactions ##################
# Aggregate all in memory daily end-user spend transactions and commit to db
await self._flush_daily_spend_queue(

12
litellm/proxy/db/prisma_client.py
View file

@ -1127,16 +1127,16 @@ class PrismaManager:
return ProxyExtrasDBManager.build_request_log_indexes()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> str | None:
"""The read-only pre-migration row check on its own, for `--check_migrations_against_existing_data`: the
refusal the migration step would print, or None when the pending schema changes
can be applied to the rows the database has."""
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
"""`--check_migrations_against_existing_data` on its own: the pending migrations applied inside a
transaction that is rolled back, as (whether they all applied, what the operator reads)."""
from litellm_proxy_extras.utils import (
ProxyExtrasDBManager, # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install
)
from pydantic import TypeAdapter
refusal: Final[object] = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install
return refusal if isinstance(refusal, str) else None
summary: Final[object] = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install
return TypeAdapter(tuple[bool, str]).validate_python(summary)
@staticmethod
def start_request_log_index_build() -> None:

36
litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py
View file

@ -82,6 +82,22 @@ def _as_plain_dict(item: object) -> Mapping[str, object]:
return cast("Mapping[str, object]", item) # cast-ok: wire rows are message/tool-call dicts
_STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object])
def _str_keyed(value: object) -> Mapping[str, object]:
return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {}
def _request_route(request_data: Mapping[str, object]) -> str | None:
litellm_metadata: Final = request_data.get("litellm_metadata")
metadata_bucket: Final = _str_keyed(
litellm_metadata if litellm_metadata is not None else request_data.get("metadata")
)
route: Final = metadata_bucket.get("user_api_key_request_route")
return route if isinstance(route, str) else None
class GraySwanGuardrailMissingSecrets(Exception):
"""Raised when the Gray Swan API key is missing."""
@ -625,11 +641,8 @@ class GraySwanGuardrail(CustomGuardrail):
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.llms import load_guardrail_translation_mappings
litellm_metadata: Final = request_data.get("litellm_metadata")
request_route: Final = (
litellm_metadata.get("user_api_key_request_route") if isinstance(litellm_metadata, Mapping) else None
)
route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None
request_route: Final = _request_route(request_data)
route_call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None
call_type: Final = (
(route_call_types[0].value if route_call_types else None)
or (logging_obj.call_type if logging_obj is not None else None)
@ -698,13 +711,12 @@ class GraySwanGuardrail(CustomGuardrail):
inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj)
litellm_metadata: Final = request_data.get("litellm_metadata")
cleaned_litellm_metadata: Final = dict(litellm_metadata) if isinstance(litellm_metadata, dict) else {}
if inbound_headers:
existing_headers: Final = cleaned_litellm_metadata.get("headers")
cleaned_litellm_metadata["headers"] = (
{**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers
)
litellm_metadata: Final = _str_keyed(request_data.get("litellm_metadata"))
cleaned_litellm_metadata: Final = (
{**litellm_metadata, "headers": {**_str_keyed(litellm_metadata.get("headers")), **inbound_headers}}
if inbound_headers
else litellm_metadata
)
if cleaned_litellm_metadata:
sanitized: Final = _DECODED_JSON.validate_python(
safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={})

194
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
View file

@ -12,11 +12,11 @@ import json
import os
import re
import time
from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass, replace
from datetime import datetime
from re import Pattern
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast
import yaml
from fastapi import HTTPException
@ -40,6 +40,7 @@ from litellm.proxy.guardrails.content_filter_data import (
)
from litellm.types.utils import (
CallTypes,
Delta,
Function,
GenericGuardrailAPIInputs,
GuardrailStatus,
@ -47,6 +48,9 @@ from litellm.types.utils import (
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
TextChoices,
TextCompletionResponse,
)
if TYPE_CHECKING:
@ -78,6 +82,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern
MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1
_DECODED_ARGUMENTS: Final = TypeAdapter(object)
_StreamItemT: Final = TypeVar("_StreamItemT")
_JSON_ENTRIES: Final = TypeAdapter(list[object])
GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b")
SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+")
@ -139,6 +144,26 @@ class _CategoryConfigView(TypedDict):
category_file: str | None
def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None:
if isinstance(choice, TextChoices):
return choice.text if isinstance(choice.text, str) else ""
delta: Final = getattr(choice, "delta", None)
if not delta:
return None
content: Final = getattr(delta, "content", None)
return content if isinstance(content, str) else ""
def _streamed_choices(item: object) -> Sequence[StreamingChoices | TextChoices]:
if isinstance(item, (ModelResponseStream, TextCompletionResponse)):
return item.choices
return ()
def _is_choiceless_chunk(item: object) -> bool:
return isinstance(item, (ModelResponseStream, TextCompletionResponse)) and not item.choices
@dataclass(frozen=True, slots=True)
class _StreamedChoiceState:
buffered_text: str = ""
@ -146,6 +171,52 @@ class _StreamedChoiceState:
committed_detections: tuple[ContentFilterDetection, ...] = ()
latest_detections: tuple[ContentFilterDetection, ...] = ()
next_trim_len: int = 0
finished: bool = False
def _end_of_stream_chunk(
last_item: object, state_by_choice: Mapping[int, _StreamedChoiceState]
) -> ModelResponseStream | TextCompletionResponse | None:
unfinished: Final = tuple(index for index, state in state_by_choice.items() if not state.finished)
if not unfinished:
return None
if isinstance(last_item, TextCompletionResponse):
return TextCompletionResponse(
id=last_item.id,
created=last_item.created,
model=last_item.model,
object="text_completion",
choices=[TextChoices(index=index, text="") for index in unfinished],
).model_copy(update={"usage": None})
if isinstance(last_item, ModelResponseStream):
return ModelResponseStream(
id=last_item.id,
created=last_item.created,
model=last_item.model,
choices=[StreamingChoices(index=index, delta=Delta(content="")) for index in unfinished],
)
return None
async def _stream_then_end_of_stream_chunk(
response: AsyncIterable[_StreamItemT], end_of_stream_chunk: Callable[[_StreamItemT], _StreamItemT | None]
) -> AsyncGenerator[tuple[_StreamItemT, bool], None]:
last_item: _StreamItemT | None = None # rebind-ok: the end-of-stream chunk copies the last chunk's id and model
held_items: tuple[_StreamItemT, ...] = () # rebind-ok: the usage chunk must stay after the end-of-stream chunk
async for item in response:
last_item = item
if _is_choiceless_chunk(item):
held_items = (*held_items, item)
continue
for held_item in held_items:
yield held_item, False
held_items = ()
yield item, False
closing_item: Final = None if last_item is None else end_of_stream_chunk(last_item)
if closing_item is not None:
yield closing_item, True
for held_item in held_items:
yield held_item, False
@dataclass(frozen=True, slots=True)
@ -2104,72 +2175,75 @@ class ContentFilterGuardrail(CustomGuardrail):
)
try:
async for item in response:
if isinstance(item, ModelResponseStream) and item.choices:
for choice in item.choices:
if not (hasattr(choice, "delta") and choice.delta):
continue
async for item, stream_ended in _stream_then_end_of_stream_chunk(
response, lambda last_item: _end_of_stream_chunk(last_item, state_by_choice)
):
for choice in _streamed_choices(item):
new_content = _streamed_choice_content(choice)
if new_content is None:
continue
choice_index = getattr(choice, "index", 0)
if not isinstance(choice_index, int):
choice_index = 0
choice_index = getattr(choice, "index", 0)
if not isinstance(choice_index, int):
choice_index = 0
content = getattr(choice.delta, "content", None)
is_final = bool(getattr(choice, "finish_reason", None))
new_content = content if isinstance(content, str) else ""
if not new_content and not is_final:
continue
is_final = stream_ended or bool(getattr(choice, "finish_reason", None))
if not new_content and not is_final:
continue
previous_state = state_by_choice.get(choice_index, _StreamedChoiceState())
buffered_text = previous_state.buffered_text + new_content
if not buffered_text:
continue
previous_state = state_by_choice.get(choice_index, _StreamedChoiceState())
buffered_text = previous_state.buffered_text + new_content
if not buffered_text:
continue
# Add a space at the end if it's the final chunk to trigger word boundaries (\b)
text_to_scan = buffered_text + (" " if is_final else "")
choice_detections: list[ContentFilterDetection] = []
scan_started = time.perf_counter()
# Add a space at the end if it's the final chunk to trigger word boundaries (\b)
text_to_scan = buffered_text + (" " if is_final else "")
choice_detections: list[ContentFilterDetection] = []
scan_started = time.perf_counter()
try:
masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
if is_final and masked_text.endswith(" "):
masked_text = masked_text[:-1]
latest_detections = tuple(choice_detections)
except HTTPException:
state_by_choice[choice_index] = replace(
previous_state, latest_detections=tuple(choice_detections)
)
raise
except Exception as e:
verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e)
masked_text = text_to_scan # Fallback to current text
latest_detections = previous_state.latest_detections
finally:
scan_seconds += time.perf_counter() - scan_started
safe_to_yield_len = max(
previous_state.yielded_masked_text_len,
len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
try:
masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
if is_final and masked_text.endswith(" "):
masked_text = masked_text[:-1]
latest_detections = tuple(choice_detections)
except HTTPException:
state_by_choice[choice_index] = replace(
previous_state, latest_detections=tuple(choice_detections)
)
choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
next_state = replace(
previous_state,
buffered_text=buffered_text,
yielded_masked_text_len=safe_to_yield_len,
latest_detections=latest_detections,
)
if is_final:
state_by_choice[choice_index] = next_state
continue
raise
except Exception as e:
verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e)
masked_text = text_to_scan # Fallback to current text
latest_detections = previous_state.latest_detections
finally:
scan_seconds += time.perf_counter() - scan_started
trim_started = time.perf_counter()
state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan)
scan_seconds += time.perf_counter() - trim_started
safe_to_yield_len = max(
previous_state.yielded_masked_text_len,
len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
)
released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
if isinstance(choice, TextChoices):
choice.text = released_text
else:
choice.delta.content = released_text
next_state = replace(
previous_state,
buffered_text=buffered_text,
yielded_masked_text_len=safe_to_yield_len,
latest_detections=latest_detections,
)
if is_final:
state_by_choice[choice_index] = replace(next_state, finished=True)
continue
yield item
else:
# Not a ModelResponseStream or no choices - yield as is
yield item
trim_started = time.perf_counter()
state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan)
scan_seconds += time.perf_counter() - trim_started
if stream_ended and not any(_streamed_choice_content(choice) for choice in _streamed_choices(item)):
continue
yield item
except HTTPException:
status = "guardrail_intervened"
raise

2
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -865,7 +865,7 @@ async def _recorded_autorouter_savings(
entity_id_field="user_id",
entity_id=user_id,
exclude_entity_ids=None,
api_key=api_key,
api_key=(api_key,) if api_key is not None else None,
start_date=start_day,
end_date=end_day,
model=None,

168
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -23,6 +23,12 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import (
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import PrismaClient
from litellm.repositories.daily_activity_buckets import (
api_key_bucket,
model_bucket,
model_group_bucket,
raw_api_key_filter,
)
from litellm.repositories.daily_activity_repository import DailyActivityRepository
from litellm.types.proxy.management_endpoints.common_daily_activity import (
BreakdownMetrics,
@ -182,7 +188,9 @@ class _AggregatedSpendData(TypedDict):
totals: ReadOnly[SpendMetrics]
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str) -> KeyMetadata:
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata:
if not api_key:
return KeyMetadata(key_alias="Unattributed", key_exists=False)
meta: Final = api_key_metadata.get(api_key)
return KeyMetadata(
key_alias=meta.key_alias if meta is not None else None,
@ -302,13 +310,14 @@ def update_breakdown_metrics(
parent bucket but never appear as an api_key row, and are kept out of the
per-request provider breakdown."""
api_key: Final = api_key_bucket(record.api_key)
is_ptu_sentinel: Final = record.api_key == PTU_SENTINEL_API_KEY
# A PTU sentinel row keys on the deployment id so a rename cannot move it, and carries
# the operator-facing name in model_group. The breakdown key is rendered directly as a
# label, so display the name; two deployments sharing one name merge here, which is
# what the write path used to do by collapsing them into a single row.
model_key: Final = (record.model_group or record.model) if is_ptu_sentinel else record.model
model_key: Final = (record.model_group or record.model) if is_ptu_sentinel else model_bucket(record.model)
# Update model breakdown
if model_key and model_key not in breakdown.models:
@ -321,18 +330,18 @@ def update_breakdown_metrics(
if not is_ptu_sentinel:
# Update API key breakdown for this model
if record.api_key not in breakdown.models[model_key].api_key_breakdown:
breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
if api_key not in breakdown.models[model_key].api_key_breakdown:
breakdown.models[model_key].api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics,
breakdown.models[model_key].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.models[model_key].api_key_breakdown[api_key].metrics,
record,
)
# Update model group breakdown
model_group_key: Final = record.model_group or record.model
model_group_key: Final = model_group_bucket(record.model_group, record.model)
if model_group_key and model_group_key not in breakdown.model_groups:
breakdown.model_groups[model_group_key] = MetricWithMetadata(
metrics=SpendMetrics(),
@ -345,13 +354,13 @@ def update_breakdown_metrics(
if not is_ptu_sentinel:
# Update API key breakdown for this model
if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown:
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
if api_key not in breakdown.model_groups[model_group_key].api_key_breakdown:
breakdown.model_groups[model_group_key].api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics,
breakdown.model_groups[model_group_key].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.model_groups[model_group_key].api_key_breakdown[api_key].metrics,
record,
)
@ -365,21 +374,20 @@ def update_breakdown_metrics(
breakdown.mcp_servers[record.mcp_namespaced_tool_name].metrics, record
)
# Update API key breakdown for this MCP server
if record.api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown:
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
if not is_ptu_sentinel:
# Update API key breakdown for this MCP server
if api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown:
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key] = (
KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
)
)
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[
record.api_key
].metrics = update_metrics(
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key].metrics,
record,
)
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key].metrics,
record,
)
if not is_ptu_sentinel:
# Update provider breakdown
@ -392,13 +400,13 @@ def update_breakdown_metrics(
breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record)
# Update API key breakdown for this provider
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
if api_key not in breakdown.providers[provider].api_key_breakdown:
breakdown.providers[provider].api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
breakdown.providers[provider].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.providers[provider].api_key_breakdown[api_key].metrics,
record,
)
@ -413,25 +421,26 @@ def update_breakdown_metrics(
breakdown.endpoints[record.endpoint].metrics, record
)
# Update API key breakdown for this endpoint
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
if not is_ptu_sentinel:
# Update API key breakdown for this endpoint
if api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
breakdown.endpoints[record.endpoint].api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.endpoints[record.endpoint].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.endpoints[record.endpoint].api_key_breakdown[api_key].metrics,
record,
)
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
record,
)
if not is_ptu_sentinel:
# Update api key breakdown
if record.api_key not in breakdown.api_keys:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
if api_key not in breakdown.api_keys:
breakdown.api_keys[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
breakdown.api_keys[api_key].metrics = update_metrics(breakdown.api_keys[api_key].metrics, record)
# Update entity-specific metrics if entity_id_field is provided
if entity_id_field:
@ -445,13 +454,13 @@ def update_breakdown_metrics(
if not is_ptu_sentinel:
# Update API key breakdown for this entity
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
if api_key not in breakdown.entities[entity_value].api_key_breakdown:
breakdown.entities[entity_value].api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
breakdown.entities[entity_value].api_key_breakdown[api_key].metrics = update_metrics(
breakdown.entities[entity_value].api_key_breakdown[api_key].metrics,
record,
)
@ -566,7 +575,7 @@ def daily_activity_scope(
entity_id_field: str,
entity_id: str | list[str] | None,
exclude_entity_ids: list[str] | None,
api_key: str | list[str] | None,
api_key: str | Sequence[str] | None,
start_date: str,
end_date: str,
model: str | None,
@ -577,9 +586,7 @@ def daily_activity_scope(
entity_ids: tuple[str, ...] | None = (
(entity_id,) if isinstance(entity_id, str) else tuple(entity_id) if entity_id is not None else None
)
api_keys: tuple[str, ...] | None = (
None if api_key in (None, "") else (api_key,) if isinstance(api_key, str) else tuple(api_key)
)
api_keys: Final = raw_api_key_filter(api_key)
return DailyActivityScope(
table=table_value,
entity_id_field=entity_id_field,
@ -715,6 +722,29 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
def sum_spend_metrics(metrics: Sequence[SpendMetrics]) -> SpendMetrics:
return SpendMetrics(
spend=sum(metric.spend for metric in metrics),
flat_cost=sum(metric.flat_cost for metric in metrics),
ptu_hours=sum(metric.ptu_hours for metric in metrics),
prompt_tokens=sum(metric.prompt_tokens for metric in metrics),
completion_tokens=sum(metric.completion_tokens for metric in metrics),
cache_read_input_tokens=sum(metric.cache_read_input_tokens for metric in metrics),
cache_creation_input_tokens=sum(metric.cache_creation_input_tokens for metric in metrics),
compression_saved_tokens=sum(metric.compression_saved_tokens for metric in metrics),
compression_savings_spend=sum(metric.compression_savings_spend for metric in metrics),
prompt_caching_savings_spend=sum(metric.prompt_caching_savings_spend for metric in metrics),
gateway_injected_caching_savings_spend=sum(metric.gateway_injected_caching_savings_spend for metric in metrics),
autorouter_savings_spend=sum(metric.autorouter_savings_spend for metric in metrics),
total_tokens=sum(metric.total_tokens for metric in metrics),
successful_requests=sum(metric.successful_requests for metric in metrics),
failed_requests=sum(metric.failed_requests for metric in metrics),
api_requests=sum(metric.api_requests for metric in metrics),
total_response_time_ms=sum(metric.total_response_time_ms for metric in metrics),
timed_requests=sum(metric.timed_requests for metric in metrics),
)
def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row.
@ -753,8 +783,7 @@ def _aggregate_grouping_sets_records_sync(
Each row carries a `group_level` bitmask (from Postgres GROUPING()) that
identifies which rollup level it belongs to. We dispatch the row's
pre-aggregated metrics straight into the matching bucket — no per-row
summing in Python and no nested update_metrics calls.
pre-aggregated metrics into the matching bucket, combining normalized names.
"""
total_metrics = SpendMetrics()
grouped_data: Final[dict[str, GroupedData]] = {}
@ -771,20 +800,23 @@ def _aggregate_grouping_sets_records_sync(
if existing is None:
target[key] = MetricWithMetadata(metrics=metrics, metadata={})
else:
existing.metrics = metrics
existing.metrics = sum_spend_metrics((existing.metrics, metrics))
def assign_api_key_breakdown(
target: dict[str, MetricWithMetadata],
parent_key: str,
api_key: str,
api_key: str | None,
metrics: SpendMetrics,
) -> None:
parent = target.get(parent_key)
if parent is None:
parent = MetricWithMetadata(metrics=SpendMetrics(), metadata={})
target[parent_key] = parent
parent.api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, api_key)
bucket_key: Final = api_key_bucket(api_key)
previous: Final = parent.api_key_breakdown.get(bucket_key)
parent.api_key_breakdown[bucket_key] = KeyMetricWithMetadata(
metrics=sum_spend_metrics((previous.metrics, metrics)) if previous is not None else metrics,
metadata=_key_metadata(api_key_metadata, api_key),
)
for record in records:
@ -803,25 +835,23 @@ def _aggregate_grouping_sets_records_sync(
breakdown = ensure_date(record.date)["breakdown"]
if level == _GROUP_DATE_API_KEY:
if record.api_key and not is_ptu_sentinel:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
if not is_ptu_sentinel:
breakdown.api_keys[api_key_bucket(record.api_key)] = KeyMetricWithMetadata(
metrics=metrics,
metadata=_key_metadata(api_key_metadata, record.api_key),
)
elif level == _GROUP_DATE_MODEL:
if record.model:
assign_metric_with_metadata(breakdown.models, record.model, metrics)
assign_metric_with_metadata(breakdown.models, model_bucket(record.model), metrics)
elif level == _GROUP_DATE_MODEL_API_KEY:
if record.model and record.api_key and not is_ptu_sentinel:
assign_api_key_breakdown(breakdown.models, record.model, record.api_key, metrics)
if not is_ptu_sentinel:
assign_api_key_breakdown(breakdown.models, model_bucket(record.model), record.api_key, metrics)
elif level == _GROUP_DATE_MODEL_GROUP:
if record.model_group:
assign_metric_with_metadata(breakdown.model_groups, record.model_group, metrics)
assign_metric_with_metadata(breakdown.model_groups, model_group_bucket(record.model_group, None), metrics)
elif level == _GROUP_DATE_MODEL_GROUP_API_KEY:
if record.model_group and record.api_key and not is_ptu_sentinel:
if not is_ptu_sentinel:
assign_api_key_breakdown(
breakdown.model_groups,
record.model_group,
model_group_bucket(record.model_group, None),
record.api_key,
metrics,
)
@ -836,14 +866,14 @@ def _aggregate_grouping_sets_records_sync(
provider = record.custom_llm_provider or "unknown"
assign_metric_with_metadata(breakdown.providers, provider, provider_metrics)
elif level == _GROUP_DATE_PROVIDER_API_KEY:
if record.api_key and not is_ptu_sentinel:
if not is_ptu_sentinel:
provider = record.custom_llm_provider or "unknown"
assign_api_key_breakdown(breakdown.providers, provider, record.api_key, metrics)
elif level == _GROUP_DATE_MCP:
if record.mcp_namespaced_tool_name:
assign_metric_with_metadata(breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics)
elif level == _GROUP_DATE_MCP_API_KEY:
if record.mcp_namespaced_tool_name and record.api_key:
if record.mcp_namespaced_tool_name and not is_ptu_sentinel:
assign_api_key_breakdown(
breakdown.mcp_servers,
record.mcp_namespaced_tool_name,
@ -854,7 +884,7 @@ def _aggregate_grouping_sets_records_sync(
if record.endpoint:
assign_metric_with_metadata(breakdown.endpoints, record.endpoint, metrics)
elif level == _GROUP_DATE_ENDPOINT_API_KEY:
if record.endpoint and record.api_key:
if record.endpoint and not is_ptu_sentinel:
assign_api_key_breakdown(breakdown.endpoints, record.endpoint, record.api_key, metrics)
results: Final = [
@ -1013,8 +1043,8 @@ def _fold_entity_rollups_sync(
metrics = _record_to_spend_metrics(row)
if row.api_key_rolled:
bucket.metrics = metrics
elif row.api_key and row.api_key != PTU_SENTINEL_API_KEY:
bucket.api_key_breakdown[row.api_key] = KeyMetricWithMetadata(
elif row.api_key != PTU_SENTINEL_API_KEY:
bucket.api_key_breakdown[api_key_bucket(row.api_key)] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, row.api_key)
)

34
litellm/proxy/management_endpoints/daily_activity_routes.py
View file

@ -22,6 +22,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
parse_canonical_date_range,
raise_public,
spend_logs_window,
sum_spend_metrics,
)
from litellm.proxy.management_endpoints.daily_activity_scopes import (
AGENT_RESOLVER,
@ -38,6 +39,7 @@ from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, w
from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw
from litellm.repositories.daily_activity_buckets import api_key_bucket
from litellm.repositories.daily_activity_repository import DailyActivityRepository
from litellm.router import Router
from litellm.types.proxy.management_endpoints.common_daily_activity import (
@ -106,29 +108,9 @@ async def _resolved_scope(
return result
def _sum_metrics(metrics: Sequence[SpendMetrics]) -> SpendMetrics:
return SpendMetrics(
spend=sum(metric.spend for metric in metrics),
flat_cost=sum(metric.flat_cost for metric in metrics),
prompt_tokens=sum(metric.prompt_tokens for metric in metrics),
completion_tokens=sum(metric.completion_tokens for metric in metrics),
cache_read_input_tokens=sum(metric.cache_read_input_tokens for metric in metrics),
cache_creation_input_tokens=sum(metric.cache_creation_input_tokens for metric in metrics),
compression_saved_tokens=sum(metric.compression_saved_tokens for metric in metrics),
compression_savings_spend=sum(metric.compression_savings_spend for metric in metrics),
prompt_caching_savings_spend=sum(metric.prompt_caching_savings_spend for metric in metrics),
gateway_injected_caching_savings_spend=sum(metric.gateway_injected_caching_savings_spend for metric in metrics),
autorouter_savings_spend=sum(metric.autorouter_savings_spend for metric in metrics),
total_tokens=sum(metric.total_tokens for metric in metrics),
successful_requests=sum(metric.successful_requests for metric in metrics),
failed_requests=sum(metric.failed_requests for metric in metrics),
api_requests=sum(metric.api_requests for metric in metrics),
total_response_time_ms=sum(metric.total_response_time_ms for metric in metrics),
timed_requests=sum(metric.timed_requests for metric in metrics),
)
def _key_metadata(api_key: str, metadata: Mapping[str, KeyMetadataRow]) -> KeyMetadata:
if not api_key:
return KeyMetadata(key_alias="Unattributed", key_exists=False)
row: Final[KeyMetadataRow | None] = metadata.get(api_key)
if row is None:
return KeyMetadata()
@ -143,7 +125,7 @@ def _key_metadata(api_key: str, metadata: Mapping[str, KeyMetadataRow]) -> KeyMe
def _key_activity_row(row: KeySpendRow, metadata: Mapping[str, KeyMetadataRow]) -> KeySpendActivityRow:
return KeySpendActivityRow(
api_key=row.api_key,
api_key=api_key_bucket(row.api_key),
metrics=KeySpendMetrics(
spend=row.spend,
prompt_tokens=row.prompt_tokens,
@ -168,7 +150,7 @@ async def _key_activity_rows(
frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date))
)
metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata(
frozenset(row.api_key for row in rows),
frozenset(row.api_key for row in rows if row.api_key),
spend_window,
)
return [_key_activity_row(row, metadata) for row in rows]
@ -231,11 +213,11 @@ def _fold_key_metrics(api_key: str, response: SpendAnalyticsPaginatedResponse) -
metadata: Final[KeyMetadata] = next(
day.breakdown.api_keys[api_key].metadata for day in response.results if api_key in day.breakdown.api_keys
)
return KeyActivityRow(api_key=api_key, metrics=_sum_metrics(metrics), metadata=metadata)
return KeyActivityRow(api_key=api_key, metrics=sum_spend_metrics(metrics), metadata=metadata)
def _search_rows(keys: Sequence[str], response: SpendAnalyticsPaginatedResponse) -> list[KeyActivityRow]:
return [row for key in keys if (row := _fold_key_metrics(key, response)) is not None]
return [row for key in keys if (row := _fold_key_metrics(api_key_bucket(key), response)) is not None]
def _csv_cell(value: object) -> object:

9
litellm/proxy/management_endpoints/daily_activity_scopes.py
View file

@ -17,6 +17,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import resolve_or
from litellm.proxy.management_endpoints.tag_management_endpoints import get_tag_daily_activity_api_key_filter
from litellm.proxy.management_endpoints.team_endpoints import resolve_team_daily_activity_scope
from litellm.proxy.utils import PrismaClient
from litellm.repositories.daily_activity_buckets import raw_api_key_filter
from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable
@ -77,13 +78,7 @@ def _build_scope(
) -> ResolvedScope:
start_date: Final[str] = query.start_date or ""
end_date: Final[str] = query.end_date or ""
api_keys: Final[tuple[str, ...] | None] = (
None
if api_key_filter is None or api_key_filter == ""
else (api_key_filter,)
if isinstance(api_key_filter, str)
else tuple(api_key_filter)
)
api_keys: Final = raw_api_key_filter(api_key_filter)
return ResolvedScope(
scope=DailyActivityScope(
table=resolver.table,

8
litellm/proxy/management_endpoints/tag_management_endpoints.py
View file

@ -31,6 +31,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity,
)
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.table_repositories import (
DailyTagSpendRepository,
@ -198,16 +199,17 @@ async def get_tag_daily_activity_api_key_filter(
user_api_key_dict: UserAPIKeyAuth,
requested_api_key: str | None,
) -> str | list[str] | None:
selected_key: Final = api_key_from_bucket(requested_api_key) if requested_api_key else None
user_role: Final = user_api_key_dict.user_role
if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role):
return requested_api_key
return [selected_key] if selected_key is not None else None
scoped_api_keys: Final = await _get_internal_user_api_keys(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
if requested_api_key is not None:
return requested_api_key if requested_api_key in scoped_api_keys else []
if selected_key is not None:
return [selected_key] if selected_key in scoped_api_keys else []
return scoped_api_keys

8
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -226,6 +226,7 @@ from litellm.proxy.management_helpers.utils import (
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import delete_many_in, find_many_in
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
@ -7329,11 +7330,12 @@ async def resolve_team_daily_activity_scope(
if not user_api_keys:
user_api_keys = []
selected_key: Final = api_key_from_bucket(api_key) if api_key else None
final_api_key_filter: str | list[str] | None
if user_api_keys is None:
final_api_key_filter = api_key
elif api_key:
final_api_key_filter = api_key if api_key in user_api_keys else []
final_api_key_filter = [selected_key] if selected_key is not None else None
elif selected_key is not None:
final_api_key_filter = [selected_key] if selected_key in user_api_keys else []
else:
final_api_key_filter = user_api_keys

9
litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py
View file

@ -12,6 +12,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.secret_managers.main import (
GOOGLE_OIDC_SECRET_PREFIX,
UnresolvableSecretReference,
get_secret_str,
resolve_local_oidc_reference,
@ -47,6 +48,14 @@ def _credential_identity(credentials: VERTEX_CREDENTIALS_TYPES | None) -> str |
def _resolve_oidc_reference(credential: str | None) -> str | None:
if credential is None:
return None
if credential.startswith(GOOGLE_OIDC_SECRET_PREFIX):
raise HTTPException(
status_code=401,
detail=(
"pass-through credential oidc/google secret reference could not be resolved: pass-through routes "
"do not fetch Google ID tokens. Use a static key, oidc/file/<absolute path> or oidc/env/<VARIABLE_NAME>."
),
)
match resolve_local_oidc_reference(credential):
case UnresolvableSecretReference(kind=kind, reason=reason):
raise HTTPException(

18
litellm/proxy/proxy_cli.py
View file

@ -276,18 +276,19 @@ class ProxyInitializationHelpers:
@staticmethod
def run_migration_data_check(use_migrate: bool) -> int:
"""`--check_migrations_against_existing_data`: the pre-migration row check against DATABASE_URL, nothing applied.
0 when the pending schema changes fit the rows, 2 with the refusal when they do not."""
"""`--check_migrations_against_existing_data`: the pending migrations applied to DATABASE_URL inside one
transaction that is rolled back, so Postgres reports what the upgrade would hit. 0 when they all apply,
2 with Postgres's error when one does not."""
if os.getenv("DATABASE_URL") is None:
raise click.UsageError("--check_migrations_against_existing_data requires DATABASE_URL")
from litellm.proxy.db.prisma_client import PrismaManager
refusal: Final[str | None] = PrismaManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install
if refusal is None:
print("LiteLLM Proxy: the pending migrations and indexes can be applied to the rows this database has.")
ok, summary = PrismaManager.migration_dry_run_summary(use_migrate=use_migrate)
if ok:
print(f"LiteLLM Proxy: {summary}")
return 0
print(
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {refusal}\033[0m", file=sys.stderr, flush=True
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {summary}\033[0m", file=sys.stderr, flush=True
)
return 2
@ -1059,8 +1060,9 @@ class ProxyInitializationHelpers:
"--check_migrations_against_existing_data",
is_flag=True,
default=False,
help="Check, read-only, that the rows in DATABASE_URL can take this version's pending migrations and indexes, "
"print what would block them, then exit without applying anything. Exit code 2 when something blocks.",
help="Dry-run this version's pending migrations on DATABASE_URL: the real SQL inside one transaction that is "
"rolled back, so Postgres reports the row or schema that would make the upgrade fail, then exit. Holds the same "
"locks the upgrade will. Exit code 2 when a migration fails.",
)
@click.option(
"--keepalive_timeout",

62
litellm/proxy/proxy_server.py
View file

@ -79,6 +79,7 @@ from litellm.constants import (
DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
DEFAULT_SHARED_HEALTH_CHECK_TTL,
DEFAULT_SLACK_ALERTING_THRESHOLD,
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS,
INTERNAL_CALLBACK_NAMES,
LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS,
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS,
@ -98,6 +99,7 @@ from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
initialize_sentry,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.token_counter import (
@ -207,7 +209,7 @@ from litellm.types.utils import (
is_server_derived_pricing_key,
pricing_override_fields,
)
from litellm.utils import cost_map_omits_token_price, load_credentials_from_list, with_named_credential
from litellm.utils import acredential_updates, cost_map_omits_token_price, with_named_credential
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -1275,15 +1277,44 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
async def _flush_spend_on_shutdown() -> None:
from litellm.proxy.utils import update_daily_tag_spend
await flush_spend_counters_on_shutdown()
if prisma_client is not None:
await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
await _flush_spend_logs_queue_on_shutdown()
async def _flush_spend_around_logging_drain_on_shutdown(
flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown,
*,
worker: LoggingWorker = GLOBAL_LOGGING_WORKER,
clock: Callable[[], float] = time.monotonic,
drain_tasks: Callable[[float], Awaitable[None]] = drain_logging_tasks,
) -> None:
await flush_spend()
await drain_logging_tasks()
remaining: Final = GracefulShutdownManager.remaining(clock=clock)
callback_timeout: Final = max(0.0, remaining - GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS)
verbose_proxy_logger.info(
"Draining success callbacks before the final spend flush remaining_s=%.3f callback_timeout_s=%.3f",
remaining,
callback_timeout,
)
async def drain_callbacks() -> None:
await drain_tasks(callback_timeout)
await worker.flush()
try:
await asyncio.wait_for(drain_callbacks(), timeout=callback_timeout)
except asyncio.TimeoutError:
verbose_proxy_logger.warning(
"Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete",
worker.pending_count,
)
finally:
await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock))
await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock))
await flush_spend()
@ -1446,6 +1477,13 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
@asynccontextmanager
async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
with GracefulShutdownManager.capture_sigterm():
async with _proxy_startup_event(app) as state:
yield state
@asynccontextmanager
async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
global \
prisma_client, \
master_key, \
@ -14815,10 +14853,10 @@ def _provider_count_outcome(
return result
def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]:
async def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]:
call_params: Final = _LITELLM_PARAMS_COPY.validate_python(litellm_params)
try:
load_credentials_from_list(call_params)
return {**call_params, **(await acredential_updates(call_params))}
except litellm.AuthenticationError as e:
raise ProxyException(
message=e.message,
@ -14826,13 +14864,19 @@ def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]:
param="model",
code=401,
) from e
return call_params
except litellm.ServiceUnavailableError as e:
raise ProxyException(
message=e.message,
type="service_unavailable",
param="model",
code=503,
) from e
def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None:
async def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None:
if deployment is None:
return None
return {**deployment, "litellm_params": _with_resolved_api_key(deployment.get("litellm_params", {}))}
return {**deployment, "litellm_params": await _with_resolved_api_key(deployment.get("litellm_params", {}))}
_DEPLOYMENT_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
@ -14895,7 +14939,7 @@ async def _try_provider_token_count(
model_to_use=model_to_use or "",
messages=messages,
contents=contents,
deployment=_with_resolved_credentials(deployment_in_region),
deployment=await _with_resolved_credentials(deployment_in_region),
request_model=request_model,
tools=tools,
system=system,
@ -14929,7 +14973,7 @@ async def _try_native_responses_count(
model_to_use=model_to_use,
request=responses_request,
policy=policy,
deployment=_with_resolved_credentials(deployment),
deployment=await _with_resolved_credentials(deployment),
request_model=request_model,
)
)

16
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -34,6 +34,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
from litellm.caching.dual_cache import DualCache
from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES
from litellm.exceptions import EndpointNotSupportedByModelError
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
@ -2086,6 +2087,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec
def _responses_ws_failure_frame(failure: Exception) -> str:
if isinstance(failure, EndpointNotSupportedByModelError):
return json.dumps(
{
"type": "error",
"status": 400,
"error": {"type": "invalid_request_error", "message": failure.message.split("\n\n", 1)[0]},
}
)
proxy_exception: Final = handle_exception_on_proxy(failure)
status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500
error_type: Final = (
@ -2444,7 +2453,12 @@ async def _responses_websocket_session(
with contextlib.suppress(Exception):
await websocket.send_text(_responses_ws_failure_frame(e))
with contextlib.suppress(Exception):
await websocket.close(code=1011, reason="Internal server error")
await websocket.close(
code=1008 if isinstance(e, EndpointNotSupportedByModelError) else 1011,
reason="Pre-call error"
if isinstance(e, EndpointNotSupportedByModelError)
else "Internal server error",
)
return crash_row
return bill_socket_failure(outcome) if isinstance(outcome, Exception) else None

22
litellm/proxy/route_llm_request.py
View file

@ -431,6 +431,28 @@ def raise_if_required_body_param_missing(
)
def endpoint_unsupported_by_model(
route_type: str,
data: Mapping[str, object],
llm_router: LitellmRouter | None,
team_id: str | None,
) -> bool:
model: Final = data.get("model")
if not isinstance(model, str) or llm_router is None:
return False
metadata: Final = data.get("metadata")
request_kwargs: Final = {
**data,
"metadata": {
**(metadata if isinstance(metadata, Mapping) else {}),
"user_api_key_team_id": team_id,
},
}
return bool(llm_router.get_routable_deployments(model, request_kwargs)) and not llm_router.get_routable_deployments(
model, request_kwargs, call_type=route_type
)
class MockTestingParamsDisabledError(HTTPException):
def __init__(self, params: tuple[str, ...]):
super().__init__(

58
litellm/proxy/shutdown/graceful_shutdown_manager.py
View file

@ -5,7 +5,8 @@ Kubernetes terminates a pod by sending ``SIGTERM`` and, after
``terminationGracePeriodSeconds``, ``SIGKILL``. On ``SIGTERM`` uvicorn stops
accepting connections and waits up to ``GRACEFUL_SHUTDOWN_TIMEOUT`` for open
ones, cancelling whatever is still running, before the lifespan teardown that
runs this drain.
runs this drain. The lifespan registers a SIGTERM wrapper before serving
requests so callback draining can subtract the preceding connection drain.
A fixed ``preStop`` sleep can not solve this: it has to be sized for the
*worst-case* request, so it either wastes time on every routine shutdown or is
@ -20,8 +21,12 @@ granularity of ``InFlightRequestsMiddleware``.
import asyncio
import os
import signal
import threading
import time
from collections.abc import Callable
from collections.abc import Callable, Generator
from contextlib import contextmanager
from types import FrameType
from typing import Final
from litellm._logging import verbose_proxy_logger
@ -44,6 +49,7 @@ class GracefulShutdownManager:
_is_shutting_down: bool = False
_shutdown_started_at: float | None = None
_drain_performed: bool = False
_sigterm_handler: Callable[[int, FrameType | None], None] | None = None
@classmethod
def is_shutting_down(cls) -> bool:
@ -71,7 +77,7 @@ class GracefulShutdownManager:
return DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT
@classmethod
def start_shutdown(cls) -> None:
def start_shutdown(cls, clock: Callable[[], float] = time.monotonic) -> None:
"""
Mark the worker as draining. Idempotent — repeated calls (e.g. SIGTERM
followed by a preStop hit on /health/drain) do not reset the clock.
@ -79,12 +85,42 @@ class GracefulShutdownManager:
if cls._is_shutting_down:
return
cls._is_shutting_down = True
cls._shutdown_started_at = time.monotonic()
cls._shutdown_started_at = clock()
verbose_proxy_logger.info(
"graceful_shutdown_started in_flight_requests=%s",
get_in_flight_requests(),
)
@classmethod
def remaining(cls, cap: float | None = None, *, clock: Callable[[], float] = time.monotonic) -> float:
elapsed: Final = 0.0 if cls._shutdown_started_at is None else clock() - cls._shutdown_started_at
remaining: Final = max(0.0, cls.get_timeout() - elapsed)
return remaining if cap is None else min(max(0.0, cap), remaining)
@classmethod
@contextmanager
def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Generator[None, None, None]:
if threading.current_thread() is not threading.main_thread():
yield
return
previous: Final = signal.getsignal(signal.SIGTERM)
if not callable(previous) or previous is cls._sigterm_handler:
yield
return
def handle_sigterm(signum: int, frame: FrameType | None) -> None:
cls.start_shutdown(clock)
previous(signum, frame)
cls._sigterm_handler = handle_sigterm
signal.signal(signal.SIGTERM, handle_sigterm)
try:
yield
finally:
if signal.getsignal(signal.SIGTERM) is handle_sigterm:
signal.signal(signal.SIGTERM, previous)
cls._sigterm_handler = None
@classmethod
async def wait_for_drain(
cls,
@ -93,13 +129,14 @@ class GracefulShutdownManager:
count_fn: Callable[[], int] | None = None,
poll_interval: float = _DRAIN_POLL_INTERVAL,
log_interval: float = _DRAIN_LOG_INTERVAL,
clock: Callable[[], float] = time.monotonic,
) -> int:
"""
Poll the in-flight request counter until it reaches the drain target or
``timeout`` seconds elapse.
Args:
timeout: Max seconds to wait. Defaults to ``get_timeout()``.
timeout: Max seconds to wait. Defaults to the remaining shutdown budget.
exclude_self: When the caller is itself an in-flight HTTP request
(the /health/drain endpoint), set this so the caller's own
request is not counted as outstanding work.
@ -107,6 +144,7 @@ class GracefulShutdownManager:
live ``InFlightRequestsMiddleware`` counter; injectable for tests.
poll_interval: Seconds between counter polls.
log_interval: Minimum seconds between ``drain_waiting`` log lines.
clock: Monotonic clock shared with the shutdown timestamp.
Returns:
Number of requests that drained while waiting (>= 0).
@ -120,7 +158,7 @@ class GracefulShutdownManager:
cls._drain_performed = True
if timeout is None:
timeout = cls.get_timeout()
timeout = cls.remaining(clock=clock)
if count_fn is None:
count_fn = get_in_flight_requests
@ -128,7 +166,7 @@ class GracefulShutdownManager:
# and so counts itself; treat <=1 as "drained" in that case.
target: Final = 1 if exclude_self else 0
start: Final = time.monotonic()
start: Final = clock()
initial: Final = count_fn()
last_log = start
@ -142,11 +180,11 @@ class GracefulShutdownManager:
verbose_proxy_logger.info(
"graceful_shutdown_complete drained_requests=%s elapsed_s=%.2f",
drained,
time.monotonic() - start,
clock() - start,
)
return drained
elapsed = time.monotonic() - start
elapsed = clock() - start
if elapsed >= timeout:
verbose_proxy_logger.warning(
"graceful_shutdown_timeout in_flight_requests=%s elapsed_s=%.2f "
@ -157,7 +195,7 @@ class GracefulShutdownManager:
)
return max(0, initial - current)
now = time.monotonic()
now = clock()
if now - last_log >= log_interval:
verbose_proxy_logger.info(
"drain_waiting in_flight_requests=%s elapsed_s=%.2f",

8
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -464,7 +464,11 @@ def _sl_attribution_fallback(
) -> str:
if standard_logging_payload is None:
return ""
return standard_logging_payload.get(field) or ""
return _string_or_empty(standard_logging_payload.get(field))
def _string_or_empty(value: object) -> str:
return value if isinstance(value, str) else ""
def _deployment_provider(deployment: DeploymentTypedDict) -> str | None:
@ -682,7 +686,7 @@ def get_logging_payload(
_model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback(
standard_logging_payload, "model_id"
)
_model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback(
_model_group: Final = _string_or_empty(metadata.get("model_group")) or _sl_attribution_fallback(
standard_logging_payload, "model_group"
)
_api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback(

4
litellm/realtime_api/main.py
View file

@ -35,8 +35,8 @@ from litellm.types.utils import CallTypes, LlmProviders
from litellm.utils import (
ModelInfoLookupMiss,
ProviderConfigManager,
aresolve_api_key_secret_reference,
lookup_model_info_for_registration,
resolve_api_key_secret_reference,
)
from ..litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params
@ -845,7 +845,7 @@ async def realtime_health_check(
resolved_api_base: Final = cast( # cast-ok: provider parameters expose optional string endpoints
str | None, api_base or resolved_params.get("api_base")
)
resolved_api_key: Final = resolve_api_key_secret_reference(
resolved_api_key: Final = await aresolve_api_key_secret_reference(
cast( # cast-ok: provider parameters expose optional string credentials
str | None, api_key or resolved_params.get("api_key")
),

39
litellm/repositories/daily_activity_buckets.py Normal file
View file

@ -0,0 +1,39 @@
from collections.abc import Sequence
from typing import Final
from litellm.constants import UNKNOWN_MODEL_SPEND_LOG_MODEL
UNKNOWN_MODEL_GROUP: Final = "unknown-model-group"
UNATTRIBUTED_API_KEY: Final = "unattributed"
_RAW_API_KEY_PREFIX: Final = "raw:"
def model_bucket(model: str | None) -> str:
return model or UNKNOWN_MODEL_SPEND_LOG_MODEL
def model_group_bucket(model_group: str | None, model: str | None) -> str:
return model_group or model or UNKNOWN_MODEL_GROUP
def api_key_bucket(api_key: str | None) -> str:
if not api_key:
return UNATTRIBUTED_API_KEY
if api_key == UNATTRIBUTED_API_KEY or api_key.startswith(_RAW_API_KEY_PREFIX):
return _RAW_API_KEY_PREFIX + api_key
return api_key
def api_key_from_bucket(api_key: str) -> str:
if api_key == UNATTRIBUTED_API_KEY:
return ""
return api_key.removeprefix(_RAW_API_KEY_PREFIX)
def raw_api_key_filter(api_key: str | Sequence[str] | None) -> tuple[str, ...] | None:
"""String filters are reporting IDs; sequences are raw keys resolved by authorization."""
if api_key is None or api_key == "":
return None
if isinstance(api_key, str):
return (api_key_from_bucket(api_key),)
return tuple(api_key)

9
litellm/repositories/daily_activity_repository.py
View file

@ -10,6 +10,7 @@ from typing_extensions import assert_never
from litellm import constants
from litellm._logging import verbose_proxy_logger
from litellm.constants import UNKNOWN_MODEL_SPEND_LOG_MODEL
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.daily_activity_sql import (
ExportCursor,
@ -344,7 +345,13 @@ class DailyActivityRepository:
"date": {"gte": adjusted_start, "lte": adjusted_end},
**({scope.entity_id_field: {"in": list(scope.entity_ids)}} if scope.entity_ids is not None else {}),
**exclusion_filter,
**({"model": scope.model} if scope.model else {}),
**(
{"AND": [{"OR": [{"model": None}, {"model": ""}, {"model": scope.model}]}]}
if scope.model == UNKNOWN_MODEL_SPEND_LOG_MODEL
else {"model": scope.model}
if scope.model
else {}
),
**({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}),
}
count, rows = await asyncio.gather(

60
litellm/repositories/daily_activity_sql.py
View file

@ -8,11 +8,18 @@ from typing import Final
from typing_extensions import assert_never
from litellm import constants
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.constants import PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
from litellm.repositories.daily_activity_buckets import UNATTRIBUTED_API_KEY, UNKNOWN_MODEL_GROUP
from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType
_API_KEY_ROLLED_UP_BIT: Final = 32
_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
_MODEL_EXPR: Final = f"COALESCE(NULLIF(model, ''), '{UNKNOWN_MODEL_SPEND_LOG_MODEL}')"
_MODEL_GROUP_EXPR: Final = f"COALESCE(NULLIF(model_group, ''), NULLIF(model, ''), '{UNKNOWN_MODEL_GROUP}')"
_API_KEY_EXPR: Final = (
f"CASE WHEN api_key = '' THEN '{UNATTRIBUTED_API_KEY}' "
f"WHEN api_key = '{UNATTRIBUTED_API_KEY}' OR api_key LIKE 'raw:%' THEN 'raw:' || api_key ELSE api_key END"
)
@dataclass(frozen=True, slots=True)
@ -86,7 +93,11 @@ def build_where_clause(scope: DailyActivityScope, *, start_index: int = 1) -> tu
if scope.exclude_entity_ids
else ()
),
*((f"model = ${model_index}",) if scope.model else ()),
*(
(f"{_MODEL_EXPR if scope.model == UNKNOWN_MODEL_SPEND_LOG_MODEL else 'model'} = ${model_index}",)
if scope.model
else ()
),
*(
("FALSE",)
if scope.api_keys == ()
@ -143,7 +154,7 @@ def _top_api_keys_sql(pg_table: str, where_clause: str, *, sentinel_param: int,
FROM "{pg_table}"
WHERE {where_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR}
LIMIT ${limit_param}
"""
@ -162,13 +173,13 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq
(SELECT
date,
NULL::text AS api_key,
model,
{_MODEL_EXPR} AS model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
(GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT}
| GROUPING(model, {_MODEL_GROUP_EXPR},
| GROUPING({_MODEL_EXPR}, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
NULL::bigint AS distinct_api_keys,{metric_select}
@ -176,7 +187,7 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date),
(date, model),
(date, {_MODEL_EXPR}),
(date, {_MODEL_GROUP_EXPR}),
(date, custom_llm_provider),
(date, mcp_namespaced_tool_name),
@ -190,12 +201,12 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq
SELECT
date,
api_key,
model,
{_MODEL_EXPR} AS model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
GROUPING(date, api_key, {_MODEL_EXPR}, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select}
@ -203,7 +214,7 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date, api_key),
(date, model, api_key),
(date, {_MODEL_EXPR}, api_key),
(date, {_MODEL_GROUP_EXPR}, api_key),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name, api_key),
@ -313,7 +324,7 @@ def build_key_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) ->
FROM (SELECT 1) AS one
LEFT JOIN LATERAL (
SELECT * FROM ranked
ORDER BY rank_spend DESC, api_key
ORDER BY rank_spend DESC, {_API_KEY_EXPR}
LIMIT ${limit_param} OFFSET ${offset_param}
) AS page ON TRUE
"""
@ -376,7 +387,7 @@ def build_key_search_sql(scope: DailyActivityScope, *, search: str, limit: int)
WHERE {where_clause}
AND api_key <> ${sentinel_param}
AND (
api_key ILIKE ${search_param} ESCAPE '\\'
{_API_KEY_EXPR} ILIKE ${search_param} ESCAPE '\\'
OR api_key IN (
SELECT v.token FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = v.user_id
@ -392,7 +403,7 @@ def build_key_search_sql(scope: DailyActivityScope, *, search: str, limit: int)
)
)
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR}
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, f"%{escaped}%", PTU_SENTINEL_API_KEY, limit))
@ -406,15 +417,13 @@ def build_model_top_keys_sql(
model_param: Final = len(where_params) + 1
sentinel_param: Final = model_param + 1
limit_param: Final = sentinel_param + 1
model_clause: Final = (
f"COALESCE(NULLIF(model_group, ''), model) = ${model_param}" if by_model_group else f"model = ${model_param}"
)
model_clause: Final = f"{_MODEL_GROUP_EXPR if by_model_group else _MODEL_EXPR} = ${model_param}"
sql: Final = f"""
SELECT api_key,{_key_spend_select()}
FROM "{PRISMA_TO_PG_TABLE[scope.table]}"
WHERE {where_clause} AND {model_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR}
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, model_group, PTU_SENTINEL_API_KEY, limit))
@ -431,7 +440,7 @@ def build_cache_leakage_keys_sql(scope: DailyActivityScope, *, limit: int) -> Sq
WHERE {where_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0
ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key
ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, {_API_KEY_EXPR}
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, limit))
@ -483,7 +492,7 @@ def build_export_sql(
{alias_expression} AS entity_alias,
{output_key} AS api_key,
{user_fields},
{"NULLIF(COALESCE(scoped.model, ''), '')" if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model,
{_MODEL_EXPR if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model,
COALESCE(SUM(scoped.spend), 0)::float AS spend,
{flat_cost} AS flat_cost,
COALESCE(SUM(scoped.prompt_tokens), 0)::bigint AS prompt_tokens,
@ -521,25 +530,26 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
)
if export_type is ExportType.DAILY_WITH_KEYS:
return (
"scoped.api_key",
"NULLIF(scoped.api_key, '')",
"MAX(COALESCE(vt.key_alias, dvt.key_alias)) AS key_alias, "
_API_KEY_EXPR,
_API_KEY_EXPR,
"MAX(CASE WHEN scoped.api_key = '' THEN 'Unattributed' "
"ELSE COALESCE(vt.key_alias, dvt.key_alias) END) AS key_alias, "
"MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email",
(
'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key',
"LEFT JOIN \"LiteLLM_VerificationToken\" vt ON vt.token = scoped.api_key AND scoped.api_key <> ''",
"""LEFT JOIN LATERAL (
SELECT key_alias, user_id
FROM "LiteLLM_DeletedVerificationToken"
WHERE token = scoped.api_key
ORDER BY deleted_at DESC
LIMIT 1
) dvt ON vt.token IS NULL""",
) dvt ON vt.token IS NULL AND scoped.api_key <> ''""",
'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)',
),
)
if export_type is ExportType.DAILY_WITH_MODELS:
return (
"COALESCE(scoped.model, '')",
_MODEL_EXPR,
"NULL::text",
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
(),

12
litellm/rerank_api/main.py
View file

@ -19,6 +19,12 @@ from litellm.types.rerank import RerankResponse
from litellm.types.router import *
from litellm.utils import ProviderConfigManager, client, exception_type
_PROVIDERS_RESOLVING_THEIR_OWN_RERANK_BASE: Final = (
litellm.LlmProviders.DASHSCOPE,
litellm.LlmProviders.QWENCLOUD,
litellm.LlmProviders.QWEN_AI_PLATFORM,
)
####### ENVIRONMENT VARIABLES ###################
# Initialize any necessary instances or variables here
together_rerank: Final = TogetherAIRerank()
@ -543,7 +549,11 @@ def rerank(
logging_obj=litellm_logging_obj,
timeout=optional_params.timeout,
api_key=dynamic_api_key or optional_params.api_key,
api_base=dynamic_api_base or optional_params.api_base,
api_base=(
optional_params.api_base or dynamic_api_base
if _custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_RERANK_BASE
else dynamic_api_base or optional_params.api_base
),
_is_async=_is_async,
headers=headers or litellm.headers or {},
client=client,

2
litellm/responses/streaming_iterator.py
View file

@ -227,7 +227,7 @@ async def drain_logging_tasks(timeout: float = 5.0) -> None:
return
_, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout)
if pending:
verbose_logger.warning("Timed out draining %d logging tasks on shutdown, dropping them", len(pending))
verbose_logger.warning("Timed out draining logging tasks on shutdown: %d tasks still alive", len(pending))
_ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(

64
litellm/router.py
View file

@ -288,6 +288,11 @@ from litellm.router_utils.embedding_batch_size_check import (
embedding_deployment_selection,
log_rejection_before_provider_call,
)
from litellm.router_utils.endpoint_eligibility import (
eligible_deployments,
endpoint_selection,
warn_on_unenforced_endpoint_flag,
)
from litellm.router_utils.fallback_event_handlers import (
MID_STREAM_FALLBACK_CONTROLS_KEY,
AttemptedFallbackTargets,
@ -483,6 +488,7 @@ from litellm.utils import (
ProviderConfigManager,
Rules,
TextCompletionResponse,
aresolve_api_key_secret_reference,
function_setup,
get_llm_provider,
get_non_default_completion_params,
@ -493,7 +499,6 @@ from litellm.utils import (
lookup_model_info_for_registration,
model_info_with_pricing_overrides,
provider_rejectable_params,
resolve_api_key_secret_reference,
set_live_deployment_replay,
token_counter,
)
@ -1473,7 +1478,7 @@ class Router:
self.assistants_config = assistants_config
self.search_tools = search_tools or []
self.guardrail_list = guardrail_list or []
self.deployment_names: list = [] # names of models under litellm_params. ex. azure/chatgpt-v-2
self.deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2
self.deployment_latency_map = {}
### CACHING ###
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local" # default to an in-memory cache
@ -6089,12 +6094,13 @@ class Router:
kwargs.setdefault("metadata", {}).update({"model_group": model})
# pick the one that is available (lowest TPM/RPM)
deployment: Final = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with endpoint_selection("text_completion", model=model):
deployment: Final = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
data: Final = deployment["litellm_params"].copy()
for k, v in self.default_litellm_params.items():
@ -6950,6 +6956,7 @@ class Router:
"""
model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs)
handler: Final = retry_kwargs.get("original_generic_function")
healthy_deployments, all_deployments = await self._async_get_healthy_deployments(
model=model_group,
parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs),
@ -6980,7 +6987,10 @@ class Router:
)
carry_superseded_attempt_usage(retry_kwargs)
try:
response = await attempt_function(**retry_kwargs)
with endpoint_selection(
getattr(handler, "__name__", None), getattr(handler, "__module__", None), model=model_group
):
response = await attempt_function(**retry_kwargs)
except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry
wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider)
if wrapped is None:
@ -9350,9 +9360,18 @@ class Router:
Handler for making a call to the .completion()/.embeddings()/etc. functions.
"""
model_group: Final = kwargs.get("model")
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
generic_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("original_generic_function"))
endpoint_function: Final = (
generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function)
)
with endpoint_selection(
getattr(endpoint_function, "__name__", None),
getattr(endpoint_function, "__module__", None),
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
):
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
@ -10673,6 +10692,7 @@ class Router:
declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params)
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params)
warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params)
warn_on_unenforced_endpoint_flag(model_name=_model_name, model_info=_model_info)
deployment = without_model_capability_flags(
Deployment(
**deployment_info,
@ -11476,6 +11496,9 @@ class Router:
declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True)
warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params)
warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params)
warn_on_unenforced_endpoint_flag(
model_name=deployment.model_name, model_info=deployment.model_info.model_dump()
)
served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment))
# add to model list
@ -12091,7 +12114,7 @@ class Router:
return
if api_base is None or "*" in model or params.get("use_clientside_credentials"):
return
api_key: Final = resolve_api_key_secret_reference(
api_key: Final = await aresolve_api_key_secret_reference(
params.api_key or dynamic_api_key,
model=model,
llm_provider=provider,
@ -14832,9 +14855,15 @@ class Router:
"""
return [m for m in self.model_list if m["litellm_params"]["model"] == model]
def get_routable_deployments(self, model: str) -> tuple[DeploymentTypedDict, ...]:
def get_routable_deployments(
self,
model: str,
request_kwargs: Mapping[str, object] = MappingProxyType({}),
call_type: str | None = None,
) -> tuple[DeploymentTypedDict, ...]:
try:
_, candidates = self._common_checks_available_deployment(model=model, request_kwargs={}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped
with endpoint_selection(call_type, model=model):
_, candidates = self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped
except (litellm.BadRequestError, ValueError):
return ()
if isinstance(candidates, dict):
@ -15110,7 +15139,10 @@ class Router:
model=model,
llm_provider="",
)
return selectable
eligible: Final = eligible_deployments(model, selectable)
if isinstance(eligible, tuple):
return list(eligible)
raise eligible
def _filter_reserved_deployments(
self,

120
litellm/router_utils/endpoint_eligibility.py Normal file
View file

@ -0,0 +1,120 @@
from collections.abc import Generator, Mapping, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from types import MappingProxyType
from typing import Final, TypeVar
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator
from litellm._logging import verbose_router_logger
from litellm.exceptions import EndpointNotSupportedByModelError
CHAT_ENDPOINT_BY_CALL: Final[Mapping[str, str]] = MappingProxyType(
{
"completion": "/v1/chat/completions",
"acompletion": "/v1/chat/completions",
"_completion": "/v1/chat/completions",
"_acompletion": "/v1/chat/completions",
"text_completion": "/v1/completions",
"atext_completion": "/v1/completions",
"_atext_completion": "/v1/completions",
"anthropic_messages": "/v1/messages",
"aresponses": "/v1/responses",
"agenerate_content": "generateContent",
"agenerate_content_stream": "streamGenerateContent",
"acreate_interaction": "/v1beta/interactions",
"litellm.interactions.main.acreate": "/v1beta/interactions",
"litellm.interactions.main.create": "/v1beta/interactions",
"_aresponses_websocket": "/v1/responses",
"acompact_responses": "/v1/responses/compact",
}
)
def _endpoint_path(endpoint: str) -> str:
rooted: Final = "/" + endpoint.strip().strip("/").casefold()
path: Final = rooted.removeprefix("/openai").removeprefix("/v1beta").removeprefix("/v1")
return "/" + path.rsplit(":", 1)[-1] if ":" in path else path
CHAT_ENDPOINTS: Final = frozenset(_endpoint_path(endpoint) for endpoint in CHAT_ENDPOINT_BY_CALL.values())
_Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object])
_ENFORCE_FLAG: Final = TypeAdapter(bool | None)
_REQUEST_SELECTION: Final[ContextVar[tuple[str | None, str | None]]] = ContextVar(
"router_request_endpoint", default=(None, None)
)
class _EndpointDeclaration(BaseModel):
model_config = ConfigDict(frozen=True)
enforce_supported_endpoints: bool | None = None
supported_endpoints: tuple[str, ...] | None = None
@field_validator("supported_endpoints", mode="before")
@classmethod
def _one_endpoint_as_a_list(cls, value: object) -> object:
return (value,) if isinstance(value, str) else value
@contextmanager
def endpoint_selection(call_type: object, module: object = None, model: object = None) -> Generator[None]:
qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type
endpoint: Final = CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None
token: Final = _REQUEST_SELECTION.set((endpoint, model if isinstance(model, str) else None))
try:
yield
finally:
_REQUEST_SELECTION.reset(token)
def _enforced_endpoints(model_info: object) -> tuple[str, ...] | None:
try:
declaration: Final = _EndpointDeclaration.model_validate(model_info)
except ValidationError:
return None
return declaration.supported_endpoints if declaration.enforce_supported_endpoints else None
def warn_on_unenforced_endpoint_flag(model_name: str, model_info: Mapping[str, object]) -> str | None:
try:
flagged: Final = _ENFORCE_FLAG.validate_python(model_info.get("enforce_supported_endpoints"))
except ValidationError:
return None
if not flagged or _enforced_endpoints(model_info) is not None:
return None
warning: Final = (
f"Deployment '{model_name}' sets enforce_supported_endpoints without a supported_endpoints list of "
'strings, so litellm enforces nothing for it. Set supported_endpoints, e.g. ["/v1/embeddings"].'
)
verbose_router_logger.warning(warning)
return warning
def non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None:
declared: Final = _enforced_endpoints(deployment.get("model_info"))
if declared is None:
return None
endpoints: Final = frozenset(declared)
normalized: Final = frozenset(_endpoint_path(endpoint) for endpoint in endpoints)
return endpoints if normalized.isdisjoint(CHAT_ENDPOINTS) else None
def eligible_deployments(
model: str, deployments: Sequence[_Deployment]
) -> tuple[_Deployment, ...] | EndpointNotSupportedByModelError:
endpoint, requested_model = _REQUEST_SELECTION.get()
if endpoint is None or not deployments:
return tuple(deployments)
eligible: Final = tuple(deployment for deployment in deployments if non_chat_endpoints(deployment) is None)
if eligible:
return eligible
requested: Final = requested_model or model
supported: Final = ", ".join(
sorted(frozenset[str]().union(*(non_chat_endpoints(deployment) or frozenset() for deployment in deployments)))
)
return EndpointNotSupportedByModelError(
message=f"{endpoint} is not supported by model={requested}. Supported endpoints: {supported or 'none'}",
model=model,
llm_provider="",
)

135
litellm/secret_managers/main.py
View file

@ -5,7 +5,7 @@ import time
import traceback
from collections.abc import Callable
from dataclasses import dataclass
from typing import Final
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
@ -13,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.env_utils import get_env_float
from litellm.litellm_core_utils.single_flight import single_flight_lru_cache
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.secret_managers.get_azure_ad_token_provider import (
@ -167,12 +168,37 @@ def get_secret_str(
LOCAL_OIDC_SECRET_PREFIXES: Final = ("oidc/file/", "oidc/env/")
GOOGLE_OIDC_SECRET_PREFIX: Final = "oidc/google/"
API_KEY_OIDC_SECRET_PREFIXES: Final = (*LOCAL_OIDC_SECRET_PREFIXES, GOOGLE_OIDC_SECRET_PREFIX)
_GOOGLE_METADATA_DEFAULT_HOST: Final = "metadata.google.internal"
_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = 3.0
_GOOGLE_TOKEN_UNAVAILABLE: Final = "the metadata server did not return an ID token"
def _google_metadata_identity_url() -> str:
host: Final = os.getenv("GCE_METADATA_HOST") or _GOOGLE_METADATA_DEFAULT_HOST
return f"http://{host}/computeMetadata/v1/instance/service-accounts/default/identity"
def _google_metadata_timeout() -> float:
seconds: Final = get_env_float("GCE_METADATA_TIMEOUT", _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS)
return seconds if seconds > 0 else _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS
class GoogleMetadataStatusError(ValueError):
def __init__(self, status_code: int) -> None:
super().__init__(f"Google OIDC provider failed with HTTP {status_code}")
self.status_code: Final = status_code
@dataclass(frozen=True, slots=True)
class UnresolvableSecretReference:
kind: str
reason: str
retryable: bool = False
OidcResolver: TypeAlias = Callable[[str], "str | UnresolvableSecretReference"]
_OIDC_PROVIDERS: Final = frozenset({"google", "circleci", "circleci_v2", "github", "azure", "file", "env", "env_path"})
@ -198,11 +224,86 @@ def resolve_local_oidc_reference(value: str) -> str | UnresolvableSecretReferenc
return secret
def _logged_unresolvable(reference: str, kind: str, reason: str) -> UnresolvableSecretReference:
def resolve_api_key_oidc_reference(value: str) -> str | UnresolvableSecretReference:
if not value.startswith(GOOGLE_OIDC_SECRET_PREFIX):
return resolve_local_oidc_reference(value)
try:
token: Final = _google_id_token(
value, value.removeprefix(GOOGLE_OIDC_SECRET_PREFIX), timeout=_google_metadata_timeout()
)
except (httpx.HTTPError, httpx.InvalidURL, ValueError) as exc:
return _logged_unresolvable(
value,
"oidc/google",
_GOOGLE_TOKEN_UNAVAILABLE,
f"{type(exc).__name__}: {exc}",
retryable=_metadata_server_failed_to_answer(exc),
)
if not token:
return _logged_unresolvable(value, "oidc/google", "it resolved to an empty value")
return token
def _metadata_server_failed_to_answer(exc: httpx.HTTPError | httpx.InvalidURL | ValueError) -> bool:
match exc:
case httpx.TransportError():
return True
case GoogleMetadataStatusError(status_code=status_code):
return status_code >= 500
case _:
return False
def _google_id_token(reference: str, audience: str, timeout: float | None) -> str:
cached: Final = cached_oidc_token(reference)
if cached is not None:
return cached
oidc_client: Final = _get_oidc_http_handler()
try:
# https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature
response: Final = oidc_client.get(
_google_metadata_identity_url(),
params={"audience": audience},
headers={"Metadata-Flavor": "Google"},
timeout=timeout,
)
finally:
oidc_client.close()
if response.status_code != 200:
raise GoogleMetadataStatusError(response.status_code)
token: Final = response.text
ttl: Final = _oidc_token_cache_ttl(token, 3600 - 60)
if not token:
verbose_logger.warning("Google OIDC token for %s is empty; not caching it", _single_line(reference))
elif ttl > 0:
oidc_cache.set_cache(key=reference, value=token, ttl=ttl)
else:
verbose_logger.warning(
"Google OIDC token for %s is already expired or expires within %ss; not caching it",
_single_line(reference),
_OIDC_TOKEN_EXPIRY_MARGIN_SECONDS,
)
return token
def cached_oidc_token(reference: str) -> str | None:
cached: Final = oidc_cache.get_cache(key=reference) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache.get_cache is untyped
return cached if isinstance(cached, str) and cached else None
def _logged_unresolvable(
reference: str, kind: str, reason: str, detail: str | None = None, retryable: bool = False
) -> UnresolvableSecretReference:
verbose_logger.warning(
"Secret reference %s could not be resolved: %s", reference.replace("\r", "").replace("\n", ""), reason
"Secret reference %s could not be resolved: %s",
_single_line(reference),
reason if detail is None else f"{reason} ({_single_line(detail)})",
)
return UnresolvableSecretReference(kind=kind, reason=reason)
return UnresolvableSecretReference(kind=kind, reason=reason, retryable=retryable)
def _single_line(text: str) -> str:
return text.replace("\r", "").replace("\n", "")
def _unresolvable_reason(kind: str, exc: OSError | ValueError) -> str:
@ -273,31 +374,7 @@ def get_secret(
oidc_aud = "/".join(secret_name_split.split("/")[1:])
# TODO: Add caching for HTTP requests
if oidc_provider == "google":
oidc_token = oidc_cache.get_cache(key=secret_name)
if oidc_token is not None:
return oidc_token
oidc_client = _get_oidc_http_handler()
# https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature
response = oidc_client.get(
"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity",
params={"audience": oidc_aud},
headers={"Metadata-Flavor": "Google"},
)
if response.status_code == 200:
oidc_token = response.text
ttl: Final = _oidc_token_cache_ttl(oidc_token, 3600 - 60)
if ttl > 0:
oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=ttl)
else:
verbose_logger.warning(
"Google OIDC token for %s is already expired or expires within %ss; not caching it",
secret_name,
_OIDC_TOKEN_EXPIRY_MARGIN_SECONDS,
)
return oidc_token
else:
raise ValueError("Google OIDC provider failed")
return _google_id_token(secret_name, oidc_aud, timeout=None)
elif oidc_provider == "circleci":
# https://circleci.com/docs/openid-connect-tokens/
env_secret = os.getenv("CIRCLE_OIDC_TOKEN")

1
litellm/types/router.py
View file

@ -270,6 +270,7 @@ def _whole_float_as_int(value: object) -> object:
class ModelInfo(MirroredPricingParams):
enforce_supported_endpoints: bool | None = None
id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance
db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config.
updated_at: datetime.datetime | None = None

134
litellm/utils.py
View file

@ -28,6 +28,7 @@ import textwrap
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
from dataclasses import dataclass, field, replace
from functools import lru_cache, partial, wraps
@ -488,11 +489,14 @@ from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig
from litellm.secret_managers.main import (
LOCAL_OIDC_SECRET_PREFIXES,
API_KEY_OIDC_SECRET_PREFIXES,
GOOGLE_OIDC_SECRET_PREFIX,
OidcResolver,
UnresolvableSecretReference,
cached_oidc_token,
get_secret,
oidc_reference_kind,
resolve_local_oidc_reference,
resolve_api_key_oidc_reference,
)
from ._logging import is_debugging_on, verbose_logger
@ -795,27 +799,97 @@ def _named_credential_values(kwargs: Mapping[str, object]) -> Mapping[str, objec
)
def resolve_api_key_secret_reference(
async def aresolve_api_key_secret_reference(
api_key: str | None, model: object, llm_provider: object, upstream: Mapping[str, object]
) -> str | None:
if not isinstance(api_key, AdminApiKeySecretReference):
return api_key
return _secret_behind_api_key_reference(api_key, model=model, llm_provider=llm_provider, upstream=upstream)
return await _resolved_off_the_event_loop(
partial(_secret_behind_api_key_reference, api_key, model, llm_provider, upstream), api_key
)
_OIDC_TOKEN_FETCH_EXECUTOR: Final = ThreadPoolExecutor(thread_name_prefix="litellm-oidc-token-fetch")
_PASS_THROUGH_CALL_TYPES: Final = frozenset(
{CallTypes.allm_passthrough_route.value, CallTypes.llm_passthrough_route.value}
)
_Resolved: Final = TypeVar("_Resolved")
async def _resolved_off_the_event_loop(resolve: Callable[[OidcResolver], _Resolved], api_key: object) -> _Resolved:
if not isinstance(api_key, AdminApiKeySecretReference) or not api_key.startswith(GOOGLE_OIDC_SECRET_PREFIX):
return resolve(resolve_api_key_oidc_reference)
match cached_oidc_token(api_key):
case str() as token:
return resolve(lambda _reference: token)
case None:
_raise_if_the_reference_may_not_be_used(resolve)
fetched: Final = await asyncio.shield(
_google_token_fetch(asyncio.get_running_loop(), api_key).start_or_join()
)
return resolve(lambda _reference: fetched)
def _raise_if_the_reference_may_not_be_used(resolve: Callable[[OidcResolver], object]) -> None:
_ = resolve(lambda _reference: "")
@dataclass(slots=True)
class _GoogleTokenFetch:
reference: str
in_flight: asyncio.Future[str | UnresolvableSecretReference] | None = None
def start_or_join(self) -> asyncio.Future[str | UnresolvableSecretReference]:
if self.in_flight is not None:
return self.in_flight
caller_context: Final = contextvars.copy_context()
fetch: Final = asyncio.get_running_loop().run_in_executor(
_OIDC_TOKEN_FETCH_EXECUTOR, lambda: caller_context.run(resolve_api_key_oidc_reference, self.reference)
)
fetch.add_done_callback(self._land)
self.in_flight = fetch
return fetch
def _land(self, _fetch: asyncio.Future[str | UnresolvableSecretReference]) -> None:
self.in_flight = None
@lru_cache(maxsize=256)
def _google_token_fetch(loop: asyncio.AbstractEventLoop, reference: str) -> _GoogleTokenFetch:
return _GoogleTokenFetch(reference)
def _secret_behind_api_key_reference(
reference: AdminApiKeySecretReference, model: object, llm_provider: object, upstream: Mapping[str, object]
reference: AdminApiKeySecretReference,
model: object,
llm_provider: object,
upstream: Mapping[str, object],
resolve_oidc: OidcResolver = resolve_api_key_oidc_reference,
call_type: str | None = None,
) -> str:
if not reference.may_reach(upstream):
raise _secret_reference_error(_unreachable_upstream_message(reference), llm_provider=llm_provider, model=model)
if not reference.startswith(LOCAL_OIDC_SECRET_PREFIXES):
if not reference.startswith(API_KEY_OIDC_SECRET_PREFIXES):
raise _secret_reference_error(
f"api_key is an {oidc_reference_kind(reference)} secret reference, a kind api_key does not support. "
"Use oidc/file/<absolute path> or oidc/env/<VARIABLE_NAME>.",
"Use oidc/file/<absolute path>, oidc/env/<VARIABLE_NAME> or oidc/google/<audience>.",
llm_provider=llm_provider,
model=model,
)
match resolve_local_oidc_reference(reference):
if call_type in _PASS_THROUGH_CALL_TYPES and reference.startswith(GOOGLE_OIDC_SECRET_PREFIX):
raise _secret_reference_error(
"api_key oidc/google secret reference could not be resolved: pass-through routes do not fetch Google ID "
"tokens. Use a static key, oidc/file/<absolute path> or oidc/env/<VARIABLE_NAME>.",
llm_provider=llm_provider,
model=model,
)
match resolve_oidc(reference):
case UnresolvableSecretReference(kind=kind, reason=reason, retryable=True):
raise ServiceUnavailableError(
message=f"api_key {kind} secret reference could not be resolved: {reason}",
llm_provider=llm_provider,
model=model,
)
case UnresolvableSecretReference(kind=kind, reason=reason):
raise _secret_reference_error(
f"api_key {kind} secret reference could not be resolved: {reason}",
@ -847,21 +921,39 @@ def with_named_credential(kwargs: Mapping[str, object]) -> Mapping[str, object]:
return {**kwargs, **_named_credential_values(kwargs)}
def load_credentials_from_list(kwargs: dict[str, object]) -> None:
def load_credentials_from_list(kwargs: dict[str, object], call_type: str | None = None) -> None:
"""
Updates kwargs with the credentials if credential_name in kwarg, then swaps an admin-configured oidc/ api_key
reference for the secret it points to, read now so a rotated token is used on the next call
"""
kwargs.update(_credential_updates(kwargs))
kwargs.update(_credential_updates(kwargs, _named_credential_values(kwargs), call_type=call_type))
def _credential_updates(kwargs: Mapping[str, object]) -> Mapping[str, object]:
async def acredential_updates(kwargs: Mapping[str, object], call_type: str | None = None) -> Mapping[str, object]:
credential_values: Final = _named_credential_values(kwargs)
api_key: Final = kwargs.get("api_key", credential_values.get("api_key"))
return await _resolved_off_the_event_loop(
partial(_credential_updates, kwargs, credential_values, call_type=call_type),
_api_key_after_credentials(kwargs, credential_values),
)
def _api_key_after_credentials(kwargs: Mapping[str, object], credential_values: Mapping[str, object]) -> object:
return kwargs.get("api_key", credential_values.get("api_key"))
def _credential_updates(
kwargs: Mapping[str, object],
credential_values: Mapping[str, object],
resolve_oidc: OidcResolver = resolve_api_key_oidc_reference,
call_type: str | None = None,
) -> Mapping[str, object]:
api_key: Final = _api_key_after_credentials(kwargs, credential_values)
if not isinstance(api_key, AdminApiKeySecretReference):
return credential_values
secret: Final = _secret_behind_api_key_reference(
api_key,
resolve_oidc=resolve_oidc,
call_type=call_type,
model=kwargs.get("model"),
llm_provider=kwargs.get("custom_llm_provider"),
upstream=MappingProxyType(
@ -1664,15 +1756,19 @@ def _schedule_async_success_logging(
return
def _enqueue_async_logging() -> None:
from litellm.responses.streaming_iterator import spawn_logging_task
logging_obj.mark_async_success_log_enqueued()
asyncio.create_task(
spawn_logging_task(
asyncio.get_running_loop(),
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
),
task_name="success callback enqueue",
)
if not logging_obj.defer_async_logging:
@ -2060,7 +2156,7 @@ def client(original_function):
assert logging_obj is not None, "logging_obj should not be None after function_setup"
## LOAD CREDENTIALS
load_credentials_from_list(kwargs)
load_credentials_from_list(kwargs, call_type)
from litellm.llms.github_copilot.per_user_auth import attach_github_copilot_user_session
attach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
@ -2369,7 +2465,7 @@ def client(original_function):
kwargs["litellm_logging_obj"] = logging_obj
## LOAD CREDENTIALS
load_credentials_from_list(kwargs)
kwargs.update(await acredential_updates(kwargs, call_type)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is the untyped request dict
from litellm.llms.github_copilot.per_user_auth import aattach_github_copilot_user_session
await aattach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
@ -10124,6 +10220,12 @@ class ProviderConfigManager:
)
return TencentAnthropicMessagesConfig()
elif litellm.LlmProviders.DASHSCOPE == provider:
from litellm.llms.dashscope.messages.transformation import (
DashScopeAnthropicMessagesConfig,
)
return DashScopeAnthropicMessagesConfig()
elif litellm.LlmProviders.GITHUB_COPILOT == provider:
if "claude" in model_lower:
from litellm.llms.github_copilot.messages.transformation import (

1
tests/e2e/coverage_registry/llm_nonconversational.yaml
View file

@ -2,6 +2,7 @@
- {id: llm.completions.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: completions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_completions_endpoint_e2e.py", rationale: "Legacy text /completions endpoint, second-highest production request volume"}
- {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"}
- {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"}
- {id: llm.embeddings.openai.supported_endpoints_gate.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], fail_before_fix: proven, source: "test_embedding_only_supported_endpoints_e2e.py", rationale: "Explicit endpoint restrictions preserve real embeddings, refuse chat with a gateway 400, and route chat through a live fallback; zero-provider-call contracts live in integration/routing/test_non_chat_model_endpoint_gate.py"}
- {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"}
- {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (GitHub issue #36646)"}
- {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"}

119
tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py Normal file
View file

@ -0,0 +1,119 @@
"""Live provider checks for embeddings, the chat refusal, and chat fallback from an explicitly restricted model."""
from __future__ import annotations
import dataclasses
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, ModelInfoBody, ModelMode, ModelNewBody
from openai import BadRequestError, OpenAI
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients
pytestmark = pytest.mark.e2e
EMBEDDING_BACKEND: Final = "openai/text-embedding-3-small"
CHAT_FALLBACK_BACKEND: Final = "openai/gpt-5.4-nano"
EMBEDDINGS_ONLY: Final = ("/v1/embeddings",)
PROMPT: Final = "Say this is a test!"
def _subject(route: Route) -> Subject:
return Subject(
domain=Domain.LLM_TRANSLATION,
route=route,
providers=(Provider.OPENAI,),
models=(EMBEDDING_BACKEND,),
mode=Mode.NONSTREAM,
)
def _register(
proxy: ProxyClient,
resources: ResourceManager,
prefix: str,
supported_endpoints: tuple[str, ...] | None,
backend: str = EMBEDDING_BACKEND,
mode: ModelMode = "embedding",
) -> str:
model: Final = f"{prefix}-{unique_marker()}"
model_id: Final = proxy.register_model(
ModelNewBody(
model_name=model,
litellm_params=LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY"),
model_info=ModelInfoBody(
enforce_supported_endpoints=supported_endpoints is not None,
mode=mode,
supported_endpoints=None if supported_endpoints is None else list(supported_endpoints),
),
)
)
resources.defer(lambda: proxy.delete_model(model_id))
return model
def _openai_at(sdk: SdkClients, api_prefix: str, key: str) -> OpenAI:
return dataclasses.replace(sdk, base_url=f"{sdk.base_url}{api_prefix}").openai(key)
class TestEmbeddingOnlySupportedEndpoints:
@pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works")
@meta(_subject(Route.EMBEDDINGS))
@pytest.mark.parametrize(
("prefix", "supported_endpoints"),
(
pytest.param("e2e-embed-only", EMBEDDINGS_ONLY, id="embed-only"),
pytest.param("e2e-embed-control", None, id="control"),
),
)
def test_embeddings_return_a_vector(
self,
proxy: ProxyClient,
resources: ResourceManager,
sdk: SdkClients,
prefix: str,
supported_endpoints: tuple[str, ...] | None,
) -> None:
model: Final = _register(proxy, resources, prefix, supported_endpoints)
embeddings: Final = sdk.openai(resources.key()).embeddings.create(
model=model, input=PROMPT, extra_body=NO_PROXY_CACHE
)
assert embeddings.data, f"/embeddings returned no data for {model}: {embeddings!r}"
assert any(component != 0.0 for component in embeddings.data[0].embedding), (
f"/embeddings returned an empty or all-zero vector for {model}"
)
@pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works")
@meta(dataclasses.replace(_subject(Route.CHAT_COMPLETIONS), models=(EMBEDDING_BACKEND, CHAT_FALLBACK_BACKEND)))
def test_chat_completions_with_a_fallback_are_answered_by_the_fallback(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY)
fallback: Final = _register(proxy, resources, "e2e-chat-fallback", None, CHAT_FALLBACK_BACKEND, "chat")
completion: Final = _openai_at(sdk, "/v1", resources.key()).chat.completions.create(
model=model,
messages=[{"role": "user", "content": PROMPT}],
extra_body={**NO_PROXY_CACHE, "fallbacks": [fallback]},
)
assert completion.choices and completion.choices[0].message.content, (
f"the fallback {fallback} should answer a chat call to {model}, got {completion!r}"
)
@pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works")
@meta(_subject(Route.CHAT_COMPLETIONS))
def test_chat_completions_without_a_fallback_are_refused_by_the_gateway(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY)
with pytest.raises(BadRequestError) as refused:
_ = _openai_at(sdk, "/v1", resources.key()).chat.completions.create(
model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE
)
assert (
f"/v1/chat/completions is not supported by model={model}. Supported endpoints: /v1/embeddings"
in refused.value.message
), f"chat on {model} should be refused by the gateway before any provider call, got {refused.value!r}"

3
tests/e2e/models.py
View file

@ -1416,7 +1416,7 @@ class LiteLLMParamsBody(BaseModel):
order: int | None = None
ModelMode = Literal["chat", "batch", "realtime", "image_generation"]
ModelMode = Literal["chat", "embedding", "batch", "realtime", "image_generation"]
class TokenCountModelBody(BaseModel, frozen=True):
@ -1425,6 +1425,7 @@ class TokenCountModelBody(BaseModel, frozen=True):
class ModelInfoBody(BaseModel):
enforce_supported_endpoints: bool | None = None
# id is left unset so the proxy assigns a unique model_id per deployment.
# Pinning it to the model_name made re-registrations of a fixed-name model
# (e.g. the batch suite's openai-batch) collide on the model_id unique

248
tests/integration/database/migrations/test_request_log_indexes.py
View file

@ -16,7 +16,7 @@ from typing import Final
import psycopg
import pytest
from litellm_proxy_extras import migration_data_checks, request_log_indexes
from litellm_proxy_extras import migration_data_checks, migration_dry_run, request_log_indexes
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock
from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration
from litellm_proxy_extras.request_log_indexes import (
@ -967,33 +967,25 @@ _HASH_LEGACY_KEYS: Final = (
@requires_db
@RESOLVERS
@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_keys_are_too_long_for_the_index(
def test_the_migration_job_applies_every_migration_and_only_the_index_build_waits_for_too_long_keys(
scratch_database: str, use_v2_resolver: bool
) -> None:
"""Rows never stop an upgrade. The migrations go through, and the one request-log index
whose keys do not fit is left for the next build with the rows named in the log."""
with psycopg.connect(scratch_database, autocommit=True) as conn:
_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:
ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
refusal: Final = str(refused.value)
assert (
'"LiteLLM_SpendLogs".api_key, startTime: 2 row(s) with an index entry over 2704 bytes, largest 3232 bytes '
"(api_key up to 3207 bytes)"
) in refusal, refusal
assert f'blocks: CREATE INDEX "{API_KEY_INDEX}"' in refusal, refusal
assert _ledger(scratch_database) == ledger_before, "a refused job must not apply a single migration"
assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is False
assert len(_ledger(scratch_database)) > len(ledger_before), "the rows must not hold back a single migration"
assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute(_HASH_LEGACY_KEYS)
assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
assert len(_ledger(scratch_database)) > len(ledger_before)
@requires_db
@ -1017,49 +1009,6 @@ def test_a_serving_proxy_index_build_that_meets_too_long_keys_reports_the_rows_i
assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
_PROBE_STATEMENTS: Final = (
'CREATE UNIQUE INDEX "probe_b_key" ON "preflight_probe" ("b")',
'ALTER TABLE "preflight_probe" ALTER COLUMN "b" SET NOT NULL',
'ALTER TABLE "preflight_probe" ADD COLUMN "c" INTEGER NOT NULL',
'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_r_fkey" FOREIGN KEY ("r") REFERENCES "preflight_parent"("id")',
'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_a_check" CHECK ("a" <> \'bad\')',
'CREATE INDEX "probe_a_idx" ON "preflight_probe" ("a")',
)
@requires_db
def test_each_check_kind_counts_exactly_the_rows_postgres_would_reject(scratch_database: str) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "preflight_parent" ("id" INTEGER PRIMARY KEY)')
conn.execute('CREATE TABLE "preflight_probe" ("a" TEXT, "b" INTEGER, "r" INTEGER)')
conn.execute('INSERT INTO "preflight_parent" VALUES (1)')
conn.execute(
'INSERT INTO "preflight_probe" VALUES (%s, 1, 1), (%s, 1, 1), (%s, NULL, 2), (%s, 2, NULL), (%s, 3, 1)',
("ok", "ok", "bad", "bad", secrets.token_hex(1400)),
)
checks: Final = tuple(
check
for statement in _PROBE_STATEMENTS
for check in migration_data_checks.checks_for_statement(statement, "probe")
)
findings: Final = migration_data_checks.find_blockers(conn, "public", checks)
for statement in _PROBE_STATEMENTS:
with conn.transaction(force_rollback=True), pytest.raises(psycopg.Error):
conn.execute(statement) # pyright: ignore[reportArgumentType] # test SQL
counted: Final = {
(f.check.kind, f.check.columns): f.rows for f in findings if isinstance(f, migration_data_checks.BlockedChange)
}
assert counted == {
("duplicates", ("b",)): 1,
("nulls", ("b",)): 1,
("any_row", ("c",)): 5,
("orphans", ("r",)): 1,
("check_violations", ()): 2,
("too_long", ("a",)): 1,
}
@requires_db
@pytest.mark.parametrize(
("index_columns", "s_value", "t_value", "longest_fitting_key", "described_as"),
@ -1098,7 +1047,7 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
),
)
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
conn, "public", migration_data_checks.checks_for_index("entry_probe", "entry_probe_idx", index_columns)
)
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"):
conn.execute(create_index)
@ -1122,7 +1071,7 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random
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")
conn, "public", migration_data_checks.checks_for_index("compress_probe", "compress_probe_idx", ("k",))
)
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
conn.execute(create_index)
@ -1154,7 +1103,6 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
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),))
@ -1162,7 +1110,7 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
'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")
conn, "public", migration_data_checks.checks_for_index("stored_probe", "stored_probe_idx", ("k",))
)
rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
@ -1176,14 +1124,176 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
@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:
conn.execute('CREATE TABLE "preflight_slow" ("a" INTEGER)')
conn.execute('INSERT INTO "preflight_slow" SELECT generate_series(1, 50)')
(check,) = migration_data_checks.checks_for_statement(
'ALTER TABLE "preflight_slow" ADD CONSTRAINT "slow" CHECK (pg_sleep(0.1) IS NULL)', "probe"
conn.execute('CREATE TABLE "preflight_slow" ("k" TEXT)')
conn.execute(
"INSERT INTO \"preflight_slow\" SELECT md5(i::text) || repeat('a', 1500) FROM generate_series(1, 200000) AS i"
)
findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.5)
(check,) = migration_data_checks.checks_for_index("preflight_slow", "preflight_slow_idx", ("k",))
findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.01)
timeout_afterwards: Final = conn.execute("SHOW statement_timeout").fetchone()
assert findings == (migration_data_checks.CouldNotVerify(check, 0.5),)
assert timeout_afterwards == ("0",), "the preflight must hand the session back with the timeout it found"
assert "could not be checked within 0.5 s" in migration_data_checks.report(findings)
assert findings == (migration_data_checks.CouldNotVerify(check, 0.01),)
assert timeout_afterwards == ("0",), "the check must hand the session back with the timeout it found"
assert "could not be checked within 0.01 s" in migration_data_checks.report(findings)
CURRENT_RELEASE: Final = "99999999999999"
REQUIRE_TEAM_ALIAS: Final = "20991231000000_require_team_alias"
def _extras_with_scratch_migration(tmp_path: Path, name: str, migration_sql: str) -> Path:
"""A copy of the shipped package dir with one more migration folder, the shape `LITELLM_MIGRATION_DIR` points at."""
extras: Final = tmp_path / "extras"
shutil.copytree(PACKAGE / "migrations", extras / "migrations")
shutil.copy(PACKAGE / "schema.prisma", extras / "schema.prisma")
(extras / "migrations" / name).mkdir()
(extras / "migrations" / name / "migration.sql").write_text(migration_sql)
return extras
def _team_alias_is_nullable(database_url: str) -> bool:
with psycopg.connect(database_url) as conn:
row: Final = conn.execute(
"SELECT is_nullable FROM information_schema.columns "
"WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'team_alias'"
).fetchone()
return row == ("YES",)
@requires_db
@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
def test_the_dry_run_applies_every_pending_migration_in_order_and_rolls_all_of_it_back(scratch_database: str) -> None:
ledger_before: Final = _ledger(scratch_database)
with psycopg.connect(scratch_database, autocommit=True) as conn:
pending: Final = migration_data_checks.pending_migrations(conn, "public", PACKAGE / "migrations")
result: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations")
assert result.rejected is None, result
assert tuple(applied.name for applied in result.applied) == tuple(name for name, _ in pending)
assert len(result.applied) > 1
assert _ledger(scratch_database) == ledger_before, "a dry run must leave the ledger alone"
assert _index_validity(scratch_database, "api_key_startTime_idx") == {}, "a dry run must leave no index behind"
names: Final = tuple(name for name, _ in pending)
assert all(set(applied.replayed) <= set(names[: names.index(applied.name)]) for applied in result.applied), (
"a migration is only ever replayed behind a later one that depends on it"
)
assert "applied one at a time in a dry run" in migration_dry_run.report(result)
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_the_dry_run_reports_postgres_own_error_for_the_row_a_migration_cannot_take(
scratch_database: str, tmp_path: Path
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n'
)
ledger_before: Final = _ledger(scratch_database)
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",))
result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations")
assert result.rejected is not None, result
assert result.rejected.name == REQUIRE_TEAM_ALIAS
assert 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values' in result.rejected.error
assert result.applied == ()
assert _team_alias_is_nullable(scratch_database), "a rejected dry run must leave the schema as it was"
assert _ledger(scratch_database) == ledger_before
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_the_cli_summary_turns_from_postgres_error_to_timings_once_the_rows_are_fixed(
scratch_database: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n'
)
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(extras))
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",))
ok, summary = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True)
assert ok is False
assert summary.startswith(f"Migration {REQUIRE_TEAM_ALIAS} fails on this database after "), summary
assert "contains null values" in summary
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('UPDATE "LiteLLM_TeamTable" SET "team_alias" = \'filled\' WHERE "team_alias" IS NULL')
ok_after, summary_after = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True)
assert ok_after is True
assert summary_after.startswith("1 pending migration(s) applied one at a time in a dry run"), summary_after
assert f"\n {REQUIRE_TEAM_ALIAS}: " in summary_after
assert _team_alias_is_nullable(scratch_database)
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_a_migration_that_depends_on_the_one_before_it_is_replayed_with_it_while_independent_ones_run_alone(
scratch_database: str, tmp_path: Path
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, "20991231000000_add_team_region", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "region" TEXT;\n'
)
(extras / "migrations" / "20991231000001_index_team_region").mkdir()
(extras / "migrations" / "20991231000001_index_team_region" / "migration.sql").write_text(
'CREATE INDEX CONCURRENTLY "LiteLLM_TeamTable_region_idx" ON "LiteLLM_TeamTable"("region");\n'
)
(extras / "migrations" / "20991231000002_index_team_alias").mkdir()
(extras / "migrations" / "20991231000002_index_team_alias" / "migration.sql").write_text(
'CREATE INDEX "LiteLLM_TeamTable_alias_idx" ON "LiteLLM_TeamTable"("team_alias");\n'
)
with psycopg.connect(scratch_database, autocommit=True) as conn:
result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations")
region: Final = conn.execute(
"SELECT 1 FROM information_schema.columns WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'region'"
).fetchone()
assert result.rejected is None, result
assert tuple((applied.name, applied.replayed) for applied in result.applied) == (
("20991231000000_add_team_region", ()),
("20991231000001_index_team_region", ("20991231000000_add_team_region",)),
("20991231000002_index_team_alias", ()),
)
assert region is None, "every transaction must have been rolled back"
with psycopg.connect(scratch_database) as conn:
probe_indexes: Final = conn.execute(
"SELECT indexname FROM pg_indexes "
"WHERE indexname IN ('LiteLLM_TeamTable_region_idx', 'LiteLLM_TeamTable_alias_idx')"
).fetchall()
assert probe_indexes == [], probe_indexes
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_a_database_with_tables_but_no_ledger_is_judged_by_the_schema_diff_instead_of_replaying_history(
scratch_database: str, tmp_path: Path
) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute("DROP TABLE _prisma_migrations")
matching: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations", baseline_check=lambda: None)
differing: Final = migration_dry_run.dry_run(
conn,
"public",
PACKAGE / "migrations",
baseline_check=lambda: (
"the schema differs from this build. prisma migrate diff:\nALTER TABLE x ADD COLUMN y TEXT;"
),
)
custom: Final = migration_dry_run.dry_run(
conn,
"public",
PACKAGE / "migrations",
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(
_extras_with_scratch_migration(
tmp_path, "20991231000000_custom", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "c" TEXT;\n'
)
),
)
assert matching.ok is True and matching.applied == ()
assert "nothing to dry-run" in migration_dry_run.report(matching)
assert differing.ok is False
assert "ALTER TABLE x ADD COLUMN y TEXT;" in migration_dry_run.report(differing)
assert custom.ok is False
assert "custom migration history" in migration_dry_run.report(custom)

24
tests/integration/routing/test_auto_router_management.py
View file

@ -9,7 +9,7 @@ import httpx
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value, until_asserted
from integration._support.database import read_rows, scratch_database
from integration._support.database import read_rows, scratch_database, write_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
@ -639,3 +639,25 @@ def test_auto_router_session_tracks_key_scoped_spend_and_bounds_long_session_ids
), response.text
assert tuple(request for request in complex_tier.drain() if request.method == "POST") == (), response.text
assert tuple(request for request in baseline.drain() if request.method == "POST") == (), response.text
@pytest.mark.parametrize("api_key,expected", [("unattributed", 2.0), ("raw:unattributed", 4.0)])
def test_benchmarks_preserve_literal_key_identity(gateway: Gateway, api_key: str, expected: float) -> None:
user_id: Final = f"benchmark-key-{uuid.uuid4().hex}"
for key, savings in (("", 3.0), ("unattributed", 2.0), ("raw:unattributed", 4.0)):
write_rows(
'INSERT INTO "LiteLLM_DailyUserSpend" '
"(id, user_id, date, api_key, model, autorouter_savings_spend, updated_at) "
"VALUES (%s, %s, '2026-09-09', %s, 'model', %s, now())",
(str(uuid.uuid4()), user_id, key, str(savings)),
)
try:
response: Final = gateway.request(
"GET",
"/auto_router/benchmarks",
params={"start_date": "2026-09-09", "end_date": "2026-09-09", "user_id": user_id, "api_key": api_key},
)
assert response.status_code == 200, response.text
assert object_value(object_value(response.json())["totals"])["saved_spend"] == expected
finally:
write_rows('DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id = %s', (user_id,))

1515
tests/integration/routing/test_non_chat_model_endpoint_gate.py Normal file
View file

File diff suppressed because it is too large Load diff

183
tests/integration/security/test_api_key_secret_reference.py
View file

@ -1,6 +1,9 @@
from __future__ import annotations
import base64
import json
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
@ -46,7 +49,18 @@ _SEARCH_PATH: Final = f"/v1/vector_stores/{_VECTOR_STORE}/search"
_BATCH_DEPLOYMENT: Final = "wif-vllm-batch"
_BEDROCK_MODEL: Final = "bedrock/converse/us.anthropic.claude-sonnet-5-5"
_PROVIDER_INPUT_TOKENS: Final = 7919
_UNSUPPORTED_REFERENCE: Final = "oidc/google/https://example.invalid"
_UNSUPPORTED_REFERENCE: Final = "oidc/github/https://example.invalid"
_GOOGLE_MODEL: Final = "hosted_vllm/wif-google"
_GOOGLE_AUDIENCES: Final = MappingProxyType(
{
"config_inline": "https://vllm-inline.example.test",
"config_credential": "https://vllm-credential.example.test",
"db_credential": "https://vllm-db-credential.example.test",
}
)
_UNAVAILABLE_AUDIENCE: Final = "https://vllm-unavailable.example.test"
_HUNG_AUDIENCE: Final = "https://vllm-hung.example.test"
_IDENTITY_PATH: Final = "/computeMetadata/v1/instance/service-accounts/default/identity"
_RUN_SECRETS_TOKEN: Final = "/run/secrets/wif-absent-token"
_SOURCES: Final = ("config_credential", "config_inline", "db_credential", "db_inline")
_UNREADABLE: Final = MappingProxyType(
@ -108,6 +122,45 @@ def _reference(path: Path) -> str:
return f"oidc/file/{path}"
def _google_reference(audience: str) -> str:
return f"oidc/google/{audience}"
def _jwt_segment(part: Mapping[str, JsonValue]) -> str:
return base64.urlsafe_b64encode(json.dumps(part).encode()).decode().rstrip("=")
@dataclass(frozen=True, slots=True)
class _Google:
expires: int
def token(self, audience: str) -> str:
claims: Final = _jwt_segment({"iss": "https://accounts.google.com", "aud": audience, "exp": self.expires})
return f"{_jwt_segment({'alg': 'RS256', 'typ': 'JWT'})}.{claims}.signature"
def bearers(self) -> frozenset[str]:
return frozenset(f"Bearer {self.token(audience)}" for audience in _GOOGLE_AUDIENCES.values())
def _metadata_server(google: _Google) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
target: Final = httpx.URL(request.target)
audience: Final = target.params.get("audience")
if audience is None or (request.method, target.path, request.headers.get("metadata-flavor")) != (
"GET",
_IDENTITY_PATH,
"Google",
):
return Reply(status=404, body=b"not found", content_type="text/plain")
if audience == _UNAVAILABLE_AUDIENCE:
return Reply(status=503, body=b"metadata server unavailable", content_type="text/plain")
if audience == _HUNG_AUDIENCE:
return Reply(chunks=(b"never sent",), gate_before_first=threading.Event(), content_type="text/plain")
return Reply(body=google.token(audience).encode(), content_type="text/plain")
return respond
def _secrets(directory: Path) -> _Secrets:
allowed: Final = directory / "secrets"
outside_dir: Final = directory / "outside"
@ -169,9 +222,9 @@ def _served(request: Request) -> Reply:
return openai_error(404)
def _provider(secrets: _Secrets) -> Callable[[Request], Reply]:
def _provider(secrets: _Secrets, google: _Google) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.headers.get("authorization") != f"Bearer {secrets.current()}":
if request.headers.get("authorization") not in {f"Bearer {secrets.current()}", *google.bearers()}:
return openai_error(401)
return _served(request)
@ -224,6 +277,27 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path:
"aws_bedrock_runtime_endpoint": provider.removesuffix("/v1"),
},
),
_deployment(
"wif-google-inline",
{
"model": _GOOGLE_MODEL,
"api_key": _google_reference(_GOOGLE_AUDIENCES["config_inline"]),
"api_base": provider,
},
),
_deployment("wif-google-credential", _named(_GOOGLE_MODEL, "wif-google")),
_deployment(
"wif-google-unavailable",
{"model": _GOOGLE_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
),
_deployment(
"wif-google-unavailable-count",
{"model": _CHAT_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
),
_deployment(
"wif-google-hung",
{"model": _GOOGLE_MODEL, "api_key": _google_reference(_HUNG_AUDIENCE), "api_base": provider},
),
_deployment(
"wif-mistral-passthrough",
{
@ -239,6 +313,7 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path:
_credential("wif-outside", _reference(secrets.outside), provider),
_credential("wif-unsupported", _UNSUPPORTED_REFERENCE, provider),
_credential("wif-run-secrets", f"oidc/file/{_RUN_SECRETS_TOKEN}", provider),
_credential("wif-google", _google_reference(_GOOGLE_AUDIENCES["config_credential"]), provider),
],
"vector_store_registry": [
{
@ -312,6 +387,8 @@ class _Rig:
owned: OwnedProxy
provider: Recorder
intruder: Recorder
metadata: Recorder
google: _Google
provider_base: str
intruder_base: str
secrets: _Secrets
@ -339,17 +416,24 @@ def _is_discovery(request: Request) -> bool:
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("api_key_secret_reference").resolve()
secrets: Final = _secrets(directory)
google: Final = _Google(expires=int(time.time()) + 3600)
with (
gateway_from_environment() as gateway,
wire_server(_provider(secrets)) as provider_wire,
wire_server(_provider(secrets, google)) as provider_wire,
wire_server(_intruder) as intruder_wire,
wire_server(_metadata_server(google)) as metadata_wire,
):
provider: Final = Recorder(provider_wire)
provider_base: Final = f"{provider_wire.url}/v1"
with owned_proxy_process(
gateway,
directory,
{**secrets.overrides(), "MISTRAL_API_BASE": provider_wire.url},
{
**secrets.overrides(),
"MISTRAL_API_BASE": provider_wire.url,
"GCE_METADATA_HOST": httpx.URL(metadata_wire.url).netloc.decode(),
"GCE_METADATA_TIMEOUT": "1",
},
config=_config(directory, secrets, provider_base),
workers=2,
healthcheck=WorkerHealthcheck.LENIENT,
@ -361,6 +445,8 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
owned=owned,
provider=provider,
intruder=Recorder(intruder_wire),
metadata=Recorder(metadata_wire),
google=google,
provider_base=provider_base,
intruder_base=f"{intruder_wire.url}/v1",
secrets=secrets,
@ -708,7 +794,92 @@ def test_team_row_naming_the_credential_sends_the_token_only_to_its_host(rig: _R
def test_startup_discovery_sends_the_boot_token(rig: _Rig) -> None:
sent: Final = frozenset(_bearers(tuple(request for request in rig.boot if _is_discovery(request))))
assert sent == {f"Bearer {rig.boot_token}"}, sent
assert sent - rig.google.bearers() == {f"Bearer {rig.boot_token}"}, sent
def _discovery_bearers(rig: _Rig) -> frozenset[str | None]:
return frozenset(_bearers(tuple(request for request in rig.provider.requests() if _is_discovery(request))))
def test_startup_discovery_sends_the_google_token(rig: _Rig) -> None:
expected: Final = frozenset(
f"Bearer {rig.google.token(_GOOGLE_AUDIENCES[source])}" for source in ("config_inline", "config_credential")
)
sent: Final = eventually(
partial(_discovery_bearers, rig), lambda seen: expected <= seen, seconds=30, return_last_on_timeout=True
)
assert expected <= sent, sent
def _google_model(rig: _Rig, scenario: Scenario, source: str) -> str:
match source:
case "config_inline":
return "wif-google-inline"
case "config_credential":
return "wif-google-credential"
case "db_credential":
credential: Final = f"db-google-{uuid.uuid4().hex}"
rig.owned.gateway.post(
"/credentials",
_credential(credential, _google_reference(_GOOGLE_AUDIENCES[source]), rig.provider_base),
)
scenario.cleanups.callback(_delete_credential, rig.owned.gateway, credential)
name, identity = _created_model(
rig.owned.gateway,
scenario,
_deployment(f"db-google-{uuid.uuid4().hex}", _named(_GOOGLE_MODEL, credential)),
)
eventually(partial(rig.owned.serves_only_on_every_worker, name, identity), bool, seconds=120)
return name
case _:
pytest.fail(f"unknown source {source}")
def _asked_audiences(rig: _Rig) -> frozenset[str | None]:
return frozenset(httpx.URL(request.target).params.get("audience") for request in rig.metadata.requests())
@pytest.mark.parametrize("source", tuple(_GOOGLE_AUDIENCES))
def test_google_reference_sends_the_metadata_server_token_for_its_audience(rig: _Rig, source: str) -> None:
audience: Final = _GOOGLE_AUDIENCES[source]
with rig.owned.gateway.scenario() as scenario:
marker: Final = _marker()
response: Final = _chat(rig.owned.gateway, _google_model(rig, scenario, source), marker)
sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
assert (response.status_code, sent) == (200, (f"Bearer {rig.google.token(audience)}",)), response.text
assert audience in _asked_audiences(rig)
def test_google_reference_fails_closed_when_the_metadata_server_fails(rig: _Rig) -> None:
marker: Final = _marker()
response: Final = _chat(rig.owned.gateway, "wif-google-unavailable", marker)
reached: Final = _carrying(rig.provider.requests(), marker)
assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
assert "oidc/google secret reference could not be resolved" in response.text, response.text
assert _UNAVAILABLE_AUDIENCE.removeprefix("https://") not in response.text, response.text
assert _UNAVAILABLE_AUDIENCE in _asked_audiences(rig)
def test_google_reference_answers_503_when_the_metadata_server_never_finishes_its_reply(rig: _Rig) -> None:
marker: Final = _marker()
response: Final = _chat(rig.owned.gateway, "wif-google-hung", marker)
reached: Final = _carrying(rig.provider.requests(), marker)
assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
assert "the metadata server did not return an ID token" in response.text, response.text
assert _HUNG_AUDIENCE in _asked_audiences(rig)
def test_input_token_count_answers_503_when_the_metadata_server_fails(rig: _Rig) -> None:
marker: Final = _marker()
response: Final = rig.owned.gateway.request(
"POST",
"/v1/responses/input_tokens",
{"model": "wif-google-unavailable-count", "input": f"say {marker}"},
headers=_CLOSE,
)
reached: Final = _carrying(rig.provider.requests(), marker)
assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
assert "oidc/google secret reference could not be resolved" in response.text, response.text
def _side_call(rig: _Rig, seam: str, marker: str) -> tuple[httpx.Response, Callable[[Request], bool]]:

178
tests/integration/spend/test_daily_activity_aggregated_breakdowns.py
View file

@ -1,3 +1,5 @@
import csv
import io
import uuid
from collections.abc import Sequence
from datetime import datetime, timedelta
@ -12,6 +14,7 @@ from tests.integration._support.database import write_rows
_URL: Final = "/user/daily/activity/aggregated"
_RESULTS: Final = TypeAdapter(list[dict[str, JsonValue]])
_NUMBER: Final = TypeAdapter(float)
def _unique_day() -> str:
@ -43,6 +46,181 @@ def _row_id() -> str:
return f"agg-{uuid.uuid4().hex}"
@pytest.mark.parametrize("suffix", ["/aggregated", ""])
def test_missing_attribution_reconciles_and_filters_without_key_collisions(gateway: Gateway, suffix: str) -> None:
day: Final = _unique_day()
rows: Final = (
("", None, None, 3.0),
("", "", "", 0.0),
("known-key", "unknown-model", None, -1.0),
("unattributed", "known-model", "", 2.0),
("raw:unattributed", "known-model", "named-group", 4.0),
("known-key", "", "unknown-model-group", 5.0),
)
_seed(
day,
[
(_row_id(), "bucket-user", day, key, model, group, "", "tool", "/mcp", 10, spend, 1, 1)
for key, model, group, spend in rows
],
)
params: Final = {"start_date": day, "end_date": day}
url: Final = f"/user/daily/activity{suffix}"
try:
response: Final = gateway.request("GET", url, params=params)
assert response.status_code == 200, response.text
body: Final = object_value(response.json())
results: Final = _RESULTS.validate_python(body["results"])
assert len(results) == 1
totals: Final = object_value(results[0]["metrics"])
assert totals["spend"] == 13.0
assert totals["api_requests"] == 6
breakdown: Final = object_value(results[0]["breakdown"])
for dimension in ("models", "model_groups", "api_keys", "providers", "mcp_servers", "endpoints"):
buckets: Final = _RESULTS.validate_python(list(object_value(breakdown[dimension]).values()))
for metric in ("spend", "api_requests", "prompt_tokens"):
assert (
sum(_NUMBER.validate_python(object_value(bucket["metrics"])[metric]) for bucket in buckets)
== totals[metric]
)
models: Final = object_value(breakdown["models"])
assert set(models) == {"known-model", "unknown-model"}
assert object_value(object_value(models["unknown-model"])["metrics"])["api_requests"] == 4
groups: Final = object_value(breakdown["model_groups"])
assert set(groups) == {"unknown-model-group", "unknown-model", "known-model", "named-group"}
assert object_value(object_value(groups["unknown-model-group"])["metrics"])["spend"] == 8.0
keys: Final = object_value(breakdown["api_keys"])
assert set(keys) == {"unattributed", "raw:unattributed", "raw:raw:unattributed", "known-key"}
unknown_key: Final = object_value(keys["unattributed"])
assert object_value(unknown_key["metadata"]) == {
"key_alias": "Unattributed",
"key_exists": False,
"team_id": None,
"user_id": None,
"user_email": None,
}
for dimension in ("models", "model_groups", "providers", "mcp_servers", "endpoints"):
for bucket in _RESULTS.validate_python(list(object_value(breakdown[dimension]).values())):
nested: Final = _RESULTS.validate_python(list(object_value(bucket["api_key_breakdown"]).values()))
assert (
sum(_NUMBER.validate_python(object_value(key["metrics"])["spend"]) for key in nested)
== object_value(bucket["metrics"])["spend"]
)
for key, spend, count in (
("unattributed", 3.0, 2),
("raw:unattributed", 2.0, 1),
("raw:raw:unattributed", 4.0, 1),
):
filtered: Final = gateway.request("GET", url, params={**params, "api_key": key})
assert filtered.status_code == 200, filtered.text
metadata: Final = object_value(object_value(filtered.json())["metadata"])
assert metadata["total_spend"] == spend
assert metadata["total_api_requests"] == count
model_response: Final = gateway.request("GET", url, params={**params, "model": "unknown-model"})
assert model_response.status_code == 200, model_response.text
model_metadata: Final = object_value(object_value(model_response.json())["metadata"])
assert model_metadata["total_spend"] == 7.0
assert model_metadata["total_api_requests"] == 4
finally:
_clean(day)
@pytest.mark.parametrize(
("route", "extra", "expected"),
[
("keys", {}, {"unattributed": 3.0, "raw:unattributed": 2.0, "known-key": 7.0}),
("search", {"search": "Unattributed"}, {"unattributed": 3.0, "raw:unattributed": 2.0}),
("model_top_keys", {"model_group": "unknown-model-group"}, {"unattributed": 3.0}),
(
"model_top_keys",
{"model_group": "unknown-model", "by_model_group": "false"},
{"unattributed": 3.0, "raw:unattributed": 2.0},
),
],
)
def test_missing_attribution_key_drilldowns(
gateway: Gateway,
route: str,
extra: dict[str, str],
expected: dict[str, float],
) -> None:
day: Final = _unique_day()
_seed(
day,
[
(_row_id(), "bucket-user", day, "", None, None, "", None, None, 0, 3.0, 1, 1),
(_row_id(), "bucket-user", day, "unattributed", "unknown-model", "", "", None, None, 0, 2.0, 1, 1),
(_row_id(), "bucket-user", day, "known-key", "known-model", "", "", None, None, 0, 7.0, 1, 1),
],
)
try:
response: Final = gateway.request(
"GET", f"{_URL}/{route}", params={"start_date": day, "end_date": day, **extra}
)
assert response.status_code == 200, response.text
rows: Final = _RESULTS.validate_python(object_value(response.json())["api_keys"])
assert {str(row["api_key"]): object_value(row["metrics"])["spend"] for row in rows} == expected
unattributed: Final = next(row for row in rows if row["api_key"] == "unattributed")
assert object_value(unattributed["metadata"])["key_alias"] == "Unattributed"
assert object_value(unattributed["metadata"])["key_exists"] is False
finally:
_clean(day)
@pytest.mark.parametrize("format", ["json", "csv"])
@pytest.mark.parametrize(
("export_type", "field", "expected"),
[
("daily_with_keys", "api_key", {"unattributed": 3.0, "raw:unattributed": 2.0, "raw:raw:unattributed": 4.0}),
("daily_with_models", "model", {"unknown-model": 5.0, "known-model": 4.0}),
],
)
def test_export_bucket_ids_round_trip_across_streaming_batches(
gateway: Gateway, format: str, export_type: str, field: str, expected: dict[str, float]
) -> None:
day: Final = _unique_day()
_seed(
day,
[
(_row_id(), "export-user", day, "", None, None, "", None, None, 0, 3.0, 1, 1),
(_row_id(), "export-user", day, "", "", "", "", None, None, 0, 0.0, 1, 1),
(_row_id(), "export-user", day, "unattributed", "unknown-model", "", "", None, None, 0, 2.0, 1, 1),
(_row_id(), "export-user", day, "raw:unattributed", "known-model", "", "", None, None, 0, 4.0, 1, 1),
],
)
write_rows(
'INSERT INTO "LiteLLM_DailyUserSpend" (id, user_id, date, api_key, model, spend, api_requests, updated_at) '
"SELECT %s || i, 'export-user', %s, 'raw:a-key-' || i, 'a-model-' || i, 1, 1, now() "
"FROM generate_series(1, 1001) i",
(str(uuid.uuid4()), day),
)
params: Final = {"start_date": day, "end_date": day}
try:
response: Final = gateway.request(
"GET", "/user/daily/activity/export", params={**params, "export_type": export_type, "format": format}
)
assert response.status_code == 200, response.text
rows: Final = (
_RESULTS.validate_python(response.json())
if format == "json"
else _RESULTS.validate_python(list(csv.DictReader(io.StringIO(response.text))))
)
buckets: Final = {str(row[field]): _NUMBER.validate_python(row["spend"]) for row in rows}
filler: Final = "raw:raw:a-key-" if field == "api_key" else "a-model-"
assert buckets == {**{f"{filler}{i}": 1.0 for i in range(1, 1002)}, **expected}
assert len(rows) == len(buckets)
assert sum(buckets.values()) == 1010.0
for bucket, spend in expected.items():
filtered: Final = _activity(gateway, day, **{field: bucket})
assert object_value(filtered["metadata"])["total_spend"] == spend
if field == "api_key":
unattributed: Final = next(row for row in rows if row["api_key"] == "unattributed")
assert unattributed["key_alias"] == "Unattributed"
assert unattributed["user_id"] in (None, "")
finally:
_clean(day)
def _ranked_key_rows(day: str, count: int) -> list[tuple[object, ...]]:
return [
(

14
tests/integration/spend/test_daily_activity_routes.py
View file

@ -451,9 +451,23 @@ def test_tag_daily_activity_scopes_internal_user_keys_and_filters(
seeded_row(TAG_SPEND, "tag", "beta", digest_one_b, DAY),
seeded_row(TAG_SPEND, "tag", "gamma", digest_one_a, DAY),
seeded_row(TAG_SPEND, "tag", "alpha", digest_two, DAY),
seeded_row(TAG_SPEND, "tag", "alpha", "", DAY),
replace(standard_model_row, model="gpt-4o"),
)
with daily_rows(rows):
for suffix in ("", "/aggregated"):
bucket_params: Final = {"start_date": DAY, "end_date": DAY, "tags": "alpha", "api_key": "unattributed"}
bucket_url: Final = f"/tag/daily/activity{suffix}"
denied_bucket: Final = gateway.request("GET", bucket_url, params=bucket_params, key=key_one_a)
admin_bucket: Final = gateway.request("GET", bucket_url, params=bucket_params)
assert denied_bucket.status_code == 200, denied_bucket.text
assert admin_bucket.status_code == 200, admin_bucket.text
denied_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(denied_bucket.content)
allowed_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(admin_bucket.content)
assert denied_body.results == []
assert denied_body.metadata.total_api_requests == 0
assert allowed_body.metadata.total_api_requests == 1
assert set(allowed_body.results[0].breakdown.api_keys) == {"unattributed"}
scoped: Final = _tag_activity(
gateway,
key=key_one_a,

224
tests/integration/spend/test_shutdown_flush.py
View file

@ -1,11 +1,14 @@
import hashlib
import json
import os
import re
import threading
import time
import uuid
from collections.abc import Callable, Iterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Final
@ -26,6 +29,225 @@ MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new"
COMMIT_DELAY_SECONDS: Final = 15
BURST_REQUESTS: Final = 30
_GATED_LOGGER: Final = """
import os
import re
import httpx
from litellm.integrations.custom_logger import CustomLogger
class ShutdownLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
async with httpx.AsyncClient(timeout=25, trust_env=False) as client:
await client.post(os.environ["SHUTDOWN_CALLBACK_URL"])
logger = ShutdownLogger()
"""
def _tag_shutdown_config(directory: Path, upstream: str, redis_buffer: bool, *, gated: bool = False) -> Path:
path: Final = directory / "tag-shutdown.yaml"
if gated:
(directory / "shutdown_callbacks.py").write_text(_GATED_LOGGER)
path.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": "shutdown-model",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": upstream + "/v1",
"api_key": "synthetic-key",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
},
}
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
"proxy_batch_write_at": 3600,
"use_redis_transaction_buffer": redis_buffer,
"disable_model_info_refresh": True,
},
"litellm_settings": {
"cache": True,
"cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
"callbacks": ["shutdown_callbacks.logger"] if gated else [],
},
}
)
)
return path
def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None:
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": "shutdown-model",
"messages": [{"role": "user", "content": uuid.uuid4().hex}],
"metadata": {"tags": list(tags)},
},
key=key,
)
assert response.status_code == 200, response.text
def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None:
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")):
rows: Final = read_rows(
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
)
assert rows[0]["requests"] == requests, (table, rows)
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,))
assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9)
for tag in tags:
tag_rows: Final = read_rows(
'SELECT sum(api_requests)::int AS requests, sum(spend) AS spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s',
(hashed, tag),
)
assert tag_rows[0]["requests"] == requests, (tag, tag_rows)
assert tag_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (tag, tag_rows)
@pytest.mark.parametrize("redis_buffer", [False, True])
def test_tagged_request_survives_shutdown_before_the_tag_timer(
gateway: Gateway, tmp_path: Path, redis_buffer: bool
) -> None:
tags: Final = (f"shutdown-a-{uuid.uuid4().hex}", f"shutdown-b-{uuid.uuid4().hex}")
overrides: Final = {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG"}
with wire_server(_provider) as wire:
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer)
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned:
key: Final = string_value(
owned.gateway.post("/key/generate", {"user_id": f"shutdown-{uuid.uuid4().hex}"})["key"]
)
_tagged_chat(owned.gateway, key, tags)
owned.begin_shutdown()
owned.process.wait(timeout=90)
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted:
_assert_shutdown_accounting(key, tags, 1)
restarted.begin_shutdown()
restarted.process.wait(timeout=90)
_assert_shutdown_accounting(key, tags, 1)
@pytest.mark.parametrize("redis_buffer", [False, True])
@pytest.mark.parametrize("overflow", [False, True])
def test_callback_backlog_is_accounted_before_shutdown_disconnects(
gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool
) -> None:
release: Final = threading.Event()
requests: Final = 6 if overflow else 3
tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}")
with (
wire_server(_provider) as wire,
wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger,
):
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
with owned_proxy_process(
gateway,
tmp_path,
{
"DATABASE_URL": os.environ["DATABASE_URL"],
"LOGGING_WORKER_CONCURRENCY": "2",
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
"SHUTDOWN_CALLBACK_URL": logger.url,
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
"LITELLM_LOG": "DEBUG",
},
config=config,
) as owned:
key: Final = string_value(
owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"]
)
try:
for _ in range(requests):
_tagged_chat(owned.gateway, key, tags)
eventually(logger.received.qsize, lambda count: count >= 2, seconds=10)
owned.begin_shutdown()
eventually(
lambda: (
owned.process.poll() is not None
or "Draining success callbacks before the final spend flush" in owned.log.read_text()
),
lambda ready: ready,
seconds=30,
)
release.set()
owned.process.wait(timeout=90)
_assert_shutdown_accounting(key, tags, requests)
finally:
release.set()
@pytest.mark.parametrize("gunicorn", [False, True])
def test_sigterm_http_drain_consumes_callback_budget(gateway: Gateway, tmp_path: Path, gunicorn: bool) -> None:
release_request: Final = threading.Event()
release_callback: Final = threading.Event()
tags: Final = (f"deadline-a-{uuid.uuid4().hex}", f"deadline-b-{uuid.uuid4().hex}")
def held_provider(request: Request) -> Reply:
response: Final = _provider(request)
return replace(response, chunks=(response.body,), gate_before_first=release_request)
with (
wire_server(held_provider) as wire,
wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release_callback)) as logger,
):
config: Final = _tag_shutdown_config(tmp_path, wire.url, False, gated=True)
with (
owned_proxy_process(
gateway,
tmp_path,
{
"DATABASE_URL": os.environ["DATABASE_URL"],
"SHUTDOWN_CALLBACK_URL": logger.url,
"GRACEFUL_SHUTDOWN_TIMEOUT": "30",
"LITELLM_LOG": "DEBUG",
},
config=config,
extra_arguments=("--run_gunicorn",) if gunicorn else (),
) as owned,
ThreadPoolExecutor(max_workers=1) as executor,
):
key: Final = string_value(
owned.gateway.post("/key/generate", {"user_id": f"deadline-{uuid.uuid4().hex}"})["key"]
)
request: Final = executor.submit(_tagged_chat, owned.gateway, key, tags)
try:
eventually(wire.received.qsize, lambda count: count == 1, seconds=10)
owned.begin_shutdown()
eventually(
lambda: "graceful_shutdown_started" in owned.log.read_text(), lambda stamped: stamped, seconds=5
)
assert not request.done(), "The shutdown clock must start before the HTTP request drains"
stamped: Final = time.monotonic()
eventually(time.monotonic, lambda now: now - stamped >= 2, seconds=5)
held_seconds: Final = time.monotonic() - stamped
release_request.set()
request.result(timeout=10)
budget: Final = eventually(
lambda: re.search(r"remaining_s=([\d.]+) callback_timeout_s=([\d.]+)", owned.log.read_text()),
lambda match: match is not None,
seconds=15,
)
assert budget is not None
remaining: Final = float(budget[1])
allowance: Final = float(budget[2])
assert remaining <= 30 - held_seconds + 0.05
assert allowance == pytest.approx(max(0, remaining - 10), abs=0.002)
release_callback.set()
owned.process.wait(timeout=45)
_assert_shutdown_accounting(key, tags, 1)
finally:
release_request.set()
release_callback.set()
def _api_requests(table: str, column: str, identity: str) -> int:
rows: Final = read_rows(

95
tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py
View file

@ -27,7 +27,12 @@ _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_PROVIDER_KEY: Final = "synthetic-provider-key"
_GUARDRAIL: Final = "audit-text-rebuild-tool-permission"
_TEXT_MODEL: Final = f"audit-text-instruct-{uuid.uuid4().hex}"
_GOVERNED_MODEL: Final = f"audit-text-governed-{uuid.uuid4().hex}"
_KEYWORD_FILTER: Final = "audit-text-keyword-filter"
_POLICY: Final = "audit-text-post-call-governance"
_PIECES: Final = ("Hello", " streamed", " world")
_MASKED_PIECES: Final = ("A pers", "im", "mon tart")
_BLOCKED_PIECES: Final = ("A kum", "quat tart")
_PROMPT_TOKENS: Final = 11
_COMPLETION_TOKENS: Final = 3
_USAGE: Final[dict[str, JsonValue]] = {
@ -97,6 +102,14 @@ def _text_choice(text: str, finish_reason: str | None) -> dict[str, JsonValue]:
return {"text": text, "index": 0, "logprobs": None, "finish_reason": finish_reason}
def _pieces_for(identity: str) -> tuple[str, ...]:
if identity.startswith("cmpl-mask-"):
return _MASKED_PIECES
if identity.startswith("cmpl-block-"):
return _BLOCKED_PIECES
return _PIECES
def _respond(request: Request) -> Reply:
body: Final = JSON_OBJECT.validate_json(request.body)
assert body.get("stream") is True, body
@ -106,7 +119,7 @@ def _respond(request: Request) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=(
*(_text_frame(identity, [_text_choice(piece, None)]) for piece in _PIECES),
*(_text_frame(identity, [_text_choice(piece, None)]) for piece in _pieces_for(identity)),
_text_frame(identity, [_text_choice("", "length")]),
_text_frame(identity, [], usage=_USAGE),
b"data: [DONE]\n\n",
@ -114,20 +127,25 @@ def _respond(request: Request) -> Reply:
)
def _text_deployment(model_name: str, wire: Wire) -> dict[str, JsonValue]:
return {
"model_name": model_name,
"litellm_params": {
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
"api_base": f"{wire.url}/v1",
"api_key": _PROVIDER_KEY,
},
}
def _config(directory: Path, wire: Wire) -> Path:
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
config: Final = {
**base,
"model_list": [
*list_value(base["model_list"]),
{
"model_name": _TEXT_MODEL,
"litellm_params": {
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
"api_base": f"{wire.url}/v1",
"api_key": _PROVIDER_KEY,
},
},
_text_deployment(_TEXT_MODEL, wire),
_text_deployment(_GOVERNED_MODEL, wire),
],
"guardrails": [
{
@ -139,8 +157,30 @@ def _config(directory: Path, wire: Wire) -> Path:
"default_action": "allow",
"rules": [],
},
}
},
{
"guardrail_name": _KEYWORD_FILTER,
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "post_call",
"default_on": False,
"blocked_words": [
{"keyword": "persimmon", "action": "MASK"},
{"keyword": "kumquat", "action": "BLOCK"},
],
},
},
],
"policies": {
_POLICY: {
"guardrails": {"add": [_KEYWORD_FILTER]},
"pipeline": {
"mode": "post_call",
"steps": [{"guardrail": _KEYWORD_FILTER, "on_pass": "allow", "on_fail": "block"}],
},
}
},
"policy_attachments": [{"policy": _POLICY, "models": [_GOVERNED_MODEL]}],
}
path: Final = directory / f"text-rebuild-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
@ -163,14 +203,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
yield _Rig(wire, owned)
def _stream(rig: _Rig, identity: str) -> _Stream:
def _stream(rig: _Rig, identity: str, model: str = _TEXT_MODEL) -> _Stream:
gateway: Final = rig.proxy.gateway
with gateway.client.stream(
"POST",
"/v1/completions",
headers={"Authorization": f"Bearer {gateway.key}"},
json={
"model": _TEXT_MODEL,
"model": model,
"prompt": identity,
"max_tokens": 16,
"stream": True,
@ -212,6 +252,37 @@ def test_guarded_text_completion_stream_reaches_the_caller_whole_with_its_usage(
assert _GUARDRAIL in list_value(object_value(row["metadata"])["applied_guardrails"]), row
def test_post_call_pipeline_masks_a_streamed_text_completion(rig: _Rig) -> None:
identity: Final = f"cmpl-mask-{uuid.uuid4().hex}"
stream: Final = _stream(rig, identity, _GOVERNED_MODEL)
assert stream.status == 200, stream.body
assert stream.body.rstrip().endswith(b"data: [DONE]"), stream.body[-400:]
assert stream.text == "A [KEYWORD_REDACTED] tart", stream.frames
assert b"persimmon" not in stream.body, stream.body
assert stream.finish_reasons == ("length",), stream.frames
assert [(usage["prompt_tokens"], usage["completion_tokens"]) for usage in stream.usages] == [
(_PROMPT_TOKENS, _COMPLETION_TOKENS)
], stream.frames
(row,) = eventually(
lambda: rig.proxy.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,)),
lambda rows: len(rows) == 1,
seconds=70,
)
assert _KEYWORD_FILTER in list_value(object_value(row["metadata"])["applied_guardrails"]), row
def test_post_call_pipeline_blocks_a_streamed_text_completion(rig: _Rig) -> None:
identity: Final = f"cmpl-block-{uuid.uuid4().hex}"
stream: Final = _stream(rig, identity, _GOVERNED_MODEL)
assert stream.status == 400, stream.body
assert b"Content blocked: keyword 'kumquat' detected" in stream.body, stream.body
assert b"quat tart" not in stream.body, stream.body
def _together_chunk(choices: list[JsonValue], **extra: JsonValue) -> bytes:
return _frame(
{

1
tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py
View file

@ -32,6 +32,7 @@ class TestTryShortCircuitSearch:
("github_copilot", "github_copilot/claude-sonnet-5.5"),
("databricks", "databricks/databricks-claude-sonnet-5-5"),
("databricks", "databricks/system.ai.claude-sonnet-5-5"),
("dashscope", "dashscope/qwen3.7-plus"),
],
)
async def test_short_circuits_single_web_search_tool(self, provider, model):

246
tests/unit/litellm_core_utils/test_logging_worker.py
View file

@ -82,6 +82,252 @@ class _RecordCollector(logging.Handler):
self.records.append(record)
@pytest.mark.asyncio
@pytest.mark.parametrize("running", [False, True])
async def test_stop_discards_waiting_callbacks_without_executing_them(running: bool) -> None:
worker: Final = LoggingWorker(concurrency=1)
started: Final = asyncio.Event()
cancelled: Final = asyncio.Event()
executed: Final = asyncio.Event()
async def blocked() -> None:
started.set()
try:
await asyncio.Event().wait()
finally:
cancelled.set()
async def waiting() -> None:
executed.set()
if running:
worker.ensure_initialized_and_enqueue(blocked())
await started.wait()
else:
worker._ensure_queue()
worker.enqueue(waiting())
await worker.stop(discard_pending=True)
assert not executed.is_set()
assert cancelled.is_set() == running
assert worker._queue is not None
assert worker._queue.empty()
await asyncio.wait_for(worker._queue.join(), timeout=1)
@pytest.mark.asyncio
async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_started() -> None:
worker: Final = LoggingWorker(concurrency=1)
executed: Final = asyncio.Event()
async def callback() -> None:
executed.set()
worker.ensure_initialized_and_enqueue(callback())
await asyncio.sleep(0)
await worker.stop(discard_pending=True)
assert not executed.is_set()
worker._flush_on_exit()
assert not executed.is_set()
assert worker._queue is not None
await asyncio.wait_for(worker._queue.join(), timeout=1)
@pytest.mark.asyncio
async def test_stop_timeout_keeps_callbacks_observable_until_queue_cleanup_finishes() -> None:
worker: Final = LoggingWorker(concurrency=1)
started: Final = asyncio.Event()
release: Final = asyncio.Event()
completed: Final = asyncio.Event()
async def running() -> None:
started.set()
await asyncio.Event().wait()
async def queued() -> None:
await release.wait()
completed.set()
worker.ensure_initialized_and_enqueue(running())
await started.wait()
worker.enqueue(queued())
try:
await worker.stop(timeout=0)
assert not completed.is_set()
release.set()
await worker.flush()
assert completed.is_set(), "Queue cleanup must survive a bounded cancellation wait"
finally:
release.set()
await worker.stop()
@pytest.mark.asyncio
async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> None:
"""A callback rejected by a full queue is run by the aggressive clear. It must not consume a
queue.task_done() it never contributed, otherwise flush() returns while a real callback is still
running and the shutdown drain cancels it with budget to spare.
"""
worker: Final = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1)
gates: Final = {name: asyncio.Event() for name in ("first", "second", "third")}
completed: Final[list[str]] = []
async def gated(name: str) -> None:
await gates[name].wait()
completed.append(name)
async def instant() -> None:
completed.append("overflow")
worker.start()
worker.enqueue(gated("first"))
await asyncio.sleep(0)
worker.enqueue(gated("second"))
worker.enqueue(gated("third"))
assert worker._queue is not None and worker._queue.full(), "precondition: the queue is full"
worker.enqueue(instant())
for _ in range(5):
await asyncio.sleep(0)
gates["first"].set()
gates["second"].set()
asyncio.get_running_loop().call_later(0.2, gates["third"].set)
await asyncio.wait_for(worker.flush(), timeout=2.0)
assert "third" in completed, "flush returned while a queued callback was still unfinished"
assert sorted(completed) == ["first", "overflow", "second", "third"]
await worker.stop()
@pytest.mark.asyncio
@pytest.mark.parametrize("queue_size", (1, 2))
async def test_flush_waits_for_delayed_overflow_and_preserves_callback_context(queue_size: int) -> None:
retry_started: Final = asyncio.Event()
release_retry: Final = asyncio.Event()
completed: Final[asyncio.Queue[tuple[int, str]]] = asyncio.Queue()
identity: Final = contextvars.ContextVar("overflow_identity", default="outside")
async def retry_sleep(delay: float) -> None:
retry_started.set()
await release_retry.wait()
worker: Final = LoggingWorker(max_queue_size=queue_size, concurrency=1, retry_sleep=retry_sleep)
async def callback(index: int) -> None:
completed.put_nowait((index, identity.get()))
worker.start()
token: Final = identity.set("request")
for index in range(4):
worker.enqueue(callback(index))
identity.reset(token)
flushing: Final = asyncio.create_task(worker.flush())
try:
await asyncio.wait_for(retry_started.wait(), timeout=2)
assert worker._queue is not None
await asyncio.wait_for(worker._queue.join(), timeout=2)
await asyncio.sleep(0)
assert not flushing.done(), "Queue completion must not hide delayed overflow callbacks"
assert worker.pending_count > 0
release_retry.set()
await asyncio.wait_for(flushing, timeout=2)
assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == [
(index, "request") for index in range(4)
]
assert worker.pending_count == 0
finally:
release_retry.set()
await worker.stop(discard_pending=True)
await asyncio.gather(flushing, return_exceptions=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("start_retry", (False, True))
async def test_shutdown_discards_overflow_without_late_enqueue_or_replay(start_retry: bool) -> None:
import inspect
retry_started: Final = asyncio.Event()
release_retry: Final = asyncio.Event()
executed: Final[asyncio.Queue[int]] = asyncio.Queue()
async def retry_sleep(delay: float) -> None:
retry_started.set()
await release_retry.wait()
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep)
async def callback(index: int) -> None:
executed.put_nowait(index)
callbacks: Final = tuple(callback(index) for index in range(4))
worker._ensure_queue()
for callback_coro in callbacks:
worker.enqueue(callback_coro)
if start_retry:
await asyncio.wait_for(retry_started.wait(), timeout=2)
await worker.stop(discard_pending=True, timeout=1)
completed_at_stop: Final = executed.qsize()
release_retry.set()
await worker.flush()
try:
assert executed.qsize() == completed_at_stop, "Discarded overflow must not run when the worker restarts"
assert worker.pending_count == 0
assert all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in callbacks)
finally:
await worker.stop(discard_pending=True)
@pytest.mark.asyncio
async def test_shutdown_cancels_aggressive_clear_callbacks_before_returning() -> None:
started: Final = asyncio.Event()
cancelled: Final = asyncio.Event()
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1)
async def blocked() -> None:
started.set()
try:
await asyncio.Event().wait()
finally:
cancelled.set()
async def noop() -> None:
pass
worker._ensure_queue()
worker.enqueue(blocked())
worker.enqueue(noop())
worker.enqueue(noop())
await asyncio.wait_for(started.wait(), timeout=2)
await worker.stop(discard_pending=True, timeout=1)
assert cancelled.is_set(), "Overflow execution must participate in shutdown cancellation"
assert worker.pending_count == 0
assert worker._queue is not None
await asyncio.wait_for(worker._queue.join(), timeout=1)
@pytest.mark.asyncio
async def test_stopped_overflow_can_be_flushed_when_not_discarded() -> None:
release_retry: Final = asyncio.Event()
completed: Final[asyncio.Queue[int]] = asyncio.Queue()
async def retry_sleep(delay: float) -> None:
await release_retry.wait()
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep)
async def callback(index: int) -> None:
completed.put_nowait(index)
worker._ensure_queue()
for index in range(4):
worker.enqueue(callback(index))
await worker.stop()
release_retry.set()
try:
await asyncio.wait_for(worker.flush(), timeout=2)
assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == list(range(4))
finally:
await worker.stop(discard_pending=True)
class TestLoggingWorker:
"""Test cases for LoggingWorker functionality."""

35
tests/unit/litellm_core_utils/test_served_output_texts.py
View file

@ -9,7 +9,16 @@ from litellm.litellm_core_utils.served_output_texts import (
served_output_texts,
served_stream_output_texts,
)
from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices
from litellm.types.utils import (
Choices,
Delta,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
TextChoices,
TextCompletionResponse,
)
RAW = "Card: 4111 1111 1111 1111"
MASKED = "Card: <CREDIT_CARD>"
@ -83,6 +92,30 @@ def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_o
assert served_stream_output_texts(chunks) == ("first choice", MASKED)
def _text_completion_chunk(text: str | None, index: int = 0) -> TextCompletionResponse:
return TextCompletionResponse(object="text_completion", choices=[TextChoices(index=index, text=text)])
def test_served_stream_output_texts_joins_text_completion_chunks_per_choice():
chunks = [
_text_completion_chunk("Card: ", 0),
_text_completion_chunk("second", 1),
_text_completion_chunk("<CREDIT_CARD>", 0),
_text_completion_chunk(None, 1),
TextCompletionResponse(object="text_completion", choices=[]),
]
assert served_stream_output_texts(chunks) == (MASKED, "second")
def test_overlay_replaces_text_completion_choice_text_with_served_text():
logged = {"id": "x", "object": "text_completion", "choices": [{"index": 0, "text": RAW}, {"index": 1, "text": RAW}]}
assert overlay_served_output_texts(logged, (MASKED,)) == {
"id": "x",
"object": "text_completion",
"choices": [{"index": 0, "text": MASKED}, {"index": 1, "text": RAW}],
}
@pytest.mark.parametrize("content", (None, "", "first"))
def test_null_choice_index_joins_choice_zero_without_changing_other_choices(content: str | None) -> None:
chunks: Final = (

96
tests/unit/litellm_proxy_extras/test_migration_data_checks.py
View file

@ -12,98 +12,15 @@ from litellm_proxy_extras.migration_data_checks import (
_count_query,
btree_max_entry_bytes,
checks_for_index,
checks_for_migration,
checks_for_statement,
migration_data_check_timeout_seconds,
report,
statements,
)
MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0: Final = 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")'
def _shape(checks: tuple[Check, ...]) -> tuple[tuple[str, str, tuple[str, ...]], ...]:
return tuple((check.kind, check.table, check.columns) for check in checks)
@pytest.mark.parametrize(
("statement", "expected"),
(
(
API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,
(("too_long", "LiteLLM_SpendLogs", ("api_key", "startTime")),),
),
(
'CREATE UNIQUE INDEX CONCURRENTLY "T_a_b_key" ON "public"."T" USING btree ("a", "b" DESC)',
(("too_long", "T", ("a", "b")), ("duplicates", "T", ("a", "b"))),
),
('CREATE INDEX "T_a_idx" ON "T" USING gin ("a")', ()),
('CREATE INDEX "T_a_idx" ON "T" (lower("a"))', ()),
('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', (("nulls", "T", ("a",)),)),
('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL', (("any_row", "T", ("a",)),)),
('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL DEFAULT \'x\'', ()),
('ALTER TABLE "T" ADD COLUMN "a" TEXT', ()),
(
'ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") ON DELETE CASCADE',
(("orphans", "T", ("r",)),),
),
('ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") NOT VALID', ()),
('ALTER TABLE "T" ADD CONSTRAINT "T_a_key" UNIQUE ("a", "b")', (("duplicates", "T", ("a", "b")),)),
(
'ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0 AND "b" IS NOT NULL)',
(("check_violations", "T", ()),),
),
('ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0) NOT VALID', ()),
(
'ALTER TABLE "T" ADD COLUMN "c" TEXT, ALTER COLUMN "b" SET NOT NULL, ADD CONSTRAINT "u" UNIQUE ("c")',
(("nulls", "T", ("b",)), ("duplicates", "T", ("c",))),
),
('CREATE TABLE "T" ("id" TEXT NOT NULL, CONSTRAINT "T_pkey" PRIMARY KEY ("id"))', ()),
('DROP INDEX IF EXISTS "T_a_idx"', ()),
("SELECT 1", ()),
),
)
def test_a_statement_gets_one_check_per_way_existing_rows_can_make_it_fail(
statement: str, expected: tuple[tuple[str, str, tuple[str, ...]], ...]
) -> None:
assert _shape(checks_for_statement(statement, "migration x")) == expected
def test_a_foreign_key_check_keeps_the_referenced_table_and_columns() -> None:
(check,) = checks_for_statement(
'ALTER TABLE "T" ADD CONSTRAINT "fk" FOREIGN KEY ("r", "s") REFERENCES "R"("id", "sub")', "migration x"
)
assert (check.referenced_table, check.referenced_columns) == ("R", ("id", "sub"))
def test_a_check_constraint_keeps_its_expression_for_the_count_query() -> None:
(check,) = checks_for_statement('ALTER TABLE "T" ADD CONSTRAINT "c" CHECK ("a" > 0)', "migration x")
assert check.expression == '"a" > 0'
def test_comments_and_blank_statements_are_dropped_before_matching() -> None:
migration_sql: Final = "-- CreateIndex\n" + API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0 + ";\n\n-- done; really\n"
assert statements(migration_sql) == (API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,)
assert _shape(checks_for_migration("x", migration_sql)) == _shape(
checks_for_statement(API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0, "x")
)
def test_every_shipped_migration_is_readable_and_the_later_indexes_on_text_columns_get_a_check() -> None:
derived: Final = {
folder.name: checks_for_migration(folder.name, (folder / "migration.sql").read_text())
for folder in MIGRATIONS.iterdir()
if folder.is_dir()
}
assert len(derived) > 200
everywhere: Final = tuple(check for checks in derived.values() for check in checks) # comprehension-ok: flattening
assert any(
check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "session_id" in check.columns
for check in everywhere
)
def test_a_request_log_index_gets_one_entry_size_check_over_all_its_columns() -> None:
assert _shape(
checks_for_index("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime"))
@ -142,14 +59,19 @@ def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_wa
assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text
assert 'fix: Shorten or hash the long "LiteLLM_SpendLogs" (api_key, startTime) values' in text
assert text.startswith("Existing rows would make a pending schema change fail:")
assert "No schema change was made" in text
assert "No index was built" in text
def test_a_scan_that_timed_out_is_reported_as_unverified_with_the_knob_to_turn() -> None:
(check,) = checks_for_statement('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', "migration 20270101000000_x")
(check,) = checks_for_index(
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
)
text: Final = report((CouldNotVerify(check, 600.0),))
assert '"T".a: could not be checked within 600 s, refusing rather than guessing' in text
assert 'blocks: ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL (migration 20270101000000_x)' in text
assert (
'"LiteLLM_SpendLogs".api_key, startTime: could not be checked within 600 s, refusing rather than guessing'
in text
)
assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text
assert MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR in text

80
tests/unit/litellm_proxy_extras/test_migration_dry_run.py Normal file
View file

@ -0,0 +1,80 @@
from typing import Final
import pytest
from litellm_proxy_extras.migration_dry_run import Applied, Baseline, DryRun, Rejected, report, transactional
INDEX: Final = 'CREATE INDEX IF NOT EXISTS "a_idx" ON "T"("c")'
@pytest.mark.parametrize(
("migration_sql", "expected"),
(
('CREATE INDEX CONCURRENTLY IF NOT EXISTS "a_idx" ON "T"("c")', INDEX),
('create unique index concurrently "u" on "T"("d")', 'create unique index "u" on "T"("d")'),
('DROP INDEX CONCURRENTLY "x"', 'DROP INDEX "x"'),
('REINDEX INDEX CONCURRENTLY "y"', 'REINDEX INDEX "y"'),
(INDEX, INDEX),
('ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL', 'ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL'),
("-- built CONCURRENTLY by the index job\n" + INDEX, "-- built CONCURRENTLY by the index job\n" + INDEX),
),
ids=("concurrently", "unique_lowercase", "drop", "reindex", "plain_index", "alter", "word_in_comment"),
)
def test_only_concurrent_index_statements_change_so_the_migration_can_run_in_a_transaction(
migration_sql: str, expected: str
) -> None:
assert transactional(migration_sql) == expected
def test_a_clean_dry_run_lists_each_migration_with_the_time_the_upgrade_will_hold_its_locks() -> None:
text: Final = report(DryRun((Applied("20270101000000_a", 0.04), Applied("20270102000000_b", 12.345)), None))
assert text.startswith(
"2 pending migration(s) applied one at a time in a dry run, each rolled back before the next"
)
assert "\n 20270101000000_a: 0.0 s\n 20270102000000_b: 12.3 s" in text
def test_no_pending_migration_is_said_plainly() -> None:
assert report(DryRun((), None)) == "No pending migrations, this database is on this version's schema."
def test_a_rejected_migration_is_reported_with_postgres_own_words_and_what_ran_before_it() -> None:
error: Final = 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values'
text: Final = report(DryRun((Applied("20270101000000_a", 0.5),), Rejected("20270102000000_b", error, 3.25)))
assert text.startswith(
"Migration 20270102000000_b fails on this database after 3.2 s, nothing was applied.\n" + error
)
assert text.endswith("Applied before it in the dry run:\n 20270101000000_a: 0.5 s")
def test_a_rejection_on_the_first_migration_lists_nothing_before_it() -> None:
text: Final = report(DryRun((), Rejected("20270101000000_a", 'syntax error at or near "NOPE"', 0.0)))
assert "Applied before it" not in text
assert text.endswith('syntax error at or near "NOPE"')
def test_a_migration_that_needed_the_one_before_it_says_which_were_replayed_in_front() -> None:
text: Final = report(
DryRun((Applied("20270101000000_a", 0.1), Applied("20270102000000_b", 0.4, ("20270101000000_a",))), None)
)
assert text.endswith(
"\n 20270101000000_a: 0.1 s\n 20270102000000_b: 0.4 s (depends on 20270101000000_a, replayed in front of it)"
)
def test_a_database_without_history_that_matches_this_build_has_nothing_to_dry_run() -> None:
result: Final = DryRun((), None, Baseline(None))
assert result.ok is True
assert report(result) == (
"No migration history and the schema matches this build, the proxy baselines it at boot, nothing to dry-run."
)
def test_a_database_without_history_the_boot_refuses_to_baseline_reports_the_refusal() -> None:
result: Final = DryRun(
(), None, Baseline('the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"')
)
assert result.ok is False
assert report(result) == (
"No migration history and the proxy cannot baseline it at boot: "
'the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"'
)

37
tests/unit/llms/dashscope/test_dashscope_chat_transformation.py
View file

@ -5,15 +5,50 @@ These tests validate the DashScopeConfig class which extends OpenAIGPTConfig.
DashScope is an OpenAI-compatible provider with minor customizations.
"""
from typing import Final
import pytest
from litellm.types.llms.openai import AllMessageValues
import pytest
import litellm
from litellm import completion
from litellm.llms.dashscope.chat.transformation import DashScopeChatConfig
WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com"
def test_get_complete_url_derives_the_compatible_base_from_any_standard_base():
config: Final = DashScopeChatConfig()
for api_base in (
WORKSPACE,
f"{WORKSPACE}/compatible-mode/v1",
f"{WORKSPACE}/apps/anthropic",
f"{WORKSPACE}/api/v1",
):
url: Final = config.get_complete_url(
api_base=api_base,
api_key="sk-test",
model="qwen3.7-plus",
optional_params={},
litellm_params={},
)
assert url == f"{WORKSPACE}/compatible-mode/v1/chat/completions"
def test_get_complete_url_keeps_unrecognized_hosts():
url: Final = DashScopeChatConfig().get_complete_url(
api_base="https://gateway.example.com/v1",
api_key="sk-test",
model="qwen3.7-plus",
optional_params={},
litellm_params={},
)
assert url == "https://gateway.example.com/v1/chat/completions"
class TestDashScopeConfig:
"""Test class for DashScope functionality"""

171
tests/unit/llms/dashscope/test_dashscope_common_utils.py Normal file
View file

@ -0,0 +1,171 @@
"""
Unit tests for DashScope's protocol-specific base URL derivation.
"""
from typing import Final
import pytest
from litellm.llms.dashscope.common_utils import (
DASHSCOPE_DEFAULT_ORIGIN,
resolve_dashscope_anthropic_api_base,
resolve_dashscope_endpoints,
resolve_dashscope_family_rerank_api_base,
resolve_dashscope_openai_api_base,
)
WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com"
@pytest.mark.parametrize(
"api_base",
[
WORKSPACE,
f"{WORKSPACE}/",
f"{WORKSPACE}/compatible-mode/v1",
f"{WORKSPACE}/compatible-mode/v1/chat/completions",
f"{WORKSPACE}/apps/anthropic",
f"{WORKSPACE}/apps/anthropic/v1/messages",
f"{WORKSPACE}/api/v1",
f"{WORKSPACE}/api/v1/services/aigc/video-generation/video-synthesis",
],
)
def test_workspace_base_derives_every_protocol(api_base: str):
endpoints: Final = resolve_dashscope_endpoints(api_base)
assert endpoints is not None
assert endpoints.origin == WORKSPACE
assert endpoints.openai == f"{WORKSPACE}/compatible-mode/v1"
assert endpoints.anthropic == f"{WORKSPACE}/apps/anthropic"
assert endpoints.native == f"{WORKSPACE}/api/v1"
@pytest.mark.parametrize(
"origin",
[
"https://dashscope.aliyuncs.com",
"https://dashscope-intl.aliyuncs.com",
"https://dashscope-us.aliyuncs.com",
],
)
@pytest.mark.parametrize("path", ["", "/", "/compatible-mode/v1"])
def test_standard_hosts_derive_every_protocol(origin: str, path: str):
endpoints: Final = resolve_dashscope_endpoints(f"{origin}{path}")
assert endpoints is not None
assert endpoints.openai == f"{origin}/compatible-mode/v1"
assert endpoints.anthropic == f"{origin}/apps/anthropic"
assert endpoints.native == f"{origin}/api/v1"
@pytest.mark.parametrize(
"api_base",
[
None,
"",
"not-a-url",
"https://gateway.example.com",
"https://gateway.example.com/compatible-mode/v1",
f"{WORKSPACE}/custom-path",
f"{WORKSPACE}/compatible-mode/v1?api-version=1",
"https://trial.cn-beijing.maas.aliyuncs.com",
"https://token-plan.cn-beijing.maas.aliyuncs.com",
],
)
def test_unrecognized_bases_are_not_rewritten(api_base: str | None):
assert resolve_dashscope_endpoints(api_base) is None
def test_openai_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
assert resolve_dashscope_openai_api_base(WORKSPACE) == f"{WORKSPACE}/compatible-mode/v1"
assert resolve_dashscope_openai_api_base(f"{WORKSPACE}/apps/anthropic") == f"{WORKSPACE}/compatible-mode/v1"
assert (
resolve_dashscope_openai_api_base(f"{WORKSPACE}/compatible-api/v1/reranks") == f"{WORKSPACE}/compatible-mode/v1"
)
@pytest.mark.parametrize(
"endpoint_url",
[
"https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
f"{WORKSPACE}/api/v1/services/aigc/video-generation/video-synthesis",
],
)
def test_openai_base_keeps_a_native_endpoint_url(monkeypatch: pytest.MonkeyPatch, endpoint_url: str):
monkeypatch.setenv("DASHSCOPE_API_BASE", endpoint_url)
assert resolve_dashscope_openai_api_base(None) == endpoint_url
assert resolve_dashscope_openai_api_base(endpoint_url) == endpoint_url
assert resolve_dashscope_openai_api_base(f"{WORKSPACE}/api/v1") == f"{WORKSPACE}/compatible-mode/v1"
def test_openai_base_keeps_unrecognized_hosts(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
assert resolve_dashscope_openai_api_base("https://gateway.example.com/v1") == "https://gateway.example.com/v1"
def test_openai_base_defaults_to_dashscope(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
assert resolve_dashscope_openai_api_base(None) == f"{DASHSCOPE_DEFAULT_ORIGIN}/compatible-mode/v1"
def test_anthropic_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
assert resolve_dashscope_anthropic_api_base(WORKSPACE) == f"{WORKSPACE}/apps/anthropic"
assert resolve_dashscope_anthropic_api_base(f"{WORKSPACE}/compatible-mode/v1") == f"{WORKSPACE}/apps/anthropic"
def test_anthropic_base_keeps_unrecognized_hosts(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
assert (
resolve_dashscope_anthropic_api_base("https://gateway.example.com/anthropic")
== "https://gateway.example.com/anthropic"
)
def test_rerank_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
assert (
resolve_dashscope_family_rerank_api_base(WORKSPACE, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks")
== f"{WORKSPACE}/compatible-api/v1/reranks"
)
def test_rerank_base_keeps_unrecognized_hosts():
gateway: Final = "https://gateway.example.com/v1"
assert (
resolve_dashscope_family_rerank_api_base(gateway, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks")
== gateway
)
def test_rerank_base_uses_the_configured_default(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
assert (
resolve_dashscope_family_rerank_api_base(None, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks")
== "https://default/reranks"
)
@pytest.mark.parametrize(
"origin",
["https://cn-hongkong.dashscope.aliyuncs.com", "https://token-plan.cn-beijing.maas.aliyuncs.com"],
)
def test_rerank_keeps_remapping_chat_bases_on_other_aliyun_hosts(monkeypatch: pytest.MonkeyPatch, origin: str):
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
assert (
resolve_dashscope_family_rerank_api_base(
f"{origin}/compatible-mode/v1", "DASHSCOPE_API_BASE_RERANK", "https://default/reranks"
)
== f"{origin}/compatible-api/v1/reranks"
)

35
tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py
View file

@ -3,6 +3,7 @@ Unit tests for DashScope embedding transformation.
"""
import json
from typing import Final
from unittest.mock import MagicMock
import httpx
@ -16,6 +17,40 @@ from litellm.llms.dashscope.embed.transformation import (
)
from litellm.types.utils import EmbeddingResponse
WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com"
def test_get_complete_url_derives_the_compatible_base_from_any_standard_base():
config: Final = DashScopeEmbeddingConfig()
for api_base in (
WORKSPACE,
f"{WORKSPACE}/compatible-mode/v1",
f"{WORKSPACE}/apps/anthropic",
f"{WORKSPACE}/api/v1",
):
url: Final = config.get_complete_url(
api_base=api_base,
api_key="sk-test",
model="text-embedding-v4",
optional_params={},
litellm_params={},
)
assert url == f"{WORKSPACE}/compatible-mode/v1/embeddings"
def test_get_complete_url_keeps_unrecognized_hosts():
url: Final = DashScopeEmbeddingConfig().get_complete_url(
api_base="https://gateway.example.com/v1",
api_key="sk-test",
model="text-embedding-v4",
optional_params={},
litellm_params={},
)
assert url == "https://gateway.example.com/v1/embeddings"
def test_validate_environment_and_url():
config = DashScopeEmbeddingConfig()

215
tests/unit/llms/dashscope/test_dashscope_messages_transformation.py Normal file
View file

@ -0,0 +1,215 @@
"""
Unit tests for DashScope's Anthropic-compatible Messages endpoint.
"""
from typing import Final
import pytest
import respx
import litellm
from litellm.exceptions import AuthenticationError
from litellm.llms.dashscope.messages.transformation import (
DashScopeAnthropicMessagesConfig,
)
WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com"
INTL: Final = "https://dashscope-intl.aliyuncs.com"
CODING_PLAN: Final = "https://coding-intl.dashscope.aliyuncs.com"
CHAT_COMPLETION: Final = {
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1792227137,
"model": "qwen3.7-plus",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10},
}
MESSAGE: Final = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "qwen3.7-plus",
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 9, "output_tokens": 1},
}
@pytest.mark.parametrize(
("api_base", "expected"),
[
(
None,
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
),
(
"https://dashscope.aliyuncs.com",
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
),
(
"https://dashscope.aliyuncs.com/compatible-mode/v1",
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
),
(
"https://dashscope.aliyuncs.com/apps/anthropic",
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
),
(
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
"https://dashscope.aliyuncs.com/apps/anthropic/v1/messages",
),
(
WORKSPACE,
f"{WORKSPACE}/apps/anthropic/v1/messages",
),
(
f"{WORKSPACE}/compatible-mode/v1",
f"{WORKSPACE}/apps/anthropic/v1/messages",
),
(
"https://gateway.example.com/anthropic",
"https://gateway.example.com/anthropic/v1/messages",
),
],
)
def test_get_complete_url(api_base: str | None, expected: str):
config: Final = DashScopeAnthropicMessagesConfig()
url: Final = config.get_complete_url(
api_base=api_base,
api_key="sk-test",
model="qwen3.7-plus",
optional_params={},
litellm_params={},
)
assert url == expected
@pytest.mark.parametrize("dashscope_api_key", [None, " "])
def test_missing_dashscope_key_does_not_fall_back_to_anthropic(
monkeypatch: pytest.MonkeyPatch, dashscope_api_key: str | None
):
monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
if dashscope_api_key is not None:
monkeypatch.setenv("DASHSCOPE_API_KEY", dashscope_api_key)
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-secret")
monkeypatch.setattr(litellm, "api_key", None)
headers: dict[str, str] = {}
with pytest.raises(AuthenticationError) as error:
DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers=headers,
model="qwen3.7-plus",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
api_base=WORKSPACE,
)
assert "DashScope" in str(error.value)
assert "sk-ant-secret" not in str(error.value)
assert {name.lower() for name in headers}.isdisjoint({"x-api-key", "authorization"})
def test_caller_auth_header_is_honored_without_a_dashscope_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
monkeypatch.setattr(litellm, "api_key", None)
headers: Final = {"x-api-key": "caller-key"}
resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers=headers,
model="qwen3.7-plus",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
api_base=WORKSPACE,
)
assert resolved_headers["x-api-key"] == "caller-key"
def test_caller_authorization_header_is_honored_without_a_dashscope_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
monkeypatch.setattr(litellm, "api_key", None)
headers: Final = {"Authorization": "Bearer caller-token"}
resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers=headers,
model="qwen3.7-plus",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
api_base=WORKSPACE,
)
assert resolved_headers["Authorization"] == "Bearer caller-token"
def test_dashscope_key_wins_over_an_anthropic_environment_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope")
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-secret")
resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers={},
model="qwen3.7-plus",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
api_base=WORKSPACE,
)
header_values: Final = " ".join(resolved_headers.values())
assert "sk-dashscope" in header_values
assert "sk-ant-secret" not in header_values
def test_dashscope_prefix_reaches_the_provider():
model, provider, _, _ = litellm.get_llm_provider(
"dashscope/qwen3.7-plus",
api_base=WORKSPACE,
)
assert provider == "dashscope"
assert model == "qwen3.7-plus"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("api_base", "expected_url"),
[
(INTL, f"{INTL}/apps/anthropic/v1/messages"),
(f"{INTL}/apps/anthropic", f"{INTL}/apps/anthropic/v1/messages"),
(WORKSPACE, f"{WORKSPACE}/apps/anthropic/v1/messages"),
(f"{CODING_PLAN}/apps/anthropic", f"{CODING_PLAN}/apps/anthropic/v1/messages"),
(None, "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"),
(f"{INTL}/compatible-mode/v1", f"{INTL}/compatible-mode/v1/chat/completions"),
(f"{CODING_PLAN}/v1", f"{CODING_PLAN}/v1/chat/completions"),
("https://gateway.example.com/anthropic", "https://gateway.example.com/anthropic/chat/completions"),
],
)
async def test_only_a_domain_root_or_anthropic_base_leaves_the_chat_bridge(
monkeypatch: pytest.MonkeyPatch, api_base: str | None, expected_url: str
):
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock(assert_all_called=False) as upstream:
upstream.route(path__regex=r".*/chat/completions$").respond(200, json=CHAT_COMPLETION)
upstream.route(path__regex=r".*/v1/messages$").respond(200, json=MESSAGE)
await litellm.anthropic.messages.acreate(
model="dashscope/qwen3.7-plus",
messages=[{"role": "user", "content": "say ok"}],
max_tokens=16,
api_key="sk-dashscope",
api_base=api_base,
)
assert [str(call.request.url) for call in upstream.calls] == [expected_url]

52
tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py
View file

@ -3,12 +3,15 @@ Unit tests for DashScope rerank transformation.
"""
import json
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import respx
from pydantic import ValidationError
import litellm
from litellm.llms.dashscope.common_utils import DashScopeError
from litellm.llms.dashscope.rerank.transformation import (
@ -26,6 +29,21 @@ class TestDashScopeRerankURL:
url = self.config.get_complete_url(api_base=None, model="qwen3-rerank")
assert url == DEFAULT_RERANK_URL
def test_domain_root_remaps_to_rerank_route(self):
url: Final = self.config.get_complete_url(
api_base="https://llm-abc123.cn-beijing.maas.aliyuncs.com",
model="qwen3-rerank",
)
assert url == "https://llm-abc123.cn-beijing.maas.aliyuncs.com/compatible-api/v1/reranks"
def test_explicit_rerank_url_wins_over_the_rerank_env(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("DASHSCOPE_API_BASE_RERANK", DEFAULT_RERANK_URL)
explicit_url: Final = "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks"
url: Final = self.config.get_complete_url(api_base=explicit_url, model="qwen3-rerank")
assert url == explicit_url
def test_chat_shaped_base_remaps_to_rerank_route(self):
url = self.config.get_complete_url(
api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
@ -434,3 +452,37 @@ def test_transform_rerank_response_shape_errors_do_not_echo_the_payload():
_transform({"results": [["leaked document text"]]})
assert "leaked document text" not in str(exc_info.value)
@pytest.mark.parametrize(
("provider", "rerank_env"),
[
("dashscope", "DASHSCOPE_API_BASE_RERANK"),
("qwencloud", "QWENCLOUD_API_BASE_RERANK"),
("qwen_ai_platform", "QWEN_AI_PLATFORM_API_BASE_RERANK"),
],
)
@pytest.mark.parametrize(
"api_base",
[
"https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks",
"https://dashscope-intl.aliyuncs.com/compatible-api/v1",
"https://dashscope-intl.aliyuncs.com",
"https://dashscope-intl.aliyuncs.com/api/v1",
],
)
def test_rerank_sends_an_explicit_base_over_the_rerank_env(
monkeypatch: pytest.MonkeyPatch, provider: str, rerank_env: str, api_base: str
):
monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com")
monkeypatch.setenv(rerank_env, DEFAULT_RERANK_URL)
with respx.mock(assert_all_called=True) as upstream:
upstream.post("https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks").respond(
200, json={"id": "r1", "results": [{"index": 0, "relevance_score": 0.9}], "usage": {"total_tokens": 5}}
)
response: Final = litellm.rerank(
model=f"{provider}/qwen3-rerank", query="q", documents=["a"], api_base=api_base, api_key="sk-test"
)
assert response.results == [{"index": 0, "relevance_score": 0.9}]

14
tests/unit/llms/dashscope/test_dashscope_video_transformation.py
View file

@ -1119,6 +1119,20 @@ class TestDashScopeVideoEnvironment:
assert url == "https://ws-123.cn-beijing.maas.aliyuncs.com"
def test_every_protocol_base_normalizes_to_the_host(self):
workspace: Final = "https://ws-123.cn-beijing.maas.aliyuncs.com"
for api_base in (
f"{workspace}/compatible-mode/v1",
f"{workspace}/apps/anthropic",
f"{workspace}/api/v1",
):
url: Final = DashScopeVideoConfig().get_complete_url(
model="wan3.0-video", api_base=api_base, litellm_params={}
)
assert url == workspace
def test_default_host_is_beijing(self, monkeypatch):
monkeypatch.delenv("DASHSCOPE_API_BASE_VIDEO", raising=False)
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)

46
tests/unit/llms/dashscope/test_qwen_brand_aliases.py
View file

@ -2,8 +2,10 @@ import json
import math
import re
from pathlib import Path
from typing import Final, Literal
import pytest
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm import completion, get_llm_provider
@ -53,6 +55,23 @@ DASHSCOPE_FAMILY_ENV_VARS = [
"QWEN_AI_PLATFORM_API_BASE_IMAGE",
]
class BrandCase(TypedDict):
provider: ReadOnly[str]
display_name: ReadOnly[str]
enum: ReadOnly[LlmProviders]
key_env: ReadOnly[str]
base_env: ReadOnly[str]
default_base: ReadOnly[str]
default_rerank_base: ReadOnly[str]
default_image_base: ReadOnly[str]
chat_config: ReadOnly[type[DashScopeChatConfig]]
embedding_config: ReadOnly[type[DashScopeEmbeddingConfig]]
rerank_config: ReadOnly[type[DashScopeRerankConfig]]
image_config: ReadOnly[type[DashScopeImageGenerationConfig]]
BRAND_CASES = [
pytest.param(
{
@ -211,6 +230,20 @@ class TestQwenBrandDefaultUrls:
)
assert url == f"{brand['default_base']}/embeddings"
@pytest.mark.parametrize("brand", BRAND_CASES)
@pytest.mark.parametrize(("config_key", "route"), [("chat_config", "chat/completions"), ("embedding_config", "embeddings")])
def test_openai_routes_derive_from_a_domain_root(
self, brand: BrandCase, config_key: Literal["chat_config", "embedding_config"], route: str
) -> None:
url: Final = brand[config_key]().get_complete_url(
api_base="https://dashscope-intl.aliyuncs.com",
api_key="sk-test",
model="qwen-max",
optional_params={},
litellm_params={},
)
assert url == f"https://dashscope-intl.aliyuncs.com/compatible-mode/v1/{route}"
@pytest.mark.parametrize("brand", BRAND_CASES)
def test_rerank_complete_url(self, brand):
url = brand["rerank_config"]().get_complete_url(api_base=None, model="gte-rerank-v2")
@ -249,6 +282,19 @@ class TestQwenBrandDefaultUrls:
)
assert url == brand["default_image_base"]
@pytest.mark.parametrize("brand", BRAND_CASES)
def test_image_generation_uses_its_own_configured_base(self, monkeypatch: pytest.MonkeyPatch, brand: BrandCase) -> None:
monkeypatch.setenv("DASHSCOPE_API_BASE", "https://legacy.example.com/images")
monkeypatch.setenv(brand["base_env"], "https://gateway.example.com/qwen/images")
url: Final = brand["image_config"]().get_complete_url(
api_base=None,
api_key="sk-test",
model="qwen-image",
optional_params={},
litellm_params={},
)
assert url == "https://gateway.example.com/qwen/images"
@pytest.mark.parametrize("brand", BRAND_CASES)
def test_validate_environment_requires_key(self, brand):
with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"):

223
tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py
View file

@ -2,18 +2,25 @@
Unit tests for OpenAI Text Completion Guardrail Translation Handler
"""
from typing import List, Optional, Tuple
import json
from typing import List, Literal, Optional, Tuple
from unittest.mock import MagicMock
import pytest
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.openai.completion.guardrail_translation.handler import (
OpenAITextCompletionHandler,
)
from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage
class MockGuardrail(CustomGuardrail):
@ -341,3 +348,215 @@ class TestPIIMaskingScenario:
assert "[EMAIL_REDACTED]" in result["prompt"][1]
assert "alice@example.com" not in result["prompt"][0]
assert "bob@test.org" not in result["prompt"][1]
class _RecordingKeywordFilter(ContentFilterGuardrail):
def __init__(self):
super().__init__(
guardrail_name="keyword-filter",
blocked_words=[
BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK),
BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK),
],
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
self.streaming_sampling_rate = 1
self.scanned_texts: tuple[tuple[str, ...], ...] = ()
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
self.scanned_texts = (*self.scanned_texts, tuple(inputs.get("texts") or ()))
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
class _PassthroughReplaceGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
if any("kumquat" in text for text in inputs.get("texts") or []):
self.raise_passthrough_exception(violation_message="replaced kumquat", request_data=request_data)
return inputs
def _text_chunk(index: int | None, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-stream",
created=1589478378,
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)],
)
def _usage_chunk() -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-stream",
created=1589478378,
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[],
usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8),
)
def _delivered(chunks: list[TextCompletionResponse]) -> list[tuple[int, str | None, str | None]]:
return [(choice.index, choice.text, choice.finish_reason) for chunk in chunks for choice in chunk.choices]
async def _drive_flat_stream(guardrail: CustomGuardrail, chunks: list[TextCompletionResponse]) -> list[object]:
async def upstream():
for chunk in chunks:
yield chunk
return [
item
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/completions"),
response=upstream(),
request_data={"guardrail_to_apply": guardrail, "model": "gpt-3.5-turbo-instruct"},
)
]
class TestStreamingOutputProcessing:
@pytest.mark.asyncio
async def test_ended_stream_delivers_each_choice_masked_text(self):
guardrail = _RecordingKeywordFilter()
chunks = [
_text_chunk(0, "A pers"),
_text_chunk(1, "kiwi "),
_text_chunk(0, "im"),
_text_chunk(1, "pie", "stop"),
_text_chunk(0, "mon tart"),
_text_chunk(0, None, "length"),
_usage_chunk(),
]
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
request_data={},
deliver_ended_stream_rewrites=True,
)
assert guardrail.scanned_texts == (("A persimmon tart", "kiwi pie"),)
assert _delivered(chunks) == [
(0, "A [KEYWORD_REDACTED] tart", None),
(1, "kiwi ", None),
(0, "", None),
(1, "pie", "stop"),
(0, "", None),
(0, None, "length"),
]
assert chunks[-1].usage.total_tokens == 8
@pytest.mark.asyncio
async def test_ended_stream_scanned_without_delivery_keeps_chunks_and_stores_the_assembled_response(self):
chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
request_data: dict[str, object] = {}
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), request_data=request_data
)
assert _delivered(chunks) == [(0, "A pers", None), (0, "immon tart", None), (0, "", "stop")]
assembled = request_data["response"]
assert isinstance(assembled, TextCompletionResponse)
assert [(c.text, c.finish_reason) for c in assembled.choices] == [("A [KEYWORD_REDACTED] tart", "stop")]
assert assembled.usage.total_tokens == 8
@pytest.mark.asyncio
async def test_ended_stream_block_raises_the_guardrail_error(self):
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")]
with pytest.raises(HTTPException) as exc_info:
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True
)
assert exc_info.value.status_code == 400
assert "kumquat" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_ended_stream_treats_a_null_choice_index_as_choice_zero(self):
guardrail = _RecordingKeywordFilter()
chunks = [_text_chunk(None, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop")]
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=guardrail, request_data={}, deliver_ended_stream_rewrites=True
)
assert guardrail.scanned_texts == (("A persimmon tart",),)
assert [choice.text for chunk in chunks for choice in chunk.choices] == ["A [KEYWORD_REDACTED] tart", "", ""]
@pytest.mark.asyncio
async def test_ended_stream_blocks_a_keyword_split_across_a_null_and_a_zero_choice_index(self):
chunks = [_text_chunk(None, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")]
with pytest.raises(HTTPException) as exc_info:
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True
)
assert exc_info.value.status_code == 400
assert "kumquat" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_flat_stream_is_scanned_once_after_it_ends(self):
guardrail = _RecordingKeywordFilter()
chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
delivered = await _drive_flat_stream(guardrail, chunks)
assert guardrail.scanned_texts == (("A persimmon tart",),)
assert len(delivered) == len(chunks)
@pytest.mark.asyncio
async def test_flat_stream_block_ends_with_an_in_band_error_frame(self):
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
delivered = await _drive_flat_stream(_RecordingKeywordFilter(), chunks)
error_frame = delivered[-1]
assert isinstance(error_frame, bytes)
assert error_frame.startswith(b'data: {"error": ')
assert b"kumquat" in error_frame
@pytest.mark.asyncio
async def test_flat_stream_passthrough_block_appends_the_replacement_and_ends_with_content_filter(self):
guardrail = _PassthroughReplaceGuardrail(
guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True
)
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
delivered = await _drive_flat_stream(guardrail, chunks)
frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)]
assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [
("replaced kumquat", None),
("", "content_filter"),
]
assert frames[-1]["usage"]["total_tokens"] == 8
@pytest.mark.asyncio
async def test_flat_stream_passthrough_block_reports_no_usage_the_upstream_never_sent(self):
guardrail = _PassthroughReplaceGuardrail(
guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True
)
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")]
delivered = await _drive_flat_stream(guardrail, [chunk.model_copy(update={"usage": None}) for chunk in chunks])
frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)]
assert frames[-1]["choices"][0]["finish_reason"] == "content_filter"
assert [frame for frame in frames if "usage" in frame] == []

13
tests/unit/proxy/common_utils/test_callback_utils.py
View file

@ -49,6 +49,19 @@ def test_get_remaining_tokens_and_requests_from_request_data():
}
@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}])
def test_remaining_limit_headers_ignore_a_model_group_that_is_not_a_string(client_model_group: object):
casedata: Final = {
"metadata": {
"model_group": client_model_group,
f"litellm-key-remaining-requests-{client_model_group}": 100,
f"litellm-key-remaining-tokens-{client_model_group}": 200,
}
}
assert get_remaining_tokens_and_requests_from_request_data(casedata) == {}
@patch(
"litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars",
return_value=["API_KEY", "MISSING_VAR"],

16
tests/unit/proxy/db/test_prisma_client.py
View file

@ -586,3 +586,19 @@ async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain
await second
await asyncio.wait_for(drained, timeout=1)
assert engine.queries_finished == 2
def test_migration_dry_run_summary_without_migrations_has_nothing_to_dry_run(unset_database_url):
"""`prisma db push` applies no migration files, so the dry run reports that instead of connecting."""
from litellm.proxy.db.prisma_client import PrismaManager
assert PrismaManager.migration_dry_run_summary(use_migrate=False) == (
True,
"The schema is pushed rather than migrated, there is nothing to dry-run.",
)
def test_migration_dry_run_summary_needs_database_url(unset_database_url):
from litellm.proxy.db.prisma_client import PrismaManager
assert PrismaManager.migration_dry_run_summary(use_migrate=True) == (False, "DATABASE_URL is not set.")

190
tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
View file

@ -4,6 +4,7 @@ Tests for the Content Filter Guardrail
import json
import os
from collections.abc import AsyncGenerator, AsyncIterator, Callable
from typing import Final
from unittest.mock import MagicMock
@ -30,7 +31,17 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor
ContentFilterCategoryConfig,
ContentFilterDetection,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents
from litellm.types.utils import (
Delta,
ModelResponseStream,
StandardLoggingGuardrailInformation,
StreamingChoices,
TextChoices,
TextCompletionResponse,
Usage,
)
class TestContentFilterGuardrail:
@ -3523,3 +3534,180 @@ async def test_image_whose_description_is_clean_passes_through_unchanged():
)
assert result == {"images": ["https://example.com/desk.png"], "texts": []}
def _fruit_word_filter() -> ContentFilterGuardrail:
return ContentFilterGuardrail(
guardrail_name="fruit-filter",
blocked_words=[
BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK),
BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK),
],
event_hook=GuardrailEventHooks.post_call,
)
async def _text_completion_stream(*parts: tuple[str | None, str | None]):
for text, finish_reason in parts:
yield TextCompletionResponse(
id="cmpl-stream",
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)],
)
@pytest.mark.asyncio
async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks():
request_data: Final[dict] = {}
delivered: Final = [
chunk
async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(),
response=_text_completion_stream(("A pers", None), ("immon tart", None), (None, "stop")),
request_data=request_data,
)
]
assert "".join(chunk.choices[0].text or "" for chunk in delivered) == "A [KEYWORD_REDACTED] tart"
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success"
_LONG_TAIL: Final = "the quick brown fox jumps over the lazy dog and then eats a ripe persimmon"
async def _text_completion_n2_stream_without_choice_1_finish() -> AsyncGenerator[TextCompletionResponse, None]:
for index, text, finish_reason in (
(0, "A short answer", None),
(1, _LONG_TAIL[:30], None),
(1, _LONG_TAIL[30:], None),
(0, None, "stop"),
):
yield TextCompletionResponse(
id="cmpl-stream",
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)],
).model_copy(update={"usage": None})
async def _chat_n2_stream_without_choice_1_finish() -> AsyncGenerator[ModelResponseStream, None]:
for index, content, finish_reason in (
(0, "A short answer", None),
(1, _LONG_TAIL[:30], None),
(1, _LONG_TAIL[30:], None),
(0, None, "stop"),
):
yield ModelResponseStream(
id="chatcmpl-stream",
model="gpt-4.1-mini",
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)],
)
def _streamed_text(chunk: ModelResponseStream | TextCompletionResponse, index: int) -> str:
return "".join(_choice_text(choice) for choice in chunk.choices if choice.index == index)
def _choice_text(choice: TextChoices | StreamingChoices) -> str:
content: Final = stream_item_field(choice, "text") if isinstance(choice, TextChoices) else choice.delta.content
return content if isinstance(content, str) else ""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"stream", [_text_completion_n2_stream_without_choice_1_finish, _chat_n2_stream_without_choice_1_finish]
)
async def test_streaming_hook_releases_the_held_tail_of_a_choice_the_stream_never_finished(
stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]],
) -> None:
delivered: Final = [
chunk
async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(), response=stream(), request_data={}
)
]
assert "".join(_streamed_text(chunk, 0) for chunk in delivered) == "A short answer"
assert "".join(_streamed_text(chunk, 1) for chunk in delivered) == _LONG_TAIL.replace(
"persimmon", "[KEYWORD_REDACTED]"
)
assert all(getattr(chunk, "usage", None) is None for chunk in delivered)
_STREAM_USAGE: Final = Usage(prompt_tokens=5, completion_tokens=20, total_tokens=25)
async def _then_usage_chunk(
stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]],
usage_chunk: ModelResponseStream | TextCompletionResponse,
) -> AsyncGenerator[ModelResponseStream | TextCompletionResponse, None]:
async for chunk in stream():
yield chunk
yield usage_chunk
@pytest.mark.asyncio
@pytest.mark.parametrize(
("stream", "usage_chunk"),
[
(
_text_completion_n2_stream_without_choice_1_finish,
TextCompletionResponse(id="cmpl-stream", object="text_completion", choices=[], usage=_STREAM_USAGE),
),
(_chat_n2_stream_without_choice_1_finish, ModelResponseStream(id="chatcmpl-stream", choices=[], usage=_STREAM_USAGE)),
],
)
async def test_streaming_hook_keeps_the_usage_chunk_last_after_releasing_a_held_tail(
stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]],
usage_chunk: ModelResponseStream | TextCompletionResponse,
) -> None:
delivered: Final = [
chunk
async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(), response=_then_usage_chunk(stream, usage_chunk), request_data={}
)
]
assert delivered[-1] is usage_chunk
assert "".join(_streamed_text(chunk, 1) for chunk in delivered[:-1]) == _LONG_TAIL.replace(
"persimmon", "[KEYWORD_REDACTED]"
)
@pytest.mark.asyncio
async def test_streaming_hook_relays_each_responses_api_event_as_it_arrives() -> None:
timeline: Final[list[str]] = []
async def responses_api_events() -> AsyncGenerator[OutputTextDeltaEvent, None]:
for delta in ("first", "second"):
timeline.append(f"upstream sent {delta}")
yield OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_1",
output_index=0,
content_index=0,
delta=delta,
)
async for event in _fruit_word_filter().async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(), response=responses_api_events(), request_data={}
):
timeline.append(f"client got {event.delta}")
assert timeline == ["upstream sent first", "client got first", "upstream sent second", "client got second"]
@pytest.mark.asyncio
async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks():
with pytest.raises(HTTPException) as exc_info:
async for _ in _fruit_word_filter().async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(),
response=_text_completion_stream(("A kum", None), ("quat tart", None), (None, "stop")),
request_data={},
):
pass
assert exc_info.value.status_code == 400
assert "kumquat" in str(exc_info.value.detail)

28
tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py
View file

@ -5,6 +5,7 @@ import pytest
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.openai.completion.guardrail_translation.handler import OpenAITextCompletionHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.grayswan import grayswan as grayswan_module
from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
GraySwanGuardrailAPIError,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import TextChoices, TextCompletionResponse
@pytest.fixture
@ -884,3 +886,29 @@ async def test_pre_call_payload_unchanged() -> None:
{"role": "user", "content": "second"},
]
assert "tools" not in payload
@pytest.mark.asyncio
async def test_post_call_never_forwards_the_proxy_key_metadata_of_the_request_metadata() -> None:
guardrail = _post_call_guardrail()
client = _CapturingClient()
guardrail.async_handler = client
await OpenAITextCompletionHandler().process_output_response(
response=TextCompletionResponse(
id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")]
),
guardrail_to_apply=guardrail,
user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key", user_email="dev@example.com"),
request_data={
"model": "gpt-4.1-mini",
"prompt": "say hello",
"metadata": {
"user_api_key_hash": "hashed-key",
"user_api_key_user_email": "dev@example.com",
"user_api_key_request_route": "/v1/completions",
},
},
)
assert "litellm_metadata" not in client.calls[0]["json"]

124
tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -75,11 +75,13 @@ class _RecordingLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.standard_logging_object: StandardLoggingPayload | None = None
self.logged: asyncio.Event = asyncio.Event()
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.standard_logging_object = cast(StandardLoggingPayload, kwargs["standard_logging_object"])
self.logged.set()
class PostCallGuardrail(CustomGuardrail):
@ -239,52 +241,33 @@ class TestHasPostCallGuardrailsForPassthrough:
@pytest.mark.asyncio
async def test_deferred_flag_stores_and_executes_closure():
"""
When defer_async_logging is True on logging_obj:
1. wrapper_async stores a callable closure instead of calling create_task
2. Calling the closure fires create_task
3. Sync callbacks fire immediately (not deferred)
"""
mock_logging_obj = MagicMock()
mock_logging_obj.defer_async_logging = True
mock_logging_obj.enqueue_deferred_logging = None
await litellm.acompletion(
model="gpt-3.5-turbo",
@pytest.mark.parametrize("deferred", (False, True))
async def test_success_logging_waits_for_deferred_closure_when_requested(deferred: bool) -> None:
recorder: Final = _RecordingLogger()
logging_obj: Final = Logging(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
start_time=datetime(2026, 1, 1),
litellm_call_id="deferred-success",
function_id="deferred-success",
dynamic_async_success_callbacks=[recorder],
)
logging_obj.defer_async_logging = deferred
response: Final = await litellm.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
mock_response="Hello!",
litellm_logging_obj=mock_logging_obj,
litellm_logging_obj=logging_obj,
)
# Closure was stored
enqueue_fn = mock_logging_obj.enqueue_deferred_logging
assert callable(enqueue_fn), "Closure should be stored on logging_obj"
# Sync callbacks fired immediately
mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
# Calling the closure fires create_task
created_tasks = []
real_create_task = asyncio.create_task
def tracking_create_task(coro):
task = real_create_task(coro)
created_tasks.append(task)
return task
with patch("asyncio.create_task", side_effect=tracking_create_task):
enqueue_fn()
assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task"
for task in created_tasks:
if not task.done():
task.cancel()
try:
await task
except (asyncio.CancelledError, Exception):
pass
if deferred:
assert recorder.standard_logging_object is None
assert logging_obj.enqueue_deferred_logging is not None
logging_obj.enqueue_deferred_logging()
await asyncio.wait_for(recorder.logged.wait(), timeout=5)
assert recorder.standard_logging_object is not None
assert recorder.standard_logging_object["id"] == response.id
@pytest.mark.asyncio
@ -401,40 +384,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the
assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6)
# ---------------------------------------------------------------------------
# 3. Non-streaming regression: without flag, create_task fires normally
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_no_flag_fires_create_task_normally():
"""Without defer_async_logging, wrapper_async calls create_task as before."""
created_tasks = []
real_create_task = asyncio.create_task
def tracking_create_task(coro):
task = real_create_task(coro)
created_tasks.append(task)
return task
with patch("asyncio.create_task", side_effect=tracking_create_task):
await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
mock_response="Hello!",
)
assert len(created_tasks) >= 1
for task in created_tasks:
if not task.done():
task.cancel()
try:
await task
except (asyncio.CancelledError, Exception):
pass
# ---------------------------------------------------------------------------
# 4. Non-streaming: deferred success log is suppressed when guardrail raises
# ---------------------------------------------------------------------------
@ -1537,6 +1486,27 @@ class TestArmDeferredStreamDispatch:
assert recorded["cache_hit"] is False
assert recorded["prefer_async_handlers"] is True
@pytest.mark.asyncio
async def test_text_completion_stream_logs_after_the_guardrails_finish(self):
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
logging_obj, recorded = self._dispatch_recording_logging_obj()
self._processor()._arm_deferred_stream_dispatch(
response=TextCompletionStreamWrapper(completion_stream=MagicMock(), model="gpt-test"),
route_type="atext_completion",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
@pytest.mark.asyncio
async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self):
"""The router wraps iterators without _hidden_params in

143
tests/unit/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -2551,3 +2551,146 @@ async def test_get_daily_activity_rejects_non_canonical_dates_before_querying(st
assert error.value.detail == {"error": "start_date and end_date must be valid YYYY-MM-DD dates"}
mock_table.count.assert_not_awaited()
mock_table.find_many.assert_not_awaited()
@pytest.mark.parametrize("missing", [None, ""])
def test_missing_dimensions_preserve_all_paginated_metrics(missing: str | None) -> None:
records: Final = (
_daily_user_spend_record(user_id="user", api_key="", spend=-0.25, model=missing, model_group=missing),
_daily_user_spend_record(user_id="user", api_key="", spend=0.0, model=missing, model_group=missing),
)
result: Final = common_daily_activity_module._aggregate_spend_records_sync(
records=records, api_key_metadata={}, entity_id_field="user_id", entity_metadata_field=None
)
day: Final = result["results"][0]
assert day.metrics == result["totals"]
for bucket in (
day.breakdown.models["unknown-model"],
day.breakdown.model_groups["unknown-model-group"],
day.breakdown.providers["openai"],
day.breakdown.endpoints["/chat/completions"],
day.breakdown.entities["user"],
):
assert bucket.metrics == day.metrics
assert bucket.api_key_breakdown == day.breakdown.api_keys
key: Final = day.breakdown.api_keys["unattributed"]
assert key.metrics == day.metrics
assert key.metadata.model_dump(exclude_none=True) == {"key_alias": "Unattributed", "key_exists": False}
def test_normalized_rollups_add_colliding_model_metrics_and_keep_key_identities() -> None:
records: Final = (
_grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model=None, spend=1.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model="", spend=2.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model="unknown-model", spend=3.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_MODEL_API_KEY, model=None, api_key="", spend=1.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_MODEL_API_KEY, model="", api_key="", spend=2.0),
_grouping_row(
common_daily_activity_module._GROUP_DATE_MODEL_API_KEY,
model="unknown-model",
api_key="unattributed",
spend=3.0,
),
_grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="", spend=3.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="unattributed", spend=3.0),
_grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="raw:unattributed", spend=4.0),
)
result: Final = common_daily_activity_module._aggregate_grouping_sets_records_sync(
records=records, api_key_metadata={}
)
breakdown: Final = result["results"][0].breakdown
assert breakdown.models["unknown-model"].metrics.spend == 6.0
assert {key: value.metrics.spend for key, value in breakdown.models["unknown-model"].api_key_breakdown.items()} == {
"unattributed": 3.0,
"raw:unattributed": 3.0,
}
assert {key: value.metrics.spend for key, value in breakdown.api_keys.items()} == {
"unattributed": 3.0,
"raw:unattributed": 3.0,
"raw:raw:unattributed": 4.0,
}
@pytest.mark.parametrize(
("model", "group", "expected_group"),
[
(None, None, "unknown-model-group"),
("", "", "unknown-model-group"),
(None, "configured-group", "configured-group"),
("provider-model", "", "provider-model"),
("unknown-model", "", "unknown-model"),
(" ", "", " "),
],
)
def test_model_group_uses_raw_model_before_unknown_group(
model: str | None, group: str | None, expected_group: str
) -> None:
result: Final = common_daily_activity_module._aggregate_spend_records_sync(
records=(_daily_user_spend_record(user_id="user", api_key="key", spend=1.0, model=model, model_group=group),),
api_key_metadata={},
entity_id_field=None,
entity_metadata_field=None,
)
assert tuple(result["results"][0].breakdown.model_groups) == (expected_group,)
@pytest.mark.parametrize(
("raw_key", "bucket"),
[
("", "unattributed"),
("real-hash", "real-hash"),
("unattributed", "raw:unattributed"),
("raw:unattributed", "raw:raw:unattributed"),
("raw:real-hash", "raw:raw:real-hash"),
],
)
def test_reporting_key_filters_round_trip_without_losing_scope(raw_key: str, bucket: str) -> None:
from litellm.repositories.daily_activity_buckets import api_key_bucket
assert api_key_bucket(raw_key) == bucket
scope: Final = daily_activity_scope(
"litellm_dailyuserspend", "user_id", "user", None, bucket, "2024-01-01", "2024-01-01", None, None
)
assert scope.api_keys == (raw_key,)
assert scope.entity_ids == ("user",)
@pytest.mark.parametrize(("raw_key", "bucket_key"), [("", "unattributed"), ("unattributed", "raw:unattributed")])
def test_mcp_nested_key_metrics_use_reporting_identity(raw_key: str, bucket_key: str) -> None:
record: Final = SimpleNamespace(
**{
**vars(_daily_user_spend_record(user_id="user", api_key=raw_key, spend=0.25, model=None, model_group=None)),
"mcp_namespaced_tool_name": "deleted-server/tool",
}
)
result: Final = common_daily_activity_module._aggregate_spend_records_sync(
records=(record,), api_key_metadata={}, entity_id_field=None, entity_metadata_field=None
)
breakdown: Final = result["results"][0].breakdown
assert breakdown.mcp_servers["deleted-server/tool"].api_key_breakdown == breakdown.api_keys
assert breakdown.mcp_servers["deleted-server/tool"].api_key_breakdown[bucket_key].metrics.spend == 0.25
def test_ptu_rows_do_not_create_endpoint_or_mcp_key_buckets(ptu_cost_attribution_enabled: None) -> None:
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
record: Final = SimpleNamespace(
**{
**vars(_spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=2.0)),
"mcp_namespaced_tool_name": "server/tool",
"endpoint": "/mcp",
}
)
breakdown: Final = common_daily_activity_module.update_breakdown_metrics(BreakdownMetrics(), record, {}, {}, {})
assert breakdown.api_keys == {}
assert breakdown.endpoints["/mcp"].api_key_breakdown == {}
assert breakdown.mcp_servers["server/tool"].api_key_breakdown == {}
assert breakdown.endpoints["/mcp"].metrics.flat_cost == 2.0
assert breakdown.mcp_servers["server/tool"].metrics.flat_cost == 2.0
def test_metric_merges_preserve_every_additive_field() -> None:
metrics: Final = SpendMetrics.model_validate(dict.fromkeys(SpendMetrics.model_fields, 2))
merged: Final = common_daily_activity_module.sum_spend_metrics((SpendMetrics(), metrics, metrics))
assert merged.model_dump() == dict.fromkeys(SpendMetrics.model_fields, 4)

22
tests/unit/proxy/management_endpoints/test_daily_activity_routes.py
View file

@ -628,7 +628,9 @@ def test_aggregated_routes_return_scoped_results(
assert len(body["results"]) == 2, response.text
def _ptu_sized_router(model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})) -> Router:
def _ptu_sized_router(
model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})
) -> Router:
return Router(
model_list=[
{
@ -1268,13 +1270,15 @@ def test_team_scope_applies_membership_and_user_key_filter(
assert repository.aggregated.call_args.kwargs["include_entity_breakdown"] is True
@pytest.mark.parametrize("api_key", ["key-beta", "unattributed", "raw:unattributed", "raw:raw:unattributed"])
def test_team_scope_does_not_allow_an_unowned_api_key(
daily_activity_client: tuple[TestClient, _FakeRepository],
api_key: str,
) -> None:
client, _ = daily_activity_client
response: Final = client.get(
"/team/daily/activity/aggregated",
params={**_entity_params("team_ids", "team-a"), "api_key": "key-beta"},
params={**_entity_params("team_ids", "team-a"), "api_key": api_key},
headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"},
)
assert response.status_code == 200, response.text
@ -1622,3 +1626,17 @@ def test_export_content_disposition_is_ascii_and_built_from_canonical_dates(
disposition: Final = response.headers["content-disposition"]
assert disposition == 'attachment; filename="team-usage-2025-01-01-2025-01-02-daily.csv"'
assert disposition.isascii()
@pytest.mark.parametrize(("prefix", "query_name", "entity_id"), _ENTITY_CASES)
def test_unattributed_key_filter_does_not_become_an_unfiltered_request(
daily_activity_client: tuple[TestClient, _FakeRepository], prefix: str, query_name: str, entity_id: str
) -> None:
client, repository = daily_activity_client
response: Final = client.get(
f"{prefix}/daily/activity/aggregated",
params={**_entity_params(query_name, entity_id), "api_key": "unattributed"},
)
assert response.status_code == 200, response.text
assert response.json()["metadata"]["total_spend"] == 0
assert repository.aggregated.call_args.args[0].api_keys == ("",)

27
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -3617,6 +3617,33 @@ class TestDeleteTeamModels:
assert any(e[0] == "delete_many" for e in prisma.events)
@pytest.mark.parametrize(
("model_info_patch", "expected_gate"),
[
pytest.param({"tpm": 5}, True, id="unrelated-field"),
pytest.param({"supported_endpoints": ["/v1/embeddings", "/v1/rerank"]}, True, id="endpoint-list-only"),
pytest.param({"enforce_supported_endpoints": False}, False, id="explicit-false"),
],
)
def test_model_info_patch_changes_the_endpoint_gate_only_when_it_sends_the_flag(
model_info_patch: dict[str, object], expected_gate: bool
):
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
result: Final = update_db_model(
db_model=Deployment(
model_name="embedding-only",
litellm_params=LiteLLM_Params(model="openai/text-embedding-3-small"),
model_info=ModelInfo(
id="embedding-only", enforce_supported_endpoints=True, supported_endpoints=["/v1/embeddings"]
),
),
updated_patch=updateDeployment.model_validate({"model_info": model_info_patch}),
)
assert json.loads(result["model_info"]).get("enforce_supported_endpoints") is expected_gate
def _build_db_model_for_blocked_test():
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo

5
tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py
View file

@ -1058,7 +1058,8 @@ async def test_internal_user_tag_daily_activity_is_scoped_to_their_keys():
@pytest.mark.asyncio
async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter():
@pytest.mark.parametrize("api_key", ["unowned-key", "unattributed", "raw:unattributed", "raw:raw:unattributed"])
async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter(api_key: str):
"""
If an internal user filters tag usage by an API key they do not own, the
endpoint should return an empty scoped filter instead of exposing that key's
@ -1092,7 +1093,7 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter()
result = await get_tag_daily_activity(
start_date="2025-01-01",
end_date="2025-01-31",
api_key="unowned-key",
api_key=api_key,
user_api_key_dict=mock_user_auth,
)

19
tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py
View file

@ -485,7 +485,7 @@ def test_oidc_env_path_reference_is_not_resolved(tmp_path, monkeypatch):
)
def test_network_backed_oidc_reference_is_not_fetched_inline(monkeypatch):
def test_oidc_google_reference_is_refused_without_a_metadata_fetch(monkeypatch):
def fail(*args, **kwargs):
raise AssertionError("network-backed oidc reference was resolved inside the request path")
@ -494,11 +494,24 @@ def test_network_backed_oidc_reference_is_not_fetched_inline(monkeypatch):
lambda name, *a, **k: "oidc/google/https://example.com" if name == "OPENAI_API_KEY" else fail(),
)
assert (
with pytest.raises(HTTPException) as excinfo:
_passthrough_router(None).get_credentials(custom_llm_provider="openai", region_name=None)
== "oidc/google/https://example.com"
assert excinfo.value.status_code == 401
assert "pass-through credential oidc/google" in str(excinfo.value.detail)
assert "example.com" not in str(excinfo.value.detail)
def test_deployment_oidc_google_reference_is_not_sent_upstream(monkeypatch):
llm_router = litellm.Router(
model_list=[_flagged_deployment("openai/gpt-4o", api_key="oidc/google/https://vllm.example.test")]
)
with pytest.raises(HTTPException) as excinfo:
_passthrough_router(llm_router).get_credentials(custom_llm_provider="openai", region_name=None)
assert excinfo.value.status_code == 401
@pytest.fixture()
async def _drain_logging_worker():

5
tests/unit/proxy/proxy_server/test_lifecycle.py
View file

@ -35,6 +35,7 @@ from typing_extensions import TypedDict
import litellm.proxy.proxy_server as ps
from litellm.proxy.management_helpers.audit_logs import track_audit_task
from litellm.proxy.proxy_server import (
_proxy_startup_event,
ProxyStartupEvent,
_resolve_pydantic_type,
_resolve_typed_dict_type,
@ -1059,7 +1060,7 @@ def test_otel_global_provider_published_after_callback_init():
span was orphaned. The publish (``publish_global_otel_v2_provider``) must
therefore appear after ``_initialize_startup_logging`` in the lifespan source.
"""
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
source = inspect.getsource(wrapped)
init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(")
publish_pos = source.find("publish_global_otel_v2_provider(")
@ -1145,7 +1146,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
The call must follow Prisma setup so DB-backed deployments do not false-positive.
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
"""
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
source = inspect.getsource(wrapped)
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
warn_pos = source.find("warn_budget_without_db(")

16
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -2055,6 +2055,22 @@ def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure(
assert frame["error"]["message"] == handle_exception_on_proxy(failure).message
def test_responses_ws_failure_frame_shows_only_the_refusal_for_a_model_that_refuses_the_endpoint():
from litellm.exceptions import EndpointNotSupportedByModelError
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame
refusal: Final = "/v1/responses is not supported by model=embed-only. Supported endpoints: /v1/embeddings"
failure: Final = EndpointNotSupportedByModelError(
message=f"{refusal}\n\nLiteLLM: model group 'embed-only' failed with the error above", model="embed-only"
)
assert json.loads(_responses_ws_failure_frame(failure)) == {
"type": "error",
"status": 400,
"error": {"type": "invalid_request_error", "message": refusal},
}
class TestWSSessionCostTracking:
@pytest.mark.asyncio
async def test_router_budget_limiter_skips_arealtime_call_type(self):

104
tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py
View file

@ -179,3 +179,107 @@ async def test_emits_periodic_drain_waiting_log_while_waiting():
timeout=10, count_fn=count_fn, poll_interval=0, log_interval=0
)
assert drained == 3
@pytest.mark.parametrize("elapsed, cap, expected", [(4, None, 26), (25, 2, 2), (29, 5, 1), (31, None, 0)])
def test_remaining_budget_includes_time_before_lifespan(
monkeypatch: pytest.MonkeyPatch, elapsed: float, cap: float | None, expected: float
) -> None:
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0 + elapsed)
assert GracefulShutdownManager.remaining(cap, clock=lambda: 100.0 + elapsed) == expected
def test_remaining_budget_before_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
assert GracefulShutdownManager.remaining(clock=lambda: 100.0) == 30.0
assert GracefulShutdownManager.remaining(-1, clock=lambda: 100.0) == 0.0
@pytest.mark.parametrize("prestop", [False, True])
def test_sigterm_stamps_before_chaining_and_restores_server_handler(
monkeypatch: pytest.MonkeyPatch, prestop: bool
) -> None:
import signal
from queue import SimpleQueue
from types import FrameType
from typing import Final
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
calls: Final[SimpleQueue[float]] = SimpleQueue()
def server_handler(signum: int, frame: FrameType | None) -> None:
assert signum == signal.SIGTERM
calls.put(GracefulShutdownManager.remaining(clock=lambda: 105.0))
original: Final = signal.signal(signal.SIGTERM, server_handler)
try:
if prestop:
GracefulShutdownManager.start_shutdown(clock=lambda: 95.0)
with GracefulShutdownManager.capture_sigterm(clock=lambda: 100.0):
installed: Final = signal.getsignal(signal.SIGTERM)
with GracefulShutdownManager.capture_sigterm(clock=lambda: 102.0):
assert signal.getsignal(signal.SIGTERM) is installed
signal.raise_signal(signal.SIGTERM)
signal.raise_signal(signal.SIGTERM)
assert calls.get_nowait() == (20.0 if prestop else 25.0)
assert calls.get_nowait() == (20.0 if prestop else 25.0)
assert calls.empty()
assert signal.getsignal(signal.SIGTERM) is server_handler
finally:
signal.signal(signal.SIGTERM, original)
def test_sigterm_capture_preserves_a_replacement_handler() -> None:
import signal
from types import FrameType
from typing import Final
def server_handler(signum: int, frame: FrameType | None) -> None:
pass
original: Final = signal.signal(signal.SIGTERM, server_handler)
try:
with GracefulShutdownManager.capture_sigterm():
signal.signal(signal.SIGTERM, signal.SIG_IGN)
assert signal.getsignal(signal.SIGTERM) is signal.SIG_IGN
finally:
signal.signal(signal.SIGTERM, original)
@pytest.mark.parametrize("handler", ["SIG_DFL", "SIG_IGN"])
def test_sigterm_capture_preserves_default_and_ignored_dispositions(handler: str) -> None:
import signal
from typing import Final
disposition: Final = signal.SIG_DFL if handler == "SIG_DFL" else signal.SIG_IGN
original: Final = signal.signal(signal.SIGTERM, disposition)
try:
with GracefulShutdownManager.capture_sigterm():
assert signal.getsignal(signal.SIGTERM) is disposition
finally:
signal.signal(signal.SIGTERM, original)
def test_sigterm_capture_in_worker_thread_leaves_signal_ownership_unchanged() -> None:
import signal
from concurrent.futures import ThreadPoolExecutor
from typing import Final
original: Final = signal.getsignal(signal.SIGTERM)
def run_lifespan() -> None:
with GracefulShutdownManager.capture_sigterm():
assert signal.getsignal(signal.SIGTERM) is original
with ThreadPoolExecutor(max_workers=1) as executor:
executor.submit(run_lifespan).result()
assert signal.getsignal(signal.SIGTERM) is original
@pytest.mark.asyncio
async def test_lifespan_request_drain_does_not_restart_an_exhausted_budget(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
assert await GracefulShutdownManager.wait_for_drain(count_fn=lambda: 2, clock=lambda: 130.0) == 2

22
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -5384,6 +5384,28 @@ def test_get_logging_payload_router_rejected_request_takes_provider_from_model_g
assert payload["custom_llm_provider"] == expected_provider
@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}])
def test_get_logging_payload_drops_a_client_model_group_that_is_not_a_string(client_model_group: object):
payload: Final = get_logging_payload(
kwargs={
"model": "embed-only",
"litellm_params": {
"metadata": {"user_api_key": "test-key", "model_group": client_model_group, "status": "failure"}
},
"standard_logging_object": {
**_make_failed_request_standard_logging_payload(),
"model_group": client_model_group,
},
},
response_obj={},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
llm_router=None,
)
assert payload["model_group"] == ""
def test_get_logging_payload_router_rejected_request_without_router_leaves_provider_empty():
assert _router_rejected_failure_payload("openai-group", None)["custom_llm_provider"] == ""

58
tests/unit/proxy/test_common_request_processing.py
View file

@ -456,7 +456,7 @@ class TestProxyBaseLLMRequestProcessing:
mock_add_litellm_data_to_request,
)
mock_general_settings = {}
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict = ProxyUserAPIKeyAuth()
mock_proxy_config = MagicMock(spec=ProxyConfig)
route_type = "acompletion"
@ -740,7 +740,7 @@ class TestProxyBaseLLMRequestProcessing:
returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings={},
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
user_api_key_dict=ProxyUserAPIKeyAuth(),
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=MagicMock(spec=ProxyConfig),
route_type=route_type,
@ -1186,7 +1186,7 @@ class TestProxyBaseLLMRequestProcessing:
)
mock_general_settings = {}
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict = ProxyUserAPIKeyAuth()
mock_proxy_config = MagicMock(spec=ProxyConfig)
mock_router_settings = {
@ -2208,7 +2208,7 @@ class TestProxyBaseLLMRequestProcessing:
mock_add_litellm_data_to_request,
)
mock_general_settings = {}
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict = ProxyUserAPIKeyAuth()
mock_proxy_config = MagicMock(spec=ProxyConfig)
route_type = "acompletion"
@ -7706,6 +7706,56 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
assert processor.data["model"] == fallback_model
assert call_count == 2
@pytest.mark.asyncio
async def test_local_rate_limit_fallback_skips_an_aliased_embedding_only_model(self):
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
router: Final = Router(
model_list=[
{
"model_name": "embed-only",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-synthetic"},
"model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
},
{"model_name": "chat-model", "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}},
],
fallbacks=[{"rate-limited": ["embed-alias", "chat-model"]}],
)
processor: Final = ProxyBaseLLMRequestProcessing(
data={
"model": "rate-limited",
"router_settings_override": {"model_group_alias": {"embed-alias": "embed-only"}},
}
)
attempted: Final[list[object]] = []
async def pre_call_logic(**_: object):
attempted.append(processor.data.get("model"))
if processor.data.get("model") == "rate-limited":
raise ProxyRateLimitError(detail="key rpm limit exceeded", headers={"retry-after": "30"})
return processor.data, MagicMock()
with patch.object(processor, "common_processing_pre_call_logic", side_effect=pre_call_logic):
await processor._pre_call_with_fallbacks(
request=MagicMock(),
general_settings={},
proxy_logging_obj=MagicMock(),
user_api_key_dict=ProxyUserAPIKeyAuth(),
version=None,
proxy_config=MagicMock(),
user_model=None,
user_temperature=None,
user_request_timeout=None,
user_max_tokens=None,
user_api_base=None,
model="rate-limited",
route_type="acompletion",
llm_router=router,
)
assert attempted == ["rate-limited", "chat-model"]
@pytest.mark.asyncio
async def test_raises_when_no_fallbacks_configured(self):
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError

17
tests/unit/proxy/test_proxy_cli.py
View file

@ -4393,3 +4393,20 @@ def test_granian_flag_notices_are_json_records_in_json_mode(
"LiteLLM: --ciphers is not applied when using --run_granian.",
]
assert [record["level"] for record in records] == ["INFO", "WARNING", "WARNING"]
def test_migration_data_check_exits_zero_when_there_is_nothing_to_dry_run(monkeypatch, capfd):
monkeypatch.setenv("DATABASE_URL", "postgresql://user:pass@db.example/litellm")
assert ProxyInitializationHelpers.run_migration_data_check(use_migrate=False) == 0
out, err = capfd.readouterr()
assert out == "LiteLLM Proxy: The schema is pushed rather than migrated, there is nothing to dry-run.\n"
assert err == ""
def test_migration_data_check_requires_database_url(monkeypatch):
monkeypatch.delenv("DATABASE_URL", raising=False)
with pytest.raises(click.UsageError, match="requires DATABASE_URL"):
ProxyInitializationHelpers.run_migration_data_check(use_migrate=True)

199
tests/unit/proxy/test_proxy_server.py
View file

@ -263,6 +263,174 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm
token = MASTER_KEY
@pytest.mark.asyncio
async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None:
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.responses.streaming_iterator import spawn_logging_task
worker: Final = LoggingWorker()
started: Final = asyncio.Event()
release: Final = asyncio.Event()
first_flush: Final = asyncio.Event()
pending: Final[asyncio.Queue[str]] = asyncio.Queue()
persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
async def enqueue_spend() -> None:
pending.put_nowait("accounted request")
async def callback() -> None:
started.set()
await release.wait()
spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting")
async def flush_spend() -> None:
while not pending.empty():
persisted.put_nowait(pending.get_nowait())
first_flush.set()
worker.ensure_initialized_and_enqueue(callback())
await started.wait()
shutdown: Final = asyncio.create_task(
litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker)
)
try:
await first_flush.wait()
waited_for_callback: Final = not shutdown.done()
release.set()
await shutdown
assert waited_for_callback, "Shutdown returned while accounting callbacks were still running"
assert persisted.get_nowait() == "accounted request"
assert pending.empty()
finally:
release.set()
await worker.flush()
await worker.stop()
@pytest.mark.asyncio
async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
worker: Final = LoggingWorker()
started: Final = asyncio.Event()
cancelled: Final = asyncio.Event()
flushed_after_cancellation: Final = asyncio.Event()
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
async def callback() -> None:
started.set()
try:
await asyncio.Event().wait()
finally:
cancelled.set()
async def flush_spend() -> None:
if cancelled.is_set():
flushed_after_cancellation.set()
worker.ensure_initialized_and_enqueue(callback())
await started.wait()
try:
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
flush_spend, worker=worker, clock=lambda: 120.0
)
assert cancelled.is_set()
assert flushed_after_cancellation.is_set()
assert "budget exhausted with 1 unfinished callbacks" in caplog.text
finally:
GracefulShutdownManager.reset()
await worker.stop()
@pytest.mark.asyncio
async def test_shutdown_bounds_callback_cancellation_and_flushes_accounting(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
worker: Final = LoggingWorker()
started: Final = asyncio.Event()
cancelling: Final = asyncio.Event()
release: Final = asyncio.Event()
finished: Final = asyncio.Event()
flushed: Final = asyncio.Event()
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
async def callback() -> None:
started.set()
try:
await release.wait()
except asyncio.CancelledError:
cancelling.set()
await release.wait()
finally:
finished.set()
async def drain_tasks(timeout: float) -> None:
assert timeout == 0.0
await cancelling.wait()
async def flush_spend() -> None:
if cancelling.is_set():
flushed.set()
worker.ensure_initialized_and_enqueue(callback())
await started.wait()
try:
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
flush_spend, worker=worker, clock=lambda: 130.0, drain_tasks=drain_tasks
)
assert not finished.is_set(), "The cancellation-resistant callback must remain observable"
assert flushed.is_set()
assert "tasks still alive after cancellation" in caplog.text
finally:
release.set()
await worker.flush()
await worker.stop()
GracefulShutdownManager.reset()
@pytest.mark.asyncio
async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
worker: Final = LoggingWorker()
persisted: Final = asyncio.Event()
drain_budgets: Final[asyncio.Queue[float]] = asyncio.Queue()
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
async def flush_spend() -> None:
persisted.set()
async def drain_tasks(timeout: float) -> None:
drain_budgets.put_nowait(timeout)
try:
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
flush_spend, worker=worker, clock=lambda: 106.0, drain_tasks=drain_tasks
)
assert "remaining_s=24.000 callback_timeout_s=14.000" in caplog.text
assert persisted.is_set()
assert tuple(drain_budgets.get_nowait() for _ in range(drain_budgets.qsize())) == (14.0, 5.0)
finally:
await worker.stop()
GracefulShutdownManager.reset()
headers = {"Authorization": f"Bearer {token}"}
example_completion_result = {
@ -4059,3 +4227,34 @@ def test_detailed_debug_logs_request_body_only_as_received_and_as_sent(
assert len(body_dumps) == 2, [message[:120] for message in body_dumps]
assert body_dumps[0].startswith("Request received by LiteLLM:")
assert "POST Request Sent from LiteLLM" in body_dumps[1]
@pytest.mark.asyncio
async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> None:
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.responses.streaming_iterator import spawn_logging_task
worker: Final = LoggingWorker()
persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
accounting: Final[asyncio.Queue[str]] = asyncio.Queue()
async def callback() -> None:
accounting.put_nowait("completed request")
async def producer() -> None:
await asyncio.sleep(0)
await asyncio.sleep(0)
worker.ensure_initialized_and_enqueue(callback())
async def flush_spend() -> None:
while not accounting.empty():
persisted.put_nowait(accounting.get_nowait())
spawn_logging_task(asyncio.get_running_loop(), producer(), task_name="success enqueue")
try:
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker)
assert persisted.qsize() == 1, "An enqueue producer must finish before callback draining stops"
assert persisted.get_nowait() == "completed request"
assert worker.pending_count == 0
finally:
await worker.stop(discard_pending=True)

33
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -279,6 +279,39 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid
assert str(tmp_path) not in excinfo.value.message
@pytest.mark.asyncio
async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc_google_api_key_fails(
monkeypatch: pytest.MonkeyPatch,
):
audience: Final = f"https://count.test/{uuid.uuid4().hex}"
monkeypatch.setattr(
litellm.proxy.proxy_server,
"llm_router",
Router(
model_list=[
{
"model_name": "claude",
"litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": f"oidc/google/{audience}"},
}
]
),
)
with respx.mock(assert_all_called=True) as upstream:
upstream.get(url__regex=r"/computeMetadata/v1/instance/service-accounts/default/identity").respond(
500, text="metadata server unavailable"
)
with pytest.raises(ProxyException) as excinfo:
await token_counter(
request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]),
call_endpoint=True,
)
assert excinfo.value.code == "503"
assert "oidc/google secret reference could not be resolved" in excinfo.value.message
assert audience.removeprefix("https://") not in excinfo.value.message
@pytest.mark.asyncio
async def test_native_openai_responses_count_sends_the_token_an_oidc_api_key_points_to(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch

25
tests/unit/proxy/test_route_llm_request.py
View file

@ -28,6 +28,7 @@ from litellm.proxy.route_llm_request import (
JSON_OBJECT_ADAPTER,
ProxyModelNotFoundError,
UserConfigRouters,
endpoint_unsupported_by_model,
route_request,
user_config_routers,
)
@ -3027,3 +3028,27 @@ async def test_user_config_routers_discard_the_least_recently_used_router_past_c
assert first_handler() is not None
assert second_handler() is None
routers.clear()
def test_fallback_probe_skips_only_a_loaded_model_that_refuses_the_route() -> None:
router: Final = litellm.Router(
model_list=[
{
"model_name": "chat-fallback",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "unit-provider-key"},
},
{
"model_name": "embed-only",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "unit-provider-key"},
"model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
},
],
num_retries=0,
)
def skipped(model: str) -> bool:
return endpoint_unsupported_by_model("acompletion", {"model": model}, router, None)
assert skipped("embed-only") is True
assert skipped("chat-fallback") is False
assert skipped("added-on-another-worker") is False

156
tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -41,6 +41,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import TextChoices, TextCompletionResponse
@pytest.fixture(autouse=True)
@ -1886,7 +1887,7 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator
@pytest.mark.parametrize(
"request_route",
["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
["/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
)
def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response(
make_user_api_key_auth, monkeypatch, caplog, request_route
@ -2036,6 +2037,120 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_cho
assert _warnings(caplog) == []
def _text_completion_chunk(text: str, finish_reason: str | None = None) -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-1", model="m", choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)]
)
def _rewritten_text_completion(response: Any) -> TextCompletionResponse:
return TextCompletionResponse(
id=response.id,
model=response.model,
choices=[TextChoices(index=0, text="[REWRITTEN] " + response.choices[0].text, finish_reason="stop")],
)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_completion_stream(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=_rewritten_text_completion)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
chunks = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")]
auth = make_user_api_key_auth(request_route="/v1/completions")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await proxy_logging.pre_call_hook(
user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True
)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
assert isinstance(seen["response"], TextCompletionResponse)
assert seen["response"].choices[0].text == "hello world"
assert [(item.choices[0].text, item.choices[0].finish_reason) for item in delivered] == [
("[REWRITTEN] hello world", None),
("", None),
("", "stop"),
]
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
assert _warnings(caplog) == []
class _FlatBlockGuardrail(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
if any("kumquat" in text for text in inputs.get("texts") or []):
raise HTTPException(status_code=400, detail="flat guardrail blocked kumquat")
return inputs
@pytest.mark.asyncio
async def test_streaming_iterator_hook_delivers_a_flat_guardrail_block_on_a_text_completion_stream_under_a_pipeline(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Final[Dict[str, Any]] = {}
flat: Final = _FlatBlockGuardrail(
guardrail_name="flat-block", event_hook=GuardrailEventHooks.post_call, default_on=True
)
monkeypatch.setattr(litellm, "callbacks", [flat, _live_scan_guardrail(seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
chunks: Final = [
_text_completion_chunk("A kum"),
_text_completion_chunk("quat tart"),
_text_completion_chunk("", "stop"),
]
auth: Final = make_user_api_key_auth(request_route="/v1/completions")
await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
delivered: Final = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
assert seen["texts"] == ["A kumquat tart"]
assert isinstance(delivered[-1], bytes)
assert json.loads(delivered[-1].decode().removeprefix("data: "))["error"]["message"] == (
"flat guardrail blocked kumquat"
)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_live_pipeline_scans_what_a_disconnected_text_completion_client_received(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Final[Dict[str, Any]] = {}
monkeypatch.setattr(litellm, "callbacks", [_live_scan_guardrail(seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data: Final = {**_post_call_pipeline_data(step=_DETECT_ONLY_STEP, stream=True), "prompt": "hi"}
chunks: Final = [
_text_completion_chunk("hello "),
_text_completion_chunk("world"),
_text_completion_chunk("", "stop"),
]
auth: Final = make_user_api_key_auth(request_route="/v1/completions")
await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
stream: Final = proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
first: Final = await stream.__anext__()
await stream.aclose()
assert first.choices[0].text == "hello "
assert seen["texts"] == ["hello "]
assert _post_call_statuses(data) == ["success"]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
proxy_logging, make_user_api_key_auth, monkeypatch
@ -2615,6 +2730,45 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated
assert not any(item is chunk for item in delivered for chunk in chunks)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_modify_response_replaces_a_text_completion_stream(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Final[Dict[str, Any]] = {}
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
pipeline: Final = GuardrailPipeline(
mode="post_call",
steps=[
PipelineStep(
guardrail="gr-post",
on_pass="allow",
on_fail="modify_response",
modify_response_message="content policy block",
)
],
)
data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)]
chunks: Final = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")]
auth: Final = make_user_api_key_auth(request_route="/v1/completions")
await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
delivered: Final = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
frames: Final = [json.loads(item.decode().removeprefix("data: ")) for item in delivered]
assert seen["count"] == 1
assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [
("content policy block", None),
("", "content_filter"),
]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse(
proxy_logging, make_user_api_key_auth, monkeypatch

17
tests/unit/repositories/test_daily_activity_sql.py
View file

@ -123,7 +123,6 @@ def test_key_page_sql_orders_exact_spend_and_binds_scope_before_page() -> None:
7,
)
assert "SUM(spend::numeric) AS rank_spend" in query.sql
assert "ORDER BY rank_spend DESC, api_key" in query.sql
assert "(SELECT COUNT(*) FROM ranked)::bigint AS total_api_keys" in query.sql
@ -329,7 +328,6 @@ def test_entity_rollup_bounds_keys_and_reuses_scope_filters() -> None:
assert "api_key = ANY($3::text[])" in query.sql
assert query.sql.count("api_key = ANY($3::text[])") == 4
assert query.sql.count("api_key <> $4") == 2
assert query.sql.count("ORDER BY SUM(spend::numeric) DESC, api_key") == 1
assert "k.entity_id = e.entity_id" in query.sql
assert "LIMIT $5" in query.sql
assert query.params == ("2026-01-01", "2026-01-31", ["key-1", "key-2"], PTU_SENTINEL_API_KEY, 3)
@ -351,13 +349,11 @@ def test_search_query_escapes_pattern_metacharacters_and_binds_limit() -> None:
assert 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = d.user_id' in query.sql
assert "d.key_alias ILIKE $3 ESCAPE" in query.sql
assert "d.user_id ILIKE $3 ESCAPE" in query.sql
assert "api_key ILIKE $3 ESCAPE" in query.sql
assert "v.key_alias ILIKE $3 ESCAPE" in query.sql
assert "v.user_id ILIKE $3 ESCAPE" in query.sql
assert "u.user_email ILIKE $3 ESCAPE" in query.sql
assert query.sql.count("ILIKE $3 ESCAPE") == 7
assert "api_key <> $4" in query.sql
assert "ORDER BY SUM(spend::numeric) DESC, api_key" in query.sql
assert "LIMIT $5" in query.sql
assert query.params == (
"2026-01-01",
@ -374,12 +370,9 @@ def test_model_and_cache_key_queries_bind_filters_sentinel_and_limits() -> None:
)
leakage_query = build_cache_leakage_keys_sql(_scope(entity_ids=None), limit=20)
assert "COALESCE(NULLIF(model_group, ''), model) = $3" in model_query.sql
assert "api_key <> $4" in model_query.sql
assert "ORDER BY SUM(spend::numeric) DESC, api_key" in model_query.sql
assert model_query.params == ("2026-01-01", "2026-01-31", "public-model", PTU_SENTINEL_API_KEY, 5)
assert "HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0" in leakage_query.sql
assert "ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key" in leakage_query.sql
assert leakage_query.params == ("2026-01-01", "2026-01-31", PTU_SENTINEL_API_KEY, 20)
@ -401,8 +394,14 @@ def test_query_builders_reject_nonpositive_limits(builder) -> None:
("export_type", "group_key", "key_filter", "joins"),
[
(ExportType.DAILY, "''", "", ""),
(ExportType.DAILY_WITH_KEYS, "scoped.api_key", "api_key <> $3", 'LEFT JOIN "LiteLLM_VerificationToken"'),
(ExportType.DAILY_WITH_MODELS, "COALESCE(scoped.model, '')", "api_key <> $3", ""),
(
ExportType.DAILY_WITH_KEYS,
"CASE WHEN api_key = '' THEN 'unattributed' "
"WHEN api_key = 'unattributed' OR api_key LIKE 'raw:%' THEN 'raw:' || api_key ELSE api_key END",
"api_key <> $3",
'LEFT JOIN "LiteLLM_VerificationToken"',
),
(ExportType.DAILY_WITH_MODELS, "COALESCE(NULLIF(model, ''), 'unknown-model')", "api_key <> $3", ""),
(
ExportType.DAILY_WITH_USERS,
"COALESCE(vt.user_id, dvt.user_id, '')",

139
tests/unit/router_utils/test_endpoint_eligibility.py Normal file
View file

@ -0,0 +1,139 @@
from collections.abc import Mapping
from typing import Final
import pytest
from pydantic import TypeAdapter
from litellm.exceptions import EndpointNotSupportedByModelError
from litellm.router_utils.endpoint_eligibility import (
eligible_deployments,
endpoint_selection,
warn_on_unenforced_endpoint_flag,
)
_DEPLOYMENT: Final = TypeAdapter(list[dict[str, object]])
def _deployment(info: Mapping[str, object], identity: str = "embedding") -> dict[str, object]:
return _DEPLOYMENT.validate_python(
[
{
"model_name": "shared",
"litellm_params": {"model": "openai/test-model"},
"model_info": {"id": identity, **info},
}
]
)[0]
@pytest.mark.parametrize("call_type", ["acompletion", "anthropic_messages", "aresponses", "acreate_interaction"])
def test_explicit_restriction_excludes_an_embedding_deployment_from_a_mixed_group(call_type: str) -> None:
embedding: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]})
chat: Final = _deployment({"supported_endpoints": ["/v1/chat/completions"]}, "chat")
with endpoint_selection(call_type):
result: Final = eligible_deployments("shared", (embedding, chat))
assert result == (chat,)
@pytest.mark.parametrize(
"info",
[
{"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]},
{"enforce_supported_endpoints": False, "supported_endpoints": ["/v1/embeddings"]},
{"enforce_supported_endpoints": True},
{"enforce_supported_endpoints": True, "supported_endpoints": None},
{"enforce_supported_endpoints": True, "supported_endpoints": 7},
{"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/messages"]},
],
)
def test_advisory_or_incomplete_metadata_preserves_routing(info: Mapping[str, object]) -> None:
deployment: Final = _deployment(info)
with endpoint_selection("acompletion"):
result: Final = eligible_deployments("shared", (deployment,))
assert result == (deployment,)
def test_a_single_endpoint_string_is_enforced_like_a_one_item_list() -> None:
deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": "/v1/embeddings"})
with endpoint_selection("acompletion"):
result: Final = eligible_deployments("shared", (deployment,))
assert isinstance(result, EndpointNotSupportedByModelError)
assert result.message.endswith("Supported endpoints: /v1/embeddings")
def test_unsupported_error_explains_the_requested_operation_and_available_endpoints() -> None:
deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]})
with endpoint_selection("acreate_interaction"):
result: Final = eligible_deployments("shared", (deployment,))
assert isinstance(result, EndpointNotSupportedByModelError)
assert result.status_code == 400
assert "/v1beta/interactions" in result.message
assert "/v1/embeddings" in result.message
@pytest.mark.parametrize("call_type", ["aembedding", "aspeech", "_arealtime", "aget_responses"])
def test_chat_restriction_does_not_block_other_operations(call_type: str) -> None:
deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]})
with endpoint_selection(call_type):
result: Final = eligible_deployments("shared", (deployment,))
assert result == (deployment,)
def test_nested_non_chat_calls_do_not_inherit_the_outer_chat_restriction() -> None:
deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]})
with endpoint_selection("acompletion"):
with endpoint_selection("aembedding"):
embedding_result: Final = eligible_deployments("shared", (deployment,))
chat_result: Final = eligible_deployments("shared", (deployment,))
assert embedding_result == (deployment,)
assert isinstance(chat_result, EndpointNotSupportedByModelError)
@pytest.mark.parametrize(
"endpoint",
[
"/messages",
"/interactions",
"/responses/compact",
"/openai/v1/responses/compact",
"/openai/v1/chat/completions",
"chat/completions",
"v1/responses",
" /v1/chat/completions ",
"/v1/chat/completions/",
"/V1/Chat/Completions",
"/v1beta/models/gemini-3-flash:generateContent",
],
)
def test_chat_endpoint_declarations_keep_their_meaning_across_prefixes(endpoint: str) -> None:
deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": [endpoint]})
with endpoint_selection("acompletion"):
result: Final = eligible_deployments("shared", (deployment,))
assert result == (deployment,)
@pytest.mark.parametrize(
"info",
[
{"enforce_supported_endpoints": True},
{"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings", None]},
{"enforce_supported_endpoints": "true", "supported_endpoints": 7},
],
)
def test_a_flag_without_a_usable_endpoint_list_is_reported(info: Mapping[str, object]) -> None:
warning: Final = warn_on_unenforced_endpoint_flag("embed-only", info)
assert warning is not None
assert "'embed-only'" in warning
@pytest.mark.parametrize(
"info",
[
{"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
{"enforce_supported_endpoints": True, "supported_endpoints": "/v1/embeddings"},
{"enforce_supported_endpoints": False},
{"supported_endpoints": ["/v1/embeddings", None]},
],
)
def test_an_enforceable_or_unflagged_deployment_is_not_reported(info: Mapping[str, object]) -> None:
assert warn_on_unenforced_endpoint_flag("embed-only", info) is None

42
tests/unit/secret_managers/test_secret_managers_main.py
View file

@ -4,6 +4,7 @@ import logging
import os
import threading
import time
import uuid
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from queue import SimpleQueue
@ -11,6 +12,7 @@ from typing import Final
from unittest.mock import Mock, patch
import pytest
import respx
import litellm
from litellm._logging import redact_internal_details_from_client_message
@ -37,7 +39,7 @@ class MockHTTPHandler:
self.text = "mocked_token"
self.json_data = {"value": "mocked_token"}
def get(self, url, params=None, headers=None):
def get(self, url, params=None, headers=None, timeout=None):
# Store params for audience verification
self.last_params = params
logger.debug(
@ -49,6 +51,9 @@ class MockHTTPHandler:
mock_response.json.return_value = self.json_data
return mock_response
def close(self):
pass
@pytest.fixture
def mock_oidc_cache():
@ -229,6 +234,41 @@ def test_oidc_google_failure():
get_secret(secret_name)
def test_oidc_google_fetches_the_token_from_the_metadata_host_gce_metadata_host_names(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("GCE_METADATA_HOST", "127.0.0.1:8930")
audience: Final = f"https://vllm.example.test/{uuid.uuid4().hex}"
with respx.mock(assert_all_called=True) as metadata:
metadata.get(
"http://127.0.0.1:8930/computeMetadata/v1/instance/service-accounts/default/identity",
params={"audience": audience},
headers={"Metadata-Flavor": "Google"},
).respond(200, text="metadata-host-token")
token: Final = get_secret(f"oidc/google/{audience}")
assert token == "metadata-host-token"
def test_oidc_google_secret_outside_a_model_api_key_keeps_its_metadata_wait_when_gce_metadata_timeout_is_set(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("GCE_METADATA_TIMEOUT", "1")
audience: Final = f"https://sts.example.test/{uuid.uuid4().hex}"
with respx.mock(assert_all_called=True) as metadata:
identity: Final = metadata.get(
"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity",
params={"audience": audience},
).respond(200, text="web-identity-token")
token: Final = get_secret(f"oidc/google/{audience}")
waits: Final = identity.calls.last.request.extensions["timeout"]
assert token == "web-identity-token"
assert waits == {"connect": 5.0, "read": 600.0, "write": 600.0, "pool": 600.0}
def test_oidc_circleci_success(monkeypatch):
monkeypatch.setenv("CIRCLE_OIDC_TOKEN", "circleci_token")

127
tests/unit/test_dashscope_image_generation.py
View file

@ -5,21 +5,26 @@ qwen-image-3.0, qwen-image-3.0-pro).
Run in docker: pytest tests/unit/test_dashscope_image_generation.py -v
"""
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
from pydantic import ValidationError
import litellm
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.dashscope.common_utils import (
DASHSCOPE_DEFAULT_ORIGIN,
DASHSCOPE_IMAGE_PATH,
)
from litellm.llms.dashscope.image_generation.transformation import (
DashScopeImageGenerationConfig,
DEFAULT_API_BASE,
DashScopeImageGenerationConfig,
)
from litellm.types.utils import ImageResponse
from litellm.utils import get_llm_provider
from litellm.llms.base_llm.chat.transformation import BaseLLMException
# ---------------------------------------------------------------------------
# 1. Provider detection
@ -77,6 +82,47 @@ class TestDashScopeImageGenerationConfig:
url = self.cfg.get_complete_url(chat_api_base, None, "qwen-image-3.0", {}, {})
assert url == DEFAULT_API_BASE
@pytest.mark.parametrize(
("api_base", "expected"),
[
(
"https://dashscope.aliyuncs.com",
"https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
),
(
"https://llm-abc123.cn-beijing.maas.aliyuncs.com",
"https://llm-abc123.cn-beijing.maas.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
),
(
"https://llm-abc123.cn-beijing.maas.aliyuncs.com/apps/anthropic",
"https://llm-abc123.cn-beijing.maas.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
),
(
"https://dashscope-us.aliyuncs.com/api/v1/",
"https://dashscope-us.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
),
],
)
def test_get_complete_url_derives_from_the_domain(
self, api_base: str, expected: str
):
url: Final = self.cfg.get_complete_url(api_base, None, "qwen-image-3.0", {}, {})
assert url == expected
@pytest.mark.parametrize(
"endpoint_url",
[
"https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation",
"https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis",
],
)
def test_get_complete_url_keeps_a_full_dashscope_endpoint_url(
self, monkeypatch: pytest.MonkeyPatch, endpoint_url: str
):
monkeypatch.setenv("DASHSCOPE_API_BASE_IMAGE", "https://images.example.com/generate")
url: Final = self.cfg.get_complete_url(endpoint_url, None, "qwen-image-3.0", {}, {})
assert url == endpoint_url
def test_validate_environment_sets_auth_header(self):
headers = self.cfg.validate_environment(
headers={},
@ -458,6 +504,81 @@ def test_litellm_image_generation_dashscope_end_to_end(model: str):
assert body["parameters"]["size"] == "1024*1024"
@pytest.mark.parametrize(
("api_base", "chat_base_env", "image_base_env", "expected_url"),
[
(None, "", "", DEFAULT_API_BASE),
("https://dashscope.aliyuncs.com", "", "", f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_IMAGE_PATH}"),
(None, "", "https://images.example.com/generate", "https://images.example.com/generate"),
(
"https://dashscope.aliyuncs.com/compatible-mode/v1",
"",
"https://images.example.com/generate",
"https://images.example.com/generate",
),
(None, "https://gateway.example.com/dashscope/images", "", "https://gateway.example.com/dashscope/images"),
(None, "https://dashscope-us.aliyuncs.com", "https://images.example.com/generate", "https://images.example.com/generate"),
(
None,
f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}",
"https://images.example.com/generate",
f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}",
),
(
None,
"https://gateway.example.com/dashscope/images",
"https://images.example.com/generate",
"https://gateway.example.com/dashscope/images",
),
(
None,
"https://dashscope-us.aliyuncs.com/compatible-mode/v1",
"https://images.example.com/generate",
"https://images.example.com/generate",
),
(
"https://dashscope-us.aliyuncs.com",
"",
"https://images.example.com/generate",
f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}",
),
(
None,
"https://dashscope-us.aliyuncs.com",
"",
f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}",
),
(
"https://dashscope-intl.aliyuncs.com",
"https://gateway.example.com/dashscope/images",
"",
f"https://dashscope-intl.aliyuncs.com{DASHSCOPE_IMAGE_PATH}",
),
],
)
def test_image_generation_resolves_from_the_configured_base_not_the_chat_default(
monkeypatch: pytest.MonkeyPatch,
api_base: str | None,
chat_base_env: str,
image_base_env: str,
expected_url: str,
):
monkeypatch.setenv("DASHSCOPE_API_BASE", chat_base_env)
monkeypatch.setenv("DASHSCOPE_API_BASE_IMAGE", image_base_env)
image_url: Final = "https://dashscope-result.oss.aliyuncs.com/t.png"
with respx.mock(assert_all_called=True) as upstream:
upstream.post(expected_url).respond(
200, json={"output": {"choices": [{"message": {"content": [{"image": image_url}]}}]}}
)
response: Final = litellm.image_generation(
model="dashscope/qwen-image-3.0", prompt="a puppy", api_key="sk-test-key", api_base=api_base
)
assert response.data[0].url == image_url
def _transform_response(payload: object) -> ImageResponse:
return DashScopeImageGenerationConfig().transform_image_generation_response(
model="qwen-image-2.0",

113
tests/unit/test_router/test_router.py
View file

@ -12270,6 +12270,67 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag():
)
def test_synchronous_text_completion_rejects_a_restricted_model_before_the_provider() -> None:
model: Final = "text-completion-openai/unit-text-model"
router: Final = Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": model,
"api_base": "https://sync-router.invalid/v1",
"api_key": "unit-provider-key",
},
"model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
}
],
num_retries=0,
)
with respx.mock(assert_all_called=False) as upstream:
provider: Final = upstream.post("https://sync-router.invalid/v1/completions").respond(
200,
json={
"id": "cmpl-unit",
"object": "text_completion",
"created": 1,
"model": "unit-text-model",
"choices": [{"text": "provider called", "index": 0, "finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
with pytest.raises(litellm.BadRequestError, match="/v1/completions is not supported"):
router.text_completion(model=model, prompt="test")
assert provider.call_count == 0
def test_routable_deployments_leave_out_models_that_refuse_the_call_type() -> None:
router: Final = Router(
model_list=[
{
"model_name": "available",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "unit-provider-key"},
},
{
"model_name": "embed-only",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "unit-provider-key"},
"model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
},
],
num_retries=0,
)
def routable(model: str, call_type: str | None) -> tuple[str, ...]:
return tuple(
deployment["model_name"] for deployment in router.get_routable_deployments(model, {}, call_type=call_type)
)
assert routable("unregistered-fallback", "acompletion") == ()
assert routable("available", "acompletion") == ("available",)
assert routable("embed-only", "acompletion") == ()
assert routable("embed-only", "aembedding") == ("embed-only",)
assert routable("embed-only", None) == ("embed-only",)
def test_routing_groups_legacy_property_forwards_to_public_attribute():
router = _router_with_two_deployments([False, False])
@ -18803,6 +18864,35 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi
assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)]
@pytest.mark.asyncio
async def test_anthropic_messages_same_group_retry_skips_an_embedding_only_deployment():
"""The retry of a /v1/messages stream dropped before content must not land on an enforced embedding-only
deployment of the same group, even when that deployment is the group's preferred order."""
router: Final = Router(
model_list=[
{"model_name": "glm", "litellm_params": {"model": "anthropic/glm-chat", "api_key": "sk-test", "order": 2}},
{
"model_name": "glm",
"litellm_params": {"model": "anthropic/glm-embed", "api_key": "sk-test", "order": 1},
"model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]},
},
],
num_retries=1,
)
provider: Final = _AnthropicMessagesScriptedProvider(
_anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream
)
async def anthropic_messages(**kwargs: object) -> object:
return await provider(**kwargs)
stream: Final = await _anthropic_messages_stream_through_router(router, anthropic_messages)
body: Final = [chunk async for chunk in stream]
assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")]
assert [model for model, _, _ in provider.calls] == ["anthropic/glm-chat", "anthropic/glm-chat"]
class _AnthropicMessagesRetryStashRecorder:
def __init__(self) -> None:
self.observed: tuple[Usage | None, ...] = ()
@ -31431,3 +31521,26 @@ def test_router_keeps_supported_fallback_shapes(
router: Final = Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs)
assert router.fallbacks == expected_fallbacks
def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog: pytest.LogCaptureFixture) -> None:
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
router: Final = Router(
model_list=[
{
"model_name": "embed-from-config",
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-fake"},
"model_info": {"enforce_supported_endpoints": True},
}
]
)
router.add_deployment(
Deployment(
model_name="embed-from-api",
litellm_params=LiteLLM_Params(model="openai/text-embedding-3-small", api_key="sk-fake"),
model_info=ModelInfo(enforce_supported_endpoints=True, supported_endpoints=["/v1/embeddings", 7]),
)
)
assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text
assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text

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