diff --git a/codecov.yaml b/codecov.yaml index 5b30e78a8fb..521088868bd 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -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: diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index e059d45e037..a3488d6114c 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -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` diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index e22639daaef..75c5bcf4e1e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -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." diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py new file mode 100644 index 00000000000..c24742ce145 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py @@ -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 "") + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py index 5a55b35b255..1c16eb5e6b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -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( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 528bf7def20..8b19576d86e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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 diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 61e3a8fb5dc..28211b3dab2 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -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, diff --git a/litellm/images/main.py b/litellm/images/main.py index b3247ff938b..4aad5210028 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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( diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index ff649953e69..8f217b1ea5e 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -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): """ diff --git a/litellm/litellm_core_utils/served_output_texts.py b/litellm/litellm_core_utils/served_output_texts.py index 854785f304b..b19973e06c3 100644 --- a/litellm/litellm_core_utils/served_output_texts.py +++ b/litellm/litellm_core_utils/served_output_texts.py @@ -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 diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 966d1a01b8d..dfff9cec298 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -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, diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index 3e2e718a469..5ad2aa17a3a 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -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": diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 35132346261..0f01a4b3172 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -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, diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index d1faa6c6cd0..2f528b5a2c3 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -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, diff --git a/litellm/llms/dashscope/messages/__init__.py b/litellm/llms/dashscope/messages/__init__.py new file mode 100644 index 00000000000..08cd184404c --- /dev/null +++ b/litellm/llms/dashscope/messages/__init__.py @@ -0,0 +1,7 @@ +""" +DashScope Anthropic-compatible Messages Module +""" + +from .transformation import DashScopeAnthropicMessagesConfig + +__all__ = ["DashScopeAnthropicMessagesConfig"] diff --git a/litellm/llms/dashscope/messages/transformation.py b/litellm/llms/dashscope/messages/transformation.py new file mode 100644 index 00000000000..9d0d1db91c8 --- /dev/null +++ b/litellm/llms/dashscope/messages/transformation.py @@ -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") diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py index f7cc50cc82a..db01c8a0a15 100644 --- a/litellm/llms/dashscope/qwen_ai_platform.py +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -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): diff --git a/litellm/llms/dashscope/qwencloud.py b/litellm/llms/dashscope/qwencloud.py index a12e21b0983..6ae42c62e42 100644 --- a/litellm/llms/dashscope/qwencloud.py +++ b/litellm/llms/dashscope/qwencloud.py @@ -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): diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 5c738423a5a..59e42e7b926 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -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)] diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index c4bc04d4627..27c82cfc004 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -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) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 098dd2f9816..67728e0fd63 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fc29a6e3a08..7e487ded569 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index f889748c504..d1787f3ff8b 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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]: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..16dbe2b4131 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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( diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 3d783d8aba1..be5de040fdf 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 7ff5aa2466b..e57d6e78378 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -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={}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f58c8705815..71b8976bebf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -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 diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 3dceab6b7f7..b79154bf509 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 66f386b792d..98372b64015 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -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) ) diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 79e515944ce..1a3445aae52 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -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: diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index 77246f5dda6..ca1a96f9d84 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -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, diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index c12a90ed6a2..ea7810a618c 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 50b700e712d..b712ba42a6a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 1d0b3d00434..739d4b3a01a 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -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/ or oidc/env/." + ), + ) match resolve_local_oidc_reference(credential): case UnresolvableSecretReference(kind=kind, reason=reason): raise HTTPException( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0d8ab98a0b1..f71d1960e24 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4a2d27e964d..7145126ee6b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, ) ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9e642a02d49..772bb3e53d2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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 diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index db1951ec9e2..fea4f342b91 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -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__( diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py index d7a07ec34e5..9aeee0453ff 100644 --- a/litellm/proxy/shutdown/graceful_shutdown_manager.py +++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py @@ -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", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 833ee71ae9c..df9ed611a47 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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( diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 919fd54e7d6..e03cb44efda 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -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") ), diff --git a/litellm/repositories/daily_activity_buckets.py b/litellm/repositories/daily_activity_buckets.py new file mode 100644 index 00000000000..ec6dd0f8c84 --- /dev/null +++ b/litellm/repositories/daily_activity_buckets.py @@ -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) diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 4ca23e4e89e..321468c4976 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -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( diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 7ec9fa67a67..4bb52767de6 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -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", (), diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 65d9b89cf79..e344ce46be7 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -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, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8959c32fff2..6a82304b76a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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( diff --git a/litellm/router.py b/litellm/router.py index 5aac475ba8f..046a8d97f5b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py new file mode 100644 index 00000000000..4a87e582c97 --- /dev/null +++ b/litellm/router_utils/endpoint_eligibility.py @@ -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="", + ) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 588bdb1409c..02cf53d00ec 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -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") diff --git a/litellm/types/router.py b/litellm/types/router.py index db8be9e91a3..56dcad4051a 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index ac07fe1560b..1cb28be8910 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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/ or oidc/env/.", + "Use oidc/file/, oidc/env/ or oidc/google/.", 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/ or oidc/env/.", + 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 ( diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 330ddf88fc5..367135eae26 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -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"} diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py new file mode 100644 index 00000000000..c45adfeeb84 --- /dev/null +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -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}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index afddc7e25d6..504adec5401 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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 diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index f558c355c84..d06f0b13cff 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -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) diff --git a/tests/integration/routing/test_auto_router_management.py b/tests/integration/routing/test_auto_router_management.py index 5aed328ba7b..442dae4d689 100644 --- a/tests/integration/routing/test_auto_router_management.py +++ b/tests/integration/routing/test_auto_router_management.py @@ -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,)) diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py new file mode 100644 index 00000000000..465f5801692 --- /dev/null +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -0,0 +1,1515 @@ +import asyncio +import contextlib +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import openai +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + list_value, + object_value, + string_value, +) +from integration._support import responses_ws_peer +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process +from integration._support.tls import write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from websockets.exceptions import ConnectionClosed + +_PROVIDER_KEY: Final = "integration-provider-key" +_EMBEDDING_MODEL: Final = "text-embedding-3-small" +_CHAT_MODEL: Final = "gpt-4o-mini" +_CHAT_BACKEND_PREFIX: Final = "gate-chat-backend-" +_SERVED_TEXT: Final = "served by the chat upstream" +_VECTOR: Final = (0.25, 0.5, 0.75) +_COUNTED_TOKENS: Final = 17 +_CLOSE: Final = MappingProxyType({"connection": "close"}) +_RELOAD_SECONDS: Final = 5 +_EMBED_ONLY: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"], "enforce_supported_endpoints": True} +) + + +def _embedding_reply() -> Reply: + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(_VECTOR)}], + "model": _EMBEDDING_MODEL, + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } + ).encode() + ) + + +def _not_a_chat_model() -> Reply: + return Reply( + status=404, + body=json.dumps( + {"error": {"message": "scripted: not a chat model", "type": "invalid_request_error", "code": None}} + ).encode(), + ) + + +def _text_completion_reply() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"cmpl-{uuid.uuid4().hex}", + "object": "text_completion", + "created": 1, + "model": _CHAT_MODEL, + "choices": [{"index": 0, "text": _SERVED_TEXT, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, + } + ).encode() + ) + + +def _serves_chat(model: JsonValue) -> bool: + return model == _CHAT_MODEL or (isinstance(model, str) and model.startswith(_CHAT_BACKEND_PREFIX)) + + +def _respond(request: Request) -> Reply: + path: Final = request.target.split("?")[0] + body: Final = JSON_OBJECT.validate_json(request.body or b"{}") + streamed: Final = body.get("stream") is True + if path == "/v1/embeddings": + return _embedding_reply() + if path == "/v1/responses/input_tokens": + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS}).encode()) + if not _serves_chat(body.get("model")): + return _not_a_chat_model() + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path in ("/v1/responses", "/v1/responses/compact"): + return responses_reply(f"resp_{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path == "/v1/completions": + return _text_completion_reply() + return _not_a_chat_model() + + +class _Upstream: + def __init__(self, wire: Wire) -> None: + self.wire = wire + self.seen: tuple[Request, ...] = () + + def calls(self, marker: str) -> tuple[Request, ...]: + self.seen = (*self.seen, *self.wire.drain()) + return tuple(request for request in self.seen if marker.encode() in request.body or marker in request.target) + + def targets(self, marker: str) -> tuple[str, ...]: + return tuple(request.target for request in self.calls(marker)) + + +@dataclass(frozen=True, slots=True) +class _Deployment: + litellm_model: str + declarations: tuple[Mapping[str, JsonValue], ...] + + +def _single(declaration: Mapping[str, JsonValue], litellm_model: str = f"openai/{_CHAT_MODEL}") -> _Deployment: + return _Deployment(litellm_model, (declaration,)) + + +_SHARED: Final[Mapping[str, _Deployment]] = MappingProxyType( + { + "embedding_only": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "embedding_only_second": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "image_generation": _single( + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/images/generations"]} + ), + "chat_target": _single({"mode": "chat"}), + "chat_rate_limited": _single({"mode": "chat"}), + "mode_not_a_string": _single( + {"mode": 7, "supported_endpoints": ["/v1/embeddings"]}, "openai/gate-mode-not-a-string" + ), + } +) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + api_base: str + + +def _create(gateway: Gateway, name: str, litellm_params: Mapping[str, JsonValue], info: Mapping[str, JsonValue]) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"api_key": _PROVIDER_KEY, **litellm_params}, + "model_info": dict(info), + }, + ) + return string_value(object_value(created["model_info"])["id"]) + + +def _deploy(scenario: Scenario, api_base: str, label: str, deployment: _Deployment) -> tuple[str, tuple[str, ...]]: + name: Final = f"gate-{label}-{uuid.uuid4().hex[:12]}" + identities: Final = tuple( + _create(scenario.gateway, name, {"model": deployment.litellm_model, "api_base": api_base}, declaration) + for declaration in deployment.declarations + ) + tuple(scenario.cleanups.callback(scenario.delete_model, identity) for identity in identities) + return name, identities + + +def _listed_ids(gateway: Gateway) -> frozenset[str]: + listed: Final = gateway.request("GET", "/model/info", headers=_CLOSE) + assert listed.status_code == 200, listed.text + return frozenset( + string_value(object_value(object_value(entry)["model_info"])["id"]) + for entry in list_value(JSON_OBJECT.validate_json(listed.content)["data"]) + ) + + +def _served_by_every_worker(gateway: Gateway, identities: frozenset[str]) -> None: + held( + lambda: identities <= _listed_ids(gateway), + lambda served: served, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + api_base: Final = wire.url + "/v1" + deployed: Final = {label: _deploy(scenario, api_base, label, spec) for label, spec in _SHARED.items()} + _served_by_every_worker( + gateway, frozenset(identity for _name, identities in deployed.values() for identity in identities) + ) + yield _Rig( + gateway, + _Upstream(wire), + MappingProxyType({label: name for label, (name, _ids) in deployed.items()}), + api_base, + ) + + +def _marker() -> str: + return f"gate-{uuid.uuid4().hex}" + + +def _openai_error(endpoint: str, model: str, supported: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"{endpoint} is not supported by model={model}. Supported endpoints: {supported}", + "type": "invalid_request_error", + "param": "model", + "code": "400", + } + } + + +def _anthropic_error(model: str, supported: str) -> dict[str, JsonValue]: + return { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"/v1/messages is not supported by model={model}. Supported endpoints: {supported}", + }, + } + + +def _chat_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], **({"stream": True} if stream else {})} + + +def _failure_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + return rows[0] + + +def _assert_error(response: httpx.Response, expected: dict[str, JsonValue]) -> None: + assert response.status_code == 400, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + message: Final = string_value(error["message"]).split("\n\n", 1)[0] + assert {**body, "error": {**error, "message": message}} == expected, response.text + + +def _assert_rejected_before( + upstream: _Upstream, response: httpx.Response, marker: str, expected: dict[str, JsonValue] +) -> None: + _assert_error(response, expected) + assert upstream.calls(marker) == (), upstream.targets(marker) + assert _failure_row(response.headers["x-litellm-call-id"])["status"] == "failure" + + +def _assert_rejected(rig: _Rig, response: httpx.Response, marker: str, expected: dict[str, JsonValue]) -> None: + _assert_rejected_before(rig.upstream, response, marker, expected) + + +@dataclass(frozen=True, slots=True) +class _Route: + path: str + endpoint: str + anthropic_shape: bool + stream: bool + model_in_path: bool + shape: str + + +_ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( + { + "interactions": _Route("/v1beta/interactions", "/v1beta/interactions", False, False, False, "interactions"), + "chat_v1": _Route("/v1/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_v1_stream": _Route("/v1/chat/completions", "/v1/chat/completions", False, True, False, "chat"), + "chat_bare": _Route("/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_engines": _Route("/engines/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat"), + "chat_azure_deployments": _Route( + "/openai/deployments/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat" + ), + "messages": _Route("/v1/messages", "/v1/messages", True, False, False, "messages"), + "messages_stream": _Route("/v1/messages", "/v1/messages", True, True, False, "messages"), + "responses_v1": _Route("/v1/responses", "/v1/responses", False, False, False, "responses"), + "responses_v1_stream": _Route("/v1/responses", "/v1/responses", False, True, False, "responses"), + "responses_bare": _Route("/responses", "/v1/responses", False, False, False, "responses"), + "cursor_responses_body": _Route("/cursor/chat/completions", "/v1/responses", False, False, False, "cursor"), + "cursor_chat_body": _Route("/cursor/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + } +) + + +def _route_body(route: _Route, model: str, marker: str) -> dict[str, JsonValue]: + shaped: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {"messages": [{"role": "user", "content": marker}]}, + "messages": {"max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + "responses": {"input": marker}, + "interactions": {"input": marker}, + "cursor": {"input": [{"role": "user", "content": marker}]}, + } + return { + **({} if route.model_in_path else {"model": model}), + **shaped[route.shape], + **({"stream": True} if route.stream else {}), + } + + +def _send(gateway: Gateway, route: _Route, model: str, marker: str) -> httpx.Response: + path: Final = route.path.format(model=model) + body: Final = _route_body(route, model, marker) + if not route.stream: + return gateway.request("POST", path, body) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as streamed: + streamed.read() + return streamed + + +@pytest.mark.parametrize("route", tuple(_ROUTES)) +def test_chat_shaped_route_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, route: str) -> None: + spec: Final = _ROUTES[route] + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _send(rig.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected(rig, response, marker, expected) + + +def _sdk_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + async def call() -> tuple[tuple[float, ...], str]: + async with openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", api_key=rig.gateway.key, max_retries=0 + ) as sdk: + raw: Final = await sdk.embeddings.with_raw_response.create( + model=model, input=marker, encoding_format="float" + ) + parsed: Final = raw.parse() + return tuple(parsed.data[0].embedding), raw.headers["x-litellm-call-id"] + + return asyncio.run(call()) + + +def _httpx_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + response: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"} + ) + assert response.status_code == 200, response.text + data: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + vector: Final = list_value(object_value(data[0])["embedding"]) + return tuple(float(str(value)) for value in vector), response.headers["x-litellm-call-id"] + + +@pytest.mark.parametrize("client", ("httpx_sync", "openai_async_sdk")) +def test_embeddings_still_reach_the_embedding_only_model(rig: _Rig, client: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + vector, call_id = ( + _httpx_embedding(rig, model, marker) if client == "httpx_sync" else _sdk_embedding(rig, model, marker) + ) + assert vector == _VECTOR + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + row: Final = eventually( + lambda: read_rows( + 'SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id = %s OR litellm_call_id = %s', + (call_id, call_id), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert [(entry["status"], entry["call_type"]) for entry in row] == [("success", "aembedding")], row + + +def _assert_served(rig: _Rig, response: httpx.Response, marker: str, target: str) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == (target,) + + +def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_name(rig: _Rig) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(alias, marker), key=key) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", target, "/v1/embeddings")) + + +async def _websocket_first_frame_and_close( + gateway: Gateway, key: str, model: str, marker: str +) -> tuple[dict[str, JsonValue], int | None, str | None]: + url: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://", 1) + "/v1/responses" + async with websockets.connect(url, additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10) as ws: + await ws.send(json.dumps({"type": "response.create", "model": model, "input": marker})) + first: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(ws.recv(), 30)) + with contextlib.suppress(ConnectionClosed): + await asyncio.wait_for(ws.recv(), 30) + return first, ws.close_code, ws.close_reason + + +def _websocket_rejection(model: str) -> tuple[dict[str, JsonValue], int | None, str | None]: + return ( + { + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "message": f"/v1/responses is not supported by model={model}. Supported endpoints: /v1/embeddings", + }, + }, + 1008, + "Pre-call error", + ) + + +@pytest.mark.parametrize("aliased", [False, True]) +def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstream(rig: _Rig, aliased: bool) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-ws-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + closed: Final = asyncio.run( + _websocket_first_frame_and_close(rig.gateway, key, alias if aliased else target, marker) + ) + assert closed == _websocket_rejection(target) + assert rig.upstream.calls(marker) == () + + +_RESERVING_KEYS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "max_parallel_requests": {"max_parallel_requests": 1}, + "tpm_limit": {"tpm_limit": 100, "default_estimated_output_tokens": 80}, + } +) + + +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_responses_websocket_rejections_leave_the_key_reservation_free(rig: _Rig, limit: str) -> None: + target: Final = rig.names["embedding_only"] + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rig.gateway, key, target, marker)) for marker in markers + ) + assert closed == (_websocket_rejection(target),) * 3 + served: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["chat_target"], served_marker), key=key + ) + _assert_served(rig, served, served_marker, "/v1/chat/completions") + assert [rig.upstream.calls(marker) for marker in markers] == [()] * 3 + + +def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None: + public: Final = f"gate-team-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team() + identity: Final = _create( + rig.gateway, + public, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base}, + {**_EMBED_ONLY, "team_id": team}, + ) + scenario.cleanups.callback(scenario.delete_model, identity) + key: Final = scenario.key(team_id=team) + _served_by_every_worker(rig.gateway, frozenset({identity})) + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": public, "input": marker, "encoding_format": "float"}, key=key + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + rejected_marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(public, rejected_marker), key=key + ) + _assert_rejected( + rig, response, rejected_marker, _openai_error("/v1/chat/completions", public, "/v1/embeddings") + ) + + +def test_wildcard_embedding_only_deployment_rejects_chat_for_any_matched_model(rig: _Rig) -> None: + prefix: Final = f"gate-wild-{uuid.uuid4().hex[:12]}" + requested: Final = f"{prefix}/{_EMBEDDING_MODEL}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"{prefix}/*", {"model": "openai/*", "api_base": rig.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(requested, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", requested, "/v1/embeddings")) + + +def test_key_without_access_gets_the_access_error_and_learns_nothing_about_endpoints(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[rig.names["chat_target"]]) + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["embedding_only"], marker), key=key + ) + assert response.status_code == 403, response.text + assert object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"] == "key_model_access_denied", ( + response.text + ) + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_unknown_model_keeps_the_invalid_model_error(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(f"gate-missing-{uuid.uuid4().hex}", marker) + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text and "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_request_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": [rig.names["chat_target"]]} + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": []} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_fallbacks_that_cannot_answer_chat_keep_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + fallbacks: Final = [rig.names["image_generation"], model] + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": fallbacks} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_request_fallbacks_do_not_rescue_an_unknown_primary(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(f"gate-missing-{uuid.uuid4().hex}", marker), "fallbacks": [rig.names["chat_target"]]}, + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def _chat(rig: _Rig, model: str, marker: str, *, key: str | None = None, **extra: JsonValue) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", {**_chat_body(model, marker), **extra}, key=key) + + +def _assert_rejected_chat(rig: _Rig, response: httpx.Response, marker: str, model: str) -> None: + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_a_fallback_chain_that_reaches_chat_two_hops_away_answers(rig: _Rig) -> None: + first, second, chat = rig.names["embedding_only"], rig.names["embedding_only_second"], rig.names["chat_target"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [chat]}]) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_a_fallback_chain_that_loops_between_embedding_only_groups_keeps_the_rejection(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [first]}]) + _assert_rejected_chat(rig, response, marker, first) + + +def test_a_specific_chain_that_loops_beats_a_catch_all_that_could_answer(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat( + rig, first, marker, fallbacks=[{"*": [rig.names["chat_target"]]}, {first: [second]}, {second: [first]}] + ) + _assert_rejected_chat(rig, response, marker, first) + + +def test_disable_fallbacks_keeps_the_rejection_over_request_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, fallbacks=[rig.names["chat_target"]], disable_fallbacks=True) + _assert_rejected_chat(rig, response, marker, model) + + +def _assert_fallback_answered( + upstream: _Upstream, response: httpx.Response, marker: str, target: str = "/v1/chat/completions" +) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert upstream.targets(marker) == (target,), response.text + + +def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: + ghost: Final = f"gate-ghost-{uuid.uuid4().hex}" + return {"fallbacks": [{rig.names["embedding_only"]: [ghost]}, {ghost: [rig.names["chat_target"]]}]} + + +def test_request_fallbacks_naming_a_chat_backend_name_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + backend: Final = f"openai/{_CHAT_BACKEND_PREFIX}{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"gate-backend-{uuid.uuid4().hex[:12]}", {"model": backend, "api_base": rig.api_base}, {} + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = _chat(rig, model, marker, fallbacks=[{model: [backend]}]) + _assert_fallback_answered(rig.upstream, response, marker) + + +@pytest.mark.parametrize("level", ("key", "team")) +def test_router_settings_fallbacks_through_a_group_with_no_deployments_answer(rig: _Rig, level: str) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + settings: Final = {"router_settings": dict(_through_a_group_with_no_deployments(rig))} + key: Final = scenario.key(**settings) if level == "key" else scenario.key(team_id=scenario.team(**settings)) + response: Final = _chat(rig, rig.names["embedding_only"], marker, key=key) + _assert_fallback_answered(rig.upstream, response, marker) + + +def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=scenario.key(team_id=team)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_key_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=key) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_key_router_settings_that_cannot_answer_override_team_fallbacks_that_could(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + key: Final = scenario.key( + team_id=team, router_settings={"fallbacks": [{model: [rig.names["embedding_only_second"]]}]} + ) + response: Final = _chat(rig, model, marker, key=key) + _assert_rejected_chat(rig, response, marker, model) + + +def test_an_empty_request_fallback_list_overrides_key_router_settings_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + response: Final = _chat(rig, model, marker, key=key, fallbacks=[]) + _assert_rejected_chat(rig, response, marker, model) + + +def test_local_rate_limit_fallback_skips_an_embedding_only_hop(rig: _Rig) -> None: + limited, chat, model = rig.names["chat_rate_limited"], rig.names["chat_target"], rig.names["embedding_only"] + first, second = _marker(), _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: [model, chat]}]} + ) + served_first: Final = _chat(rig, limited, first, key=key) + served_second: Final = _chat(rig, limited, second, key=key) + assert (served_first.status_code, served_first.headers.get("x-litellm-model-group")) == (200, limited), ( + served_first.text + ) + assert (served_second.status_code, served_second.headers.get("x-litellm-model-group")) == (200, chat), ( + served_second.text + ) + assert (rig.upstream.targets(first), rig.upstream.targets(second)) == ( + ("/v1/chat/completions",), + ("/v1/chat/completions",), + ) + + +_COMPACT_PATHS: Final = ("/v1/responses/compact", "/responses/compact", "/openai/v1/responses/compact") + + +def _compact_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": [{"role": "user", "content": marker}]} + + +@pytest.mark.parametrize("path", _COMPACT_PATHS) +def test_responses_compact_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, path: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", path, _compact_body(model, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/responses/compact", model, "/v1/embeddings")) + + +def test_responses_compact_keeps_reaching_the_upstream_for_a_chat_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/compact", _compact_body(rig.names["chat_target"], marker) + ) + _assert_served(rig, response, marker, "/v1/responses/compact") + + +def test_count_tokens_keeps_counting_for_an_embedding_only_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": rig.names["embedding_only"], "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNTED_TOKENS, response.text + assert rig.upstream.targets(marker) == ("/v1/responses/input_tokens",) + + +def test_text_completions_reject_an_embedding_only_model_before_the_upstream(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", "/v1/completions", {"model": model, "prompt": marker}) + _assert_rejected(rig, response, marker, _openai_error("/v1/completions", model, "/v1/embeddings")) + + +@pytest.mark.parametrize("method", ("generateContent", "streamGenerateContent")) +def test_gemini_routes_reject_an_embedding_only_model_before_the_upstream(rig: _Rig, method: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", f"/v1beta/models/{model}:{method}", {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + ) + assert response.status_code == 400, response.text + assert f"{method} is not supported by model={model}. Supported endpoints: /v1/embeddings" in response.text + assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) + + +def test_repeated_rejections_log_one_failure_row_per_call(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + responses: Final = tuple( + rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for _ in range(3) + ) + expected: Final = _openai_error("/v1/chat/completions", model, "/v1/embeddings") + for response in responses: + _assert_error(response, expected) + call_ids: Final = tuple(item.headers["x-litellm-call-id"] for item in responses) + assert len(set(call_ids)) == 3, call_ids + assert [_failure_row(call_id)["status"] for call_id in call_ids] == ["failure"] * 3 + assert rig.upstream.calls(marker) == () + + +def _closed_chat(rig: _Rig, model: str, marker: str) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker), headers=_CLOSE) + + +def test_live_declaration_update_turns_the_gate_on_and_off_on_every_worker(rig: _Rig) -> None: + chat_capable: Final[dict[str, JsonValue]] = { + "mode": "chat", + "supported_endpoints": ["/v1/chat/completions", "/v1/embeddings"], + } + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy(scenario, rig.api_base, "flip", _single(chat_capable)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + served: Final = _closed_chat(rig, name, marker) + _assert_served(rig, served, marker, "/v1/chat/completions") + patched: Final = rig.gateway.request( + "PATCH", f"/model/{identities[0]}/update", {"model_info": dict(_EMBED_ONLY)} + ) + assert patched.status_code == 200, patched.text + rejected: Final = held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 400, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + _assert_error(rejected, _openai_error("/v1/chat/completions", name, "/v1/embeddings")) + restored: Final = rig.gateway.request("PATCH", f"/model/{identities[0]}/update", {"model_info": chat_capable}) + assert restored.status_code == 200, restored.text + held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 200, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +_TESTS_ROOT: Final = Path(__file__).resolve().parents[2] +_CONFIG: Final = TypeAdapter(dict[str, JsonValue]) + + +def _owned_embedding(owned: Gateway, model: str, marker: str) -> httpx.Response: + return owned.request("POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"}) + + +@dataclass(frozen=True, slots=True) +class _Routed: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + + +def _routed_entry(name: str, api_base: str, litellm_model: str, info: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": litellm_model, "api_base": api_base, "api_key": _PROVIDER_KEY}, + "model_info": dict(info), + } + + +def _routed_config(directory: Path, api_base: str, names: Mapping[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["embed"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "default_fallbacks": [names["chat"]], + "model_group_alias": {names["alias_embed"]: names["embed"], names["alias_chat"]: names["chat"]}, + }, + } + path: Final = directory / "gate-routed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def routed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Routed]: + directory: Final = tmp_path_factory.mktemp("gate-routed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + {label: f"gate-routed-{label}-{suffix}" for label in ("embed", "chat", "alias_embed", "alias_chat")} + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_routed_config(directory, wire.url + "/v1", names), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Routed(owned.gateway, _Upstream(wire), names) + + +_ROUTED_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +_DEFAULT_FALLBACK_ROUTES: Final[Mapping[str, str]] = MappingProxyType( + {"chat_v1": "/v1/chat/completions", "messages": "/v1/responses", "responses_v1": "/v1/responses"} +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", tuple(_DEFAULT_FALLBACK_ROUTES)) +def test_router_default_fallbacks_answer_for_an_embedding_only_group(routed: _Routed, route: str) -> None: + marker: Final = _marker() + response: Final = _send(routed.gateway, _ROUTES[route], routed.names["embed"], marker) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == (_DEFAULT_FALLBACK_ROUTES[route],) + + +_OPT_OUTS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "empty_fallbacks": {"fallbacks": []}, + "null_fallbacks": {"fallbacks": None}, + "disable_fallbacks": {"disable_fallbacks": True}, + } +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("opt_out", tuple(_OPT_OUTS)) +def test_request_opt_outs_keep_the_rejection_over_router_default_fallbacks(routed: _Routed, opt_out: str) -> None: + model: Final = routed.names["embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), **_OPT_OUTS[opt_out]} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_falls_back_to_the_default(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_embed"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_is_rejected_without_fallbacks(routed: _Routed) -> None: + alias: Final = routed.names["alias_embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(alias, marker), "disable_fallbacks": True} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", alias, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_a_chat_group_keeps_reaching_the_upstream(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_chat"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_keeps_serving_embeddings(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(routed.gateway, routed.names["alias_embed"], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert routed.upstream.targets(marker) == ("/v1/embeddings",) + + +_REWRITER: Final = """from litellm.integrations.custom_logger import CustomLogger + + +class WebSocketModelRewriter(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if call_type == "_aresponses_websocket" and data.get("model") == "__SOURCE__": + data["model"] = "__TARGET__" + return data + + +rewriter = WebSocketModelRewriter() +""" + + +@dataclass(frozen=True, slots=True) +class _Rewritten: + gateway: Gateway + upstream: _Upstream + source: str + target: str + + +def _rewritten_config(directory: Path, api_base: str, source: str, target: str, module: str) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(source, api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(target, api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + ], + "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, + "litellm_settings": {**object_value(shared["litellm_settings"]), "callbacks": [f"{module}.rewriter"]}, + } + path: Final = directory / "gate-rewritten-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def rewritten(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rewritten]: + directory: Final = tmp_path_factory.mktemp("gate-rewritten") + suffix: Final = uuid.uuid4().hex[:12] + source: Final = f"gate-rewrite-chat-{suffix}" + target: Final = f"gate-rewrite-embed-{suffix}" + module: Final = f"gate_ws_rewriter_{suffix}" + (directory / f"{module}.py").write_text(_REWRITER.replace("__SOURCE__", source).replace("__TARGET__", target)) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_rewritten_config(directory, wire.url + "/v1", source, target, module), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(source, "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Rewritten(owned.gateway, _Upstream(wire), source, target) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_a_websocket_model_a_hook_rewrites_to_an_embedding_only_group_releases_the_key_reservation( + rewritten: _Rewritten, limit: str +) -> None: + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rewritten.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rewritten.gateway, key, rewritten.source, marker)) + for marker in markers + ) + assert closed == (_websocket_rejection(rewritten.target),) * 3 + served: Final = rewritten.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rewritten.source, served_marker), key=key + ) + assert served.status_code == 200, served.text + assert _SERVED_TEXT in served.text, served.text + assert [rewritten.upstream.calls(marker) for marker in markers] == [()] * 3 + assert rewritten.upstream.targets(served_marker) == ("/v1/chat/completions",) + + +@dataclass(frozen=True, slots=True) +class _Addressed: + gateway: Gateway + rig: Gateway + upstream: _Upstream + names: Mapping[str, str] + addresses: Mapping[str, str] + api_base: str + + +def _wildcard_backend_entry(name: str, api_base: str, backend: str) -> dict[str, JsonValue]: + entry: Final = _routed_entry(name, api_base, backend, _EMBED_ONLY) + return {**entry, "litellm_params": {**object_value(entry["litellm_params"]), "custom_llm_provider": "openai"}} + + +def _addressed_config( + directory: Path, api_base: str, names: Mapping[str, str], suffix: str, with_wildcards: bool +) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + wildcard: Final = f"gate-addr-wild-{suffix}" + wildcard_entries: Final = ( + ( + _routed_entry(f"{wildcard}/*", api_base, "openai/*", {}), + _wildcard_backend_entry(names["wildcard_backend"], api_base, f"{wildcard}/{_CHAT_MODEL}"), + ) + if with_wildcards + else () + ) + aliases: Final = {names["wildcard_alias"]: f"{wildcard}/{_CHAT_BACKEND_PREFIX}{suffix}"} if with_wildcards else {} + config: Final = { + **shared, + "model_list": [ + _routed_entry( + names["embed"], api_base, f"openai/gate-addr-backend-{suffix}", {**_EMBED_ONLY, "id": names["embed_id"]} + ), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat", "id": names["chat_id"]}), + _routed_entry(names["shared_embed"], api_base, f"openai/{_CHAT_MODEL}", _EMBED_ONLY), + *wildcard_entries, + _routed_entry(names["ghost_primary"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "model_group_alias": aliases, + "fallbacks": [{names["ghost_primary"]: [names["ghost"]]}, {names["ghost"]: [names["chat"]]}], + }, + "general_settings": { + **object_value(shared["general_settings"]), + "proxy_config_reload_interval_seconds": 3600, + }, + } + path: Final = directory / "gate-addressed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@contextlib.contextmanager +def _addressed_proxy(tmp_path_factory: pytest.TempPathFactory, with_wildcards: bool) -> Iterator[_Addressed]: + directory: Final = tmp_path_factory.mktemp("gate-addressed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + { + label: f"gate-addr-{label}-{suffix}" + for label in ( + "embed", + "embed_id", + "chat", + "chat_id", + "shared_embed", + "wildcard_backend", + "wildcard_alias", + "ghost_primary", + "ghost", + ) + } + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_addressed_config(directory, wire.url + "/v1", names, suffix, with_wildcards), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Addressed( + owned.gateway, + rig_gateway, + _Upstream(wire), + names, + MappingProxyType( + { + "deployment_id": names["embed_id"], + "backend_name": f"openai/gate-addr-backend-{suffix}", + "wildcard_backend_name": f"gate-addr-wild-{suffix}/{_CHAT_MODEL}", + } + ), + wire.url + "/v1", + ) + + +@pytest.fixture(scope="module") +def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=True) as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def addressed_without_wildcards(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=False) as proxy: + yield proxy + + +_ADDRESSED_ROUTES: Final = ("chat_v1", "messages", "responses_v1") + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", _ADDRESSED_ROUTES) +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_chat_sent_to_an_embedding_only_deployment_id_or_backend_name_is_rejected_before_the_upstream( + addressed: _Addressed, address: str, route: str +) -> None: + spec: Final = _ROUTES[route] + model: Final = addressed.addresses[address] + marker: Final = _marker() + response: Final = _send(addressed.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected_before(addressed.upstream, response, marker, expected) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_embeddings_sent_to_an_embedding_only_deployment_id_or_backend_name_keep_working( + addressed: _Addressed, address: str +) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(addressed.gateway, addressed.addresses[address], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert addressed.upstream.targets(marker) == ("/v1/embeddings",) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("chat_deployment_id", "backend_name_a_chat_deployment_shares")) +def test_chat_sent_to_a_chat_capable_deployment_id_or_backend_name_keeps_reaching_the_upstream( + addressed: _Addressed, address: str +) -> None: + model: Final = addressed.names["chat_id"] if address == "chat_deployment_id" else f"openai/{_CHAT_MODEL}" + marker: Final = _marker() + response: Final = addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_a_chat_deployment_id_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["chat_id"]]}, + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_chat_sent_to_an_embedding_only_backend_name_a_wildcard_also_matches_is_served_by_the_wildcard( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.addresses["wildcard_backend_name"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in addressed.upstream.calls(marker)] == [_CHAT_MODEL] + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_an_alias_only_a_wildcard_serves_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["wildcard_alias"]]}, + ) + _assert_fallback_answered(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_yaml_fallbacks_through_a_group_with_no_deployments_answer_for_an_embedding_only_group( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.names["ghost_primary"], marker) + ) + _assert_fallback_answered(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rejected_on_every_call( + addressed_without_wildcards: _Addressed, +) -> None: + addressed: Final = addressed_without_wildcards + name: Final = f"gate-addr-loaded-{uuid.uuid4().hex[:12]}" + markers: Final = tuple(_marker() for _ in range(6)) + with addressed.rig.scenario() as scenario: + identity: Final = _create( + addressed.rig, name, {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": addressed.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + responses: Final = tuple( + addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(name, marker), headers=_CLOSE) + for marker in markers + ) + expected: Final = _openai_error("/v1/chat/completions", name, "/v1/embeddings") + for response in responses: + _assert_error(response, expected) + assert [addressed.upstream.calls(marker) for marker in markers] == [()] * len(markers) + + +@pytest.mark.parametrize("metadata", ({"model_group": 7}, {"model_group": ["wrong-group"]})) +def test_rejection_with_invalid_metadata_is_logged_without_blocking_embeddings( + rig: _Rig, metadata: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + model: Final = rig.names["embedding_only"] + rejected: Final = _chat(rig, model, marker, metadata=metadata) + _assert_rejected_chat(rig, rejected, marker, model) + embedded_marker: Final = _marker() + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": embedded_marker, "encoding_format": "float"} + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(embedded_marker) == ("/v1/embeddings",) + + +def test_existing_non_chat_metadata_on_a_chat_backend_remains_advisory(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy( + scenario, + rig.api_base, + "legacy-audio", + _single({"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + + +def test_a_mixed_group_skips_its_restricted_embedding_deployment(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name: Final = f"gate-mixed-{uuid.uuid4().hex[:12]}" + embed: Final = _create( + scenario.gateway, + name, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base, "order": 0}, + _EMBED_ONLY, + ) + scenario.cleanups.callback(scenario.delete_model, embed) + chat: Final = _create( + scenario.gateway, name, {"model": f"openai/{_CHAT_MODEL}", "api_base": rig.api_base, "order": 1}, {} + ) + scenario.cleanups.callback(scenario.delete_model, chat) + _served_by_every_worker(rig.gateway, frozenset({embed, chat})) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + calls: Final = rig.upstream.calls(marker) + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in calls] == [_CHAT_MODEL] + + +@pytest.mark.parametrize( + "fallbacks, expected_status", + [(("embedding_only",), 429), (("embedding_only", "chat_target"), 200), (("missing",), 400)], +) +def test_local_rate_limit_skips_restricted_fallbacks_and_still_tries_unloaded_ones( + rig: _Rig, fallbacks: tuple[str, ...], expected_status: int +) -> None: + limited: Final = rig.names["chat_rate_limited"] + targets: Final = tuple(rig.names[name] if name != "missing" else _marker() for name in fallbacks) + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: list(targets)}]} + ) + first: Final = _chat(rig, limited, _marker(), key=key) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = _chat(rig, limited, marker, key=key) + assert second.status_code == expected_status, second.text + if expected_status == 200: + _assert_served(rig, second, marker, "/v1/chat/completions") + return + assert "Supported endpoints" not in second.text, second.text + assert rig.upstream.calls(marker) == () + + +def _websocket_fallback_config(directory: Path, primary: tuple[str, str], fallback: tuple[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(primary[0], primary[1], f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + _routed_entry(fallback[0], fallback[1], f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "fallbacks": [{primary[0]: [fallback[0]]}], + }, + } + path: Final = directory / "gate-websocket-fallback.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +def test_responses_websocket_uses_a_configured_chat_fallback_without_calling_the_embedding_provider( + rig: _Rig, tmp_path: Path +) -> None: + marker: Final = _marker() + suffix: Final = uuid.uuid4().hex[:12] + primary: Final = f"gate-ws-embed-{suffix}" + fallback: Final = f"gate-ws-chat-{suffix}" + certificate: Final = write_self_signed_cert(tmp_path) + with ( + responses_ws_peer.responses_peer( + certificate, lambda identity: responses_ws_peer.events(identity, _SERVED_TEXT) + ) as peer, + owned_proxy_process( + rig.gateway, + tmp_path, + {"SSL_CERT_FILE": str(certificate[0])}, + config=_websocket_fallback_config(tmp_path, (primary, f"{rig.api_base}/{marker}"), (fallback, peer.url)), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key() + (received,) = asyncio.run( + responses_ws_peer.session(str(owned.gateway.client.base_url), key, primary, (marker,)) + ) + assert received[-1]["type"] == "response.completed", received + assert responses_ws_peer.completed(received)["status"] == "completed" + assert any(event.get("delta") == _SERVED_TEXT for event in received), received + (forwarded,) = responses_ws_peer.drain(peer.frames) + assert marker in json.dumps(forwarded), forwarded + assert rig.upstream.calls(marker) == () + + +def test_client_routing_fields_cannot_disable_an_explicit_restriction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy(scenario, rig.api_base, "restricted-chat", _single(_EMBED_ONLY)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat( + rig, + model, + marker, + original_generic_function="caller-value", + additional_drop_params=["original_generic_function"], + ) + _assert_rejected_chat(rig, response, marker, model) + + +def test_a_compaction_endpoint_declaration_keeps_the_model_available_for_compaction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy( + scenario, + rig.api_base, + "compact-declaration", + _single({"supported_endpoints": ["/responses/compact"], "enforce_supported_endpoints": True}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = rig.gateway.request("POST", "/responses/compact", {"model": model, "input": marker}) + _assert_served(rig, response, marker, "/v1/responses/compact") + + +@pytest.mark.parametrize("alias_source", ["key_alias", "key_router", "team_router"]) +def test_local_rate_limit_fallbacks_resolve_saved_aliases(rig: _Rig, tmp_path: Path, alias_source: str) -> None: + suffix: Final = uuid.uuid4().hex[:12] + names: Final = {label: f"gate-fallback-alias-{label}-{suffix}" for label in ("primary", "target", "alias")} + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["primary"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(names["target"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "fallbacks": [{names["primary"]: [names["alias"]]}], + }, + } + path: Final = tmp_path / "fallback-alias.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + aliases: Final = {names["alias"]: names["target"]} + with ( + owned_proxy_process( + rig.gateway, tmp_path, {}, config=path, workers=2, state=OwnedState.SHARED_WITH_WORKER + ) as owned, + owned.gateway.scenario() as scenario, + ): + team: Final = ( + scenario.team(router_settings={"model_group_alias": aliases}) if alias_source == "team_router" else None + ) + key: Final = scenario.key( + team_id=team, + aliases=aliases if alias_source == "key_alias" else {}, + router_settings={"model_group_alias": aliases} if alias_source == "key_router" else None, + model_rpm_limit={names["primary"]: 1}, + ) + first: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], _marker()), key=key + ) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], marker), key=key + ) + _assert_fallback_answered(rig.upstream, second, marker) diff --git a/tests/integration/security/test_api_key_secret_reference.py b/tests/integration/security/test_api_key_secret_reference.py index 0b1187d2e0e..86a7d58e0d3 100644 --- a/tests/integration/security/test_api_key_secret_reference.py +++ b/tests/integration/security/test_api_key_secret_reference.py @@ -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]]: diff --git a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py index 1c315a6e694..ff60cb9d174 100644 --- a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py +++ b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py @@ -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 [ ( diff --git a/tests/integration/spend/test_daily_activity_routes.py b/tests/integration/spend/test_daily_activity_routes.py index 520c8ccc6d8..9bfe4bb50ce 100644 --- a/tests/integration/spend/test_daily_activity_routes.py +++ b/tests/integration/spend/test_daily_activity_routes.py @@ -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, diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index e96511511d9..5c2e9cf98da 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -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( diff --git a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py index 229a9be6d20..ea0a820a705 100644 --- a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py +++ b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py @@ -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( { diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index 414bcc11602..ef050d96dc5 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -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): diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index f48495d1062..fccc11617bb 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -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.""" diff --git a/tests/unit/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py index 637a20571c8..e01c710dfdb 100644 --- a/tests/unit/litellm_core_utils/test_served_output_texts.py +++ b/tests/unit/litellm_core_utils/test_served_output_texts.py @@ -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: " @@ -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("", 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 = ( diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index d4490927e29..37377498d89 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -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 diff --git a/tests/unit/litellm_proxy_extras/test_migration_dry_run.py b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py new file mode 100644 index 00000000000..75a7fea6c89 --- /dev/null +++ b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py @@ -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"' + ) diff --git a/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py index a3fd5933a83..99a8004b0a1 100644 --- a/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py @@ -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""" diff --git a/tests/unit/llms/dashscope/test_dashscope_common_utils.py b/tests/unit/llms/dashscope/test_dashscope_common_utils.py new file mode 100644 index 00000000000..2dc4d5e5a85 --- /dev/null +++ b/tests/unit/llms/dashscope/test_dashscope_common_utils.py @@ -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" + ) diff --git a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py index 1b6eea0e4c8..27afbee3173 100644 --- a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py @@ -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() diff --git a/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py b/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py new file mode 100644 index 00000000000..decc12afd87 --- /dev/null +++ b/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py @@ -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] diff --git a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py index a913c79af31..f0dbdb74b68 100644 --- a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py @@ -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}] diff --git a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py index 15b77aaee3d..c1b4157c8ab 100644 --- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py @@ -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) diff --git a/tests/unit/llms/dashscope/test_qwen_brand_aliases.py b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py index 69ae66d3818..43ef8f90b39 100644 --- a/tests/unit/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py @@ -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"): diff --git a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py index c96fbf34fe1..9beb1faf7f4 100644 --- a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -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] == [] diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index ecb2375d495..c8d909ac2eb 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -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"], diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 4b776bfbb1a..6c82628576f 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -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.") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7fb7e1fed0..9a14222c684 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -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) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index 954c57b5cb6..a5ba9dab0bd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -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"] diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 7e5d0f6342e..6d1e5c72982 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index 01d2f7007a1..c80fa3572fd 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -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) diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 2a273034a97..e487c00f25d 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -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 == ("",) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index b892dc97e5c..835e7605955 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index e5d13a381ae..a2e509b6eb8 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -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, ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py index 9c50a6d4df6..ab4f3f7a675 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py @@ -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(): diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 8b1d7c1222a..a8706c5bccd 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -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(") diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index e3d514b2ea9..b11dcc58ef6 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -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): diff --git a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py index d38852617b9..fdf1adda1f4 100644 --- a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py +++ b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py @@ -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 diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 23f9ac0e2d8..740334252d2 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -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"] == "" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..e58a1de9eb5 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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 diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index afefae0dbb1..1db334039ec 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -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) diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 22d6316c448..c835e1e27ea 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -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) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index b21730c64f6..d0cf7be794a 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -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 diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 9e31d3774f9..8f364e51e2b 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -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 diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index acc1463b072..cfd4fe197be 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -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 diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index c57db08e976..f8375b93dad 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/tests/unit/repositories/test_daily_activity_sql.py @@ -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, '')", diff --git a/tests/unit/router_utils/test_endpoint_eligibility.py b/tests/unit/router_utils/test_endpoint_eligibility.py new file mode 100644 index 00000000000..5e5d77254d0 --- /dev/null +++ b/tests/unit/router_utils/test_endpoint_eligibility.py @@ -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 diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index 862658d9169..a2e5b498f34 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -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") diff --git a/tests/unit/test_dashscope_image_generation.py b/tests/unit/test_dashscope_image_generation.py index 960807a033b..9c65399350c 100644 --- a/tests/unit/test_dashscope_image_generation.py +++ b/tests/unit/test_dashscope_image_generation.py @@ -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", diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index ee3a70c7c1c..62510dac26d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index ba1c1eb509e..76664bdeac8 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -5,6 +5,7 @@ import copy import contextlib import contextvars import io +import itertools import json import logging import os @@ -14,6 +15,7 @@ import time from collections.abc import AsyncIterator, Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone +from functools import partial from pathlib import Path, PurePath from types import MappingProxyType from typing import Final, SupportsIndex, cast @@ -83,8 +85,10 @@ from litellm.types.utils import ( all_litellm_params, bedrock_batch_litellm_params, ) +from litellm.secret_managers.main import oidc_cache from litellm.types.videos.main import VideoObject from litellm.utils import ( + acredential_updates, _invalidate_model_cost_lowercase_map, add_custom_logger_callback_to_specific_event, check_valid_key, @@ -7607,7 +7611,6 @@ def test_unresolvable_oidc_api_key_fails_closed_naming_the_reason_without_its_ta ("reference", "kind"), [ ("oidc/env_path/MISTRAL_SA_TOKEN_PATH", "oidc/env_path"), - ("oidc/google/https://api.mistral.ai", "oidc/google"), ("oidc/github/mistral", "oidc/github"), ("oidc/env_MISTRAL_SA_TOKEN_PATH", "oidc/"), ("oidc/LIT-UNKNOWN-KIND/MISTRAL_SA_TOKEN_PATH", "oidc/"), @@ -7628,6 +7631,7 @@ def test_oidc_api_key_of_a_kind_api_key_cannot_resolve_fails_closed_as_an_authen assert not route.called assert excinfo.value.status_code == 401 assert f"{kind} secret reference" in excinfo.value.message + assert "oidc/google/" in excinfo.value.message assert reference.removeprefix(kind).lstrip("/") not in excinfo.value.message assert "MISTRAL_SA_TOKEN_PATH" not in excinfo.value.message @@ -7643,7 +7647,7 @@ async def test_router_fails_over_from_a_deployment_whose_oidc_api_key_kind_canno "model_name": "devstral", "litellm_params": { "model": "mistral/devstral-latest", - "api_key": "oidc/google/https://api.mistral.ai", + "api_key": "oidc/github/mistral", "order": 1, }, }, @@ -7682,6 +7686,372 @@ def test_an_oidc_api_key_sent_with_the_call_goes_upstream_as_written( assert sent == f"Bearer {reference}" +_GOOGLE_METADATA_IDENTITY: Final = ( + "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity" +) +_VLLM_BASE: Final = "https://vllm.example.test/v1" +_VLLM_CHAT_URL: Final = f"{_VLLM_BASE}/chat/completions" +_VLLM_MODEL: Final = "hosted_vllm/qwen3-0.6b" +_VLLM_REPLY: Final = {**_MISTRAL_REPLY, "model": "qwen3-0.6b"} +_CLOCK_START: Final = 1_900_000_000.0 +_GOOGLE_TOKEN_LIFETIME: Final = 3600 + + +def _jwt_segment(claims: Mapping[str, object]) -> str: + return base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + + +def _google_id_token(audience: str, issued_at: float) -> str: + claims: Final = {"aud": audience, "iat": int(issued_at), "exp": int(issued_at) + _GOOGLE_TOKEN_LIFETIME} + return f"{_jwt_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_jwt_segment(claims)}.signature" + + +def _google_reference() -> tuple[str, str]: + audience: Final = f"https://vllm.example.test/{uuid.uuid4().hex}" + return audience, f"oidc/google/{audience}" + + +def _metadata_identity(fetched_on: queue.SimpleQueue[int], request: httpx.Request) -> httpx.Response: + fetched_on.put(threading.get_ident()) + if request.headers.get("Metadata-Flavor") != "Google": + return httpx.Response(403) + return httpx.Response(200, text=_google_id_token(request.url.params["audience"], time.time())) + + +def _vllm_credential(api_key: str) -> CredentialItem: + return CredentialItem( + credential_name="vllm-sa", credential_values={"api_key": api_key, "api_base": _VLLM_BASE}, credential_info={} + ) + + +async def _ping_vllm_through(call_path: str, credential_name: str) -> None: + if call_path == "async": + await litellm.acompletion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name=credential_name) + return + litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name=credential_name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async"]) +async def test_named_credential_oidc_google_api_key_sends_the_metadata_token_and_refreshes_it_before_it_expires( + monkeypatch: pytest.MonkeyPatch, call_path: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + refreshed_at: Final = _CLOCK_START + _GOOGLE_TOKEN_LIFETIME - 59 + + with respx.mock(assert_all_called=True) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).mock( + side_effect=partial(_metadata_identity, queue.SimpleQueue()) + ) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + await _ping_vllm_through(call_path, "vllm-sa") + await _ping_vllm_through(call_path, "vllm-sa") + fetches_within_lifetime: Final = metadata.call_count + monkeypatch.setattr(time, "time", lambda: refreshed_at) + await _ping_vllm_through(call_path, "vllm-sa") + sent: Final = [call.request.headers["Authorization"] for call in chat.calls] + audiences: Final = [call.request.url.params["audience"] for call in metadata.calls] + + first_token: Final = f"Bearer {_google_id_token(audience, _CLOCK_START)}" + assert sent == [first_token, first_token, f"Bearer {_google_id_token(audience, refreshed_at)}"] + assert (fetches_within_lifetime, audiences) == (1, [audience, audience]) + assert litellm.credential_list[0].credential_values["api_key"] == reference + + +@pytest.mark.asyncio +async def test_router_fetches_an_oidc_google_api_key_token_off_the_event_loop_and_keeps_the_reference( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + fetched_on: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "vllm", + "litellm_params": {"model": _VLLM_MODEL, "api_base": _VLLM_BASE, "api_key": reference}, + } + ] + ) + + try: + with respx.mock(assert_all_called=True) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=partial(_metadata_identity, fetched_on)) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + await router.acompletion(model="vllm", messages=_PING) + sent: Final = chat.calls.last.request.headers["Authorization"] + finally: + router.discard() + + assert sent == f"Bearer {_google_id_token(audience, _CLOCK_START)}" + assert fetched_on.qsize() == 1 + assert fetched_on.get_nowait() != threading.get_ident() + assert router.get_model_list(model_name="vllm")[0]["litellm_params"]["api_key"] == reference + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async"]) +@pytest.mark.parametrize( + ("metadata_failure", "status_code"), + [ + (lambda: httpx.Response(500, text="metadata server unavailable"), 503), + (lambda: httpx.Response(403, text="Missing Metadata-Flavor:Google header"), 401), + (lambda: httpx.Response(200, text=""), 401), + (lambda: httpx.ConnectError("[Errno -2] Name or service not known"), 503), + (lambda: httpx.ReadTimeout("timed out"), 503), + (lambda: httpx.RemoteProtocolError("Server disconnected without sending a response."), 503), + ], + ids=["server error", "refused", "empty token", "not on GCP", "timeout", "disconnected"], +) +async def test_oidc_google_api_key_fails_closed_when_the_metadata_server_fails_and_recovers_on_the_next_call( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + call_path: str, + metadata_failure: Callable[[], httpx.Response | httpx.HTTPError], + status_code: int, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + recovered_token: Final = _google_id_token(audience, _CLOCK_START) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock( + side_effect=[metadata_failure(), httpx.Response(200, text=recovered_token)] + ) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + with caplog.at_level(logging.WARNING, logger="LiteLLM"), pytest.raises( + (litellm.AuthenticationError, litellm.ServiceUnavailableError) + ) as excinfo: + await _ping_vllm_through(call_path, "vllm-sa") + sent_while_failing: Final = chat.call_count + await _ping_vllm_through(call_path, "vllm-sa") + sent_after_recovery: Final = [call.request.headers["Authorization"] for call in chat.calls] + + assert excinfo.value.status_code == status_code + assert "api_key oidc/google secret reference could not be resolved" in excinfo.value.message + assert audience.removeprefix("https://") not in excinfo.value.message + assert reference in caplog.text + assert sent_while_failing == 0 + assert sent_after_recovery == [f"Bearer {recovered_token}"] + + +@pytest.mark.asyncio +async def test_an_oidc_google_api_key_token_cached_when_the_call_starts_is_sent_though_it_expires_mid_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + cached_token: Final = _google_id_token(audience, _CLOCK_START) + oidc_cache.set_cache(key=reference, value=cached_token, ttl=60) + clock: Final = itertools.chain([_CLOCK_START + 60], itertools.repeat(_CLOCK_START + 61)) + monkeypatch.setattr(time, "time", lambda: next(clock)) + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + fetched_on: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + + with respx.mock(assert_all_called=False) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=partial(_metadata_identity, fetched_on)) + updates: Final = await acredential_updates({"model": _VLLM_MODEL, "litellm_credential_name": "vllm-sa"}) + + assert updates["api_key"] == cached_token + assert fetched_on.empty() + + +@pytest.mark.asyncio +async def test_a_pending_oidc_google_metadata_fetch_leaves_the_default_executor_free( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + loop: Final = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + release: Final = threading.Event() + + def pending_metadata(_request: httpx.Request) -> httpx.Response: + fetch_started.put(threading.get_ident()) + release.wait(timeout=30) + return httpx.Response(500, text="metadata server unavailable") + + with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=pending_metadata) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + call: Final = asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa")) + try: + await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30) + other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=30) + finally: + release.set() + with pytest.raises(litellm.ServiceUnavailableError) as excinfo: + await call + + assert other_work == "done" + assert excinfo.value.status_code == 503 + assert chat.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "metadata_answer", + [lambda token: httpx.Response(200, text=token), lambda _token: httpx.Response(500, text="metadata server down")], + ids=["token", "outage"], +) +async def test_concurrent_calls_that_miss_the_oidc_google_token_cache_share_one_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, metadata_answer: Callable[[str], httpx.Response] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + token: Final = _google_id_token(audience, _CLOCK_START) + fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + release: Final = threading.Event() + + def held_metadata(_request: httpx.Request) -> httpx.Response: + fetch_started.put(threading.get_ident()) + release.wait(timeout=30) + return metadata_answer(token) + + with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=held_metadata) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + calls: Final = [asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa")) for _ in range(24)] + loop: Final = asyncio.get_running_loop() + try: + await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30) + fetch_started_while_held: Final = await loop.run_in_executor( + observer, partial(_next_fetch_within, fetch_started, 0.5) + ) + finally: + release.set() + outcomes: Final = await asyncio.gather(*calls, return_exceptions=True) + sent: Final = [call.request.headers["Authorization"] for call in chat.calls] + + assert fetch_started_while_held is None + if metadata.calls[0].response.status_code == 200: + assert metadata.call_count == 1 + assert outcomes == [None] * 24 + assert sent == [f"Bearer {token}"] * 24 + else: + assert [getattr(outcome, "status_code", None) for outcome in outcomes] == [503] * 24 + assert sent == [] + + +def _next_fetch_within(fetches: queue.SimpleQueue[int], seconds: float) -> int | None: + try: + return fetches.get(timeout=seconds) + except queue.Empty: + return None + + +@pytest.mark.asyncio +async def test_a_call_an_oidc_google_api_key_may_not_go_to_is_refused_without_a_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + + with respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(500) + with pytest.raises(litellm.AuthenticationError) as excinfo: + await acredential_updates( + { + "model": _VLLM_MODEL, + "litellm_credential_name": "vllm-sa", + "api_base": "https://elsewhere.example.test/v1", + } + ) + + assert "only to that credential's own api_base" in excinfo.value.message + assert metadata.call_count == 0 + + +async def _pass_through_chat_call(call_path: str, reference: str) -> None: + pass_through: Final = {"method": "POST", "endpoint": "v1/chat/completions", "json": {"messages": _PING}} + if call_path == "sync": + litellm.llm_passthrough_route(model=_VLLM_MODEL, litellm_credential_name="vllm-sa", **pass_through) + return + router: Final = litellm.Router( + model_list=[ + { + "model_name": "vllm-model", + "litellm_params": {"model": _VLLM_MODEL, "api_base": _VLLM_BASE, "api_key": reference}, + } + ], + num_retries=0, + ) + await router.allm_passthrough_route(model="vllm-model", **pass_through) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async router"]) +async def test_a_pass_through_call_with_an_oidc_google_api_key_is_refused_without_a_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, call_path: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + + with respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond( + 200, text=_google_id_token(audience, _CLOCK_START) + ) + vllm: Final = upstream.route(host="vllm.example.test").respond(200, json=_VLLM_REPLY) + with pytest.raises(litellm.AuthenticationError) as excinfo: + await _pass_through_chat_call(call_path, reference) + + assert "pass-through routes do not fetch Google ID tokens" in excinfo.value.message + assert metadata.call_count == 0 + assert vllm.call_count == 0 + + +@pytest.mark.parametrize( + ("configured", "bound"), + [(None, 3.0), ("1.5", 1.5), ("three seconds", 3.0), ("0", 3.0), ("-2", 3.0)], + ids=["unset", "configured", "malformed", "zero", "negative"], +) +def test_oidc_google_api_key_metadata_fetch_waits_at_most_gce_metadata_timeout_seconds( + monkeypatch: pytest.MonkeyPatch, configured: str | None, bound: float +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if configured is None: + monkeypatch.delenv("GCE_METADATA_TIMEOUT", raising=False) + else: + monkeypatch.setenv("GCE_METADATA_TIMEOUT", configured) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + token: Final = _google_id_token(audience, _CLOCK_START) + + with respx.mock(assert_all_called=True) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(200, text=token) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name="vllm-sa") + sent: Final = chat.calls.last.request.headers["Authorization"] + waits: Final = metadata.calls.last.request.extensions["timeout"] + + assert sent == f"Bearer {token}" + assert waits == {"connect": bound, "read": bound, "write": bound, "pool": bound} + + +def test_an_oidc_google_api_key_sent_with_the_call_goes_upstream_as_written_without_a_metadata_fetch() -> None: + _, reference = _google_reference() + + with respx.mock(assert_all_called=True) as upstream: + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + litellm.completion(model=_VLLM_MODEL, api_base=_VLLM_BASE, messages=_PING, api_key=reference) + sent: Final = chat.calls.last.request.headers["Authorization"] + + assert sent == f"Bearer {reference}" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model_info", "expected"), @@ -11209,6 +11579,7 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + @pytest.mark.parametrize( ("provider", "expected_warnings"), (("streamlake", 1), ("veniceai", 1), ("my-own-llm", 0)), @@ -11227,3 +11598,34 @@ def test_custom_provider_named_like_a_built_in_warns_once( warnings: Final = [r for r in caplog.records if f"provider '{provider}'" in r.getMessage()] assert len(warnings) == expected_warnings assert provider in litellm._custom_providers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred", (False, True)) +async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronously(deferred: bool) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.utils import _schedule_async_success_logging + + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="acompletion", + start_time=now, + litellm_call_id="shutdown-enqueue", + function_id="shutdown-enqueue", + ) + logging_obj.defer_async_logging = deferred + _schedule_async_success_logging(logging_obj, ModelResponse(), now, now, False) + if deferred: + assert not logging_obj.async_success_log_enqueued + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + assert "has_logged_async_success" not in logging_obj.model_call_details + await drain_logging_tasks() + assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" + await GLOBAL_LOGGING_WORKER.flush() + assert logging_obj.model_call_details["has_logged_async_success"] is True diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx index a9c87471511..2b4b69ffe7b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx @@ -493,6 +493,26 @@ describe("TopKeyView", () => { }); }); + it.each([ + ["raw:unattributed", "unattributed"], + ["raw:raw:unattributed", "raw:unattributed"], + ["raw:raw:raw:key", "raw:raw:key"], + ])("opens the stored key behind reporting ID %s", async (bucket, key) => { + mockKeyInfoV1Call.mockResolvedValue({ key }); + mockTransformKeyInfo.mockReturnValue({ token: key } as KeyResponse); + render( + , + ); + + fireEvent.click(screen.getByRole("button", { name: bucket })); + + expect(await screen.findByText(`Key Info View for ${key}`)).toBeInTheDocument(); + expect(mockKeyInfoV1Call).toHaveBeenCalledWith("test-token", key); + }); + it("should only look up keys that still exist in the database, from both the table and the chart", async () => { mockKeyInfoV1Call.mockResolvedValue({ key: "info" }); mockTransformKeyInfo.mockReturnValue({ transformed: "data" } as unknown as KeyResponse); @@ -503,6 +523,7 @@ describe("TopKeyView", () => { {...baseProps} topKeys={[ { api_key: "session-key", key_alias: null, user: "alice@example.com", key_exists: false, spend: 100 }, + { api_key: "unattributed", key_alias: "Unattributed", user: null, key_exists: false, spend: 25 }, { api_key: "stored-key", key_alias: "Stored", user: null, key_exists: true, spend: 50 }, ]} />, @@ -510,11 +531,14 @@ describe("TopKeyView", () => { expect(screen.getByRole("button", { name: "stored-key" })).toBeInTheDocument(); expect(screen.queryByRole("button", { name: "session-key" })).not.toBeInTheDocument(); + expect(screen.getByText("Unattributed")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "unattributed" })).not.toBeInTheDocument(); + await user.click(screen.getByText("unattributed")); await user.click(screen.getByText("session-key")); await user.click(screen.getByRole("radio", { name: "Chart View" })); const bars = chartBars(container); - expect(bars).toHaveLength(2); + expect(bars).toHaveLength(3); bars.forEach((bar) => fireEvent.click(bar)); expect(await screen.findByText("Key Info View for stored-key")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index 90a26919e01..291bcb34c73 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -11,6 +11,7 @@ import { transformKeyInfo } from "../../../key_team_helpers/transform_key_info"; import { keyInfoV1Call } from "../../../networking"; import KeyInfoView from "../../../templates/key_info_view"; import { TagUsage } from "../../types"; +import { storedKeyId } from "../../keyActivityData"; const TOP_KEYS_LIMITS = [5, 10, 25, 50] as const; const LIMIT_OPTIONS = TOP_KEYS_LIMITS.map((limit) => ({ value: String(limit), label: String(limit) })); @@ -70,11 +71,12 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals if (!accessToken || !canOpenKeyInfo(item)) return; try { - const keyInfo = await keyInfoV1Call(accessToken, item.api_key); + const keyId = storedKeyId(item.api_key); + const keyInfo = await keyInfoV1Call(accessToken, keyId); const transformedKeyData = transformKeyInfo(keyInfo); setKeyData(transformedKeyData); - setSelectedKey(item.api_key); + setSelectedKey(keyId); setIsModalOpen(true); // Open modal when key is clicked } catch (error) { console.error("Error fetching key info:", error); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts index b35ea72896a..a5db5db14cc 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts @@ -7,7 +7,13 @@ import { type KeySpendActivityRow, toDailyData, } from "./dailyActivityApi"; -import { keyDetailFromResponse, mergeKeyActivityPages, overallUsageMetrics } from "./keyActivityData"; +import { + keyDetailFromResponse, + mergeKeyActivityPages, + overallUsageMetrics, + reportingKeyId, + storedKeyId, +} from "./keyActivityData"; const completeMetrics: components["schemas"]["SpendMetrics"] = { api_requests: 2, @@ -80,6 +86,50 @@ const pageRow = (api_key: string): KeySpendActivityRow => ({ }); describe("key activity data", () => { + it.each([ + ["", "unattributed"], + ["unattributed", "raw:unattributed"], + ["raw:key", "raw:raw:key"], + ["raw:raw:key", "raw:raw:raw:key"], + ["ordinary-hash", "ordinary-hash"], + [" ", " "], + ])("round trips stored key %s through reporting", (key, bucket) => { + expect(reportingKeyId(key)).toBe(bucket); + expect(storedKeyId(bucket)).toBe(key); + }); + + it("retains unattributed totals and unknown models in key details", () => { + const unattributed = { + metrics: completeMetrics, + metadata: { key_alias: "Unattributed", key_exists: false }, + }; + const response: DailyActivityAggregatedResponse = { + ...aggregatedResponse, + results: [ + { + date: "2026-09-27", + metrics: completeMetrics, + breakdown: { + api_keys: { unattributed }, + models: { + "unknown-model": { + metrics: completeMetrics, + metadata: {}, + api_key_breakdown: { unattributed }, + }, + }, + }, + }, + ], + }; + + const detail = keyDetailFromResponse(response, "unattributed", []); + expect(detail?.total_requests).toBe(2); + expect(detail?.total_spend).toBe(1.25); + expect(detail?.top_models.map((model) => model.model)).toEqual(["unknown-model"]); + expect(detail?.daily_data[0].metrics.api_requests).toBe(2); + }); + it("builds overall totals from metadata and daily data from complete daily metrics", () => { const summary = overallUsageMetrics( toDailyData(aggregatedResponse), diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts index 4b2e8c573f2..2d2977058a5 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts @@ -4,6 +4,16 @@ import type { DailyActivityAggregatedResponse, DailyActivityMetadata, KeySpendAc import { toDailyData } from "./dailyActivityApi"; import type { ModelActivityData } from "./types"; +export const reportingKeyId = (key: string): string => { + if (!key) return "unattributed"; + return key === "unattributed" || key.startsWith("raw:") ? `raw:${key}` : key; +}; + +export const storedKeyId = (bucket: string): string => { + if (bucket === "unattributed") return ""; + return bucket.startsWith("raw:") ? bucket.slice(4) : bucket; +}; + export const overallUsageMetrics = ( results: ReturnType, metadata: DailyActivityMetadata, diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx index 3567388de05..40ab0e8e77c 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx @@ -131,6 +131,18 @@ describe("KeySavingsTab", () => { expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests"); }); + it.each([ + ["unattributed", "raw:unattributed"], + ["raw:unattributed", "raw:raw:unattributed"], + ["raw:raw:key", "raw:raw:raw:key"], + ])("requests savings for the real key %s", (keyToken, bucket) => { + const hook = scopedRange().mockReturnValue(mockActivity()); + + renderTab({ keyToken, userRole: "Admin" }); + + expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: bucket }, activity); + }); + it("reads the whole key for a proxy admin, with no scope note to contradict it", () => { const hook = scopedRange().mockReturnValue(mockActivity()); diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx index d1395e0be51..03aced7358e 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx @@ -1,6 +1,7 @@ "use client"; import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; +import { reportingKeyId } from "@/components/UsagePage/keyActivityData"; import { hasProxyWideSpendView, spendScopeUserId } from "@/utils/roles"; import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; @@ -15,7 +16,7 @@ interface KeySavingsTabProps { const KeySavingsTab = ({ accessToken, keyToken, userId, userRole, activity }: KeySavingsTabProps) => (