From 94032014df175c3ec3735ded5144e91b5576547b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:24:59 -0700 Subject: [PATCH 001/127] fix(proxy): coordinate v2 migration startup and add container regression CI --- .circleci/config.yml | 137 +++++++++++- .circleci/scripts/run_migration_tests.py | 113 ++++++++++ .../litellm_proxy_extras/migration_lock.py | 89 ++++++++ .../migration_recovery.py | 158 ++++++++++++++ .../litellm_proxy_extras/prisma_toolchain.py | 5 + .../litellm_proxy_extras/utils.py | 192 +++++++++++------ .../tests/test_setup_database_fail_fast.py | 109 ++++------ tests/e2e/CLAUDE.md | 2 + tests/e2e/conftest.py | 10 +- tests/e2e/migrations/__init__.py | 0 tests/e2e/migrations/checks.py | 130 +++++++++++ tests/e2e/migrations/conftest.py | 62 ++++++ tests/e2e/migrations/containers.py | 203 ++++++++++++++++++ tests/e2e/migrations/database.py | 135 ++++++++++++ tests/e2e/migrations/startup_models.py | 25 +++ tests/e2e/migrations/test_legacy.py | 84 ++++++++ tests/e2e/migrations/test_pooling.py | 135 ++++++++++++ tests/e2e/migrations/test_recovery.py | 183 ++++++++++++++++ tests/e2e/migrations/test_startup.py | 100 +++++++++ .../test_litellm_proxy_extras_utils.py | 191 ++++++++++++---- .../test_migration_ci.py | 36 ++++ 21 files changed, 1909 insertions(+), 190 deletions(-) create mode 100644 .circleci/scripts/run_migration_tests.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_lock.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py create mode 100644 tests/e2e/migrations/__init__.py create mode 100644 tests/e2e/migrations/checks.py create mode 100644 tests/e2e/migrations/conftest.py create mode 100644 tests/e2e/migrations/containers.py create mode 100644 tests/e2e/migrations/database.py create mode 100644 tests/e2e/migrations/startup_models.py create mode 100644 tests/e2e/migrations/test_legacy.py create mode 100644 tests/e2e/migrations/test_pooling.py create mode 100644 tests/e2e/migrations/test_recovery.py create mode 100644 tests/e2e/migrations/test_startup.py create mode 100644 tests/proxy_migration_tests/test_migration_ci.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 32d2cf0390c..f6f31651306 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1,10 +1,32 @@ version: 2.1 +parameters: + run_migration_tests: + type: boolean + default: false + migration_candidate_image: + type: string + default: "" + migration_source_sha: + type: string + default: "" orbs: codecov: codecov/codecov@4.0.1 node: circleci/node@5.1.0 # Add this line to declare the node orb win: circleci/windows@5.0 # Add Windows orb commands: + checkout_migration_source: + steps: + - run: + name: Select the requested migration test revision + environment: + MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >> + command: | + if [ -n "$MIGRATION_SOURCE_SHA" ]; then + [[ "$MIGRATION_SOURCE_SHA" =~ ^[0-9a-f]{40}$ ]] || exit 1 + git fetch origin "$MIGRATION_SOURCE_SHA" + git checkout --detach "$MIGRATION_SOURCE_SHA" + fi skip_if_unrelated_changes: parameters: category: @@ -2853,14 +2875,25 @@ jobs: working_directory: ~/project steps: - checkout + - checkout_migration_source - skip_if_unrelated_changes - run: name: Build Docker image + environment: + MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >> command: | - docker build \ - -t litellm-docker-database:ci \ - -f docker/Dockerfile.database . + if [ -n "$MIGRATION_CANDIDATE_IMAGE" ]; then + [[ "$MIGRATION_CANDIDATE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+@sha256:[0-9a-f]{64}$ ]] || exit 1 + docker pull "$MIGRATION_CANDIDATE_IMAGE" + docker tag "$MIGRATION_CANDIDATE_IMAGE" litellm-docker-database:ci + else + docker build \ + --label org.opencontainers.image.revision="$(git rev-parse HEAD)" \ + -t litellm-docker-database:ci \ + -f docker/Dockerfile.database . + fi + python3 .circleci/scripts/run_migration_tests.py record-image - run: name: Save Docker image to workspace root @@ -2871,6 +2904,79 @@ jobs: root: . paths: - litellm-docker-database.tar.zst + - migration-image.json + + migration_startup_tests: + parameters: + suite: + type: enum + enum: [startup, recovery, legacy] + machine: + image: ubuntu-2204:2024.04.1 + resource_class: large + working_directory: ~/project + environment: + LITELLM_MIGRATION_TESTS: "1" + LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci + MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres + MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres + MIGRATION_TEST_OUTPUT: /tmp/migration-results + PYTHONPATH: tests/e2e + steps: + - checkout + - checkout_migration_source + - install_uv + - install_rust + - restore_cache: + keys: + - v1-uv-cache-{{ checksum "uv.lock" }} + - run: + name: Install test dependencies + command: uv sync --frozen --all-groups --all-extras --python 3.12 + - attach_workspace: + at: ~/project + - run: + name: Load the shared candidate and start PostgreSQL + command: | + zstd -d litellm-docker-database.tar.zst --stdout | docker load + docker run -d --name migration-postgres \ + -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres \ + -p 5432:5432 \ + postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Run migration startup regressions + environment: + MIGRATION_TEST_SUITE: << parameters.suite >> + MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >> + command: | + mkdir -p /tmp/migration-results + uv run --no-sync python .circleci/scripts/run_migration_tests.py + no_output_timeout: 15m + - store_test_results: + path: /tmp/migration-results/junit + - run: + name: Package migration diagnostics + when: always + command: | + mkdir -p /tmp/migration-artifacts + if [ -d /tmp/migration-results ]; then + tar -czf /tmp/migration-artifacts/diagnostics.tar.gz -C /tmp/migration-results . + if [ -f /tmp/migration-results/verdict.json ]; then + cp /tmp/migration-results/verdict.json /tmp/migration-artifacts/verdict.json + fi + fi + - store_artifacts: + path: /tmp/migration-artifacts + destination: migration-results + - run: + name: Remove migration test containers + when: always + command: | + docker ps -aq --filter label=litellm-migration-test=true | xargs -r docker rm -f + docker rm -f migration-postgres || true test_bad_database_url: machine: @@ -2915,7 +3021,32 @@ jobs: fi workflows: + migration_startup: + when: << pipeline.parameters.run_migration_tests >> + jobs: &migration_jobs + - build_docker_database_image + - migration_startup_tests: + name: migration-startup + suite: startup + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-recovery + suite: recovery + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-legacy-and-pooling + suite: legacy + requires: [build_docker_database_image] + migration_startup_scheduled: + triggers: + - schedule: + cron: "17 0,6,12,18 * * *" + filters: + branches: + only: litellm_internal_staging + jobs: *migration_jobs build_and_test: + unless: << pipeline.parameters.run_migration_tests >> jobs: - using_litellm_on_windows: filters: &main_branches diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py new file mode 100644 index 00000000000..56029c406fb --- /dev/null +++ b/.circleci/scripts/run_migration_tests.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +from pathlib import Path +from typing import Final +from xml.etree import ElementTree + +SUITES: Final = { + "startup": (("test_startup.py",), 12), + "recovery": (("test_recovery.py",), 15), + "legacy": (("test_legacy.py", "test_pooling.py"), 11), +} + + +def successful_junit(path: Path, expected: int, exit_code: int) -> bool: + if exit_code != 0 or not path.is_file(): + return False + try: + root: Final = ElementTree.parse(path).getroot() + except ElementTree.ParseError: + return False + cases: Final = tuple(root.iter("testcase")) + identities: Final = frozenset((case.get("classname"), case.get("name")) for case in cases) + return len(cases) == len(identities) == expected and all( + not any(case.find(tag) is not None for tag in ("failure", "error", "skipped")) for case in cases + ) + + +def output(*command: str) -> str: + return subprocess.check_output(command, text=True, timeout=90).strip() + + +def record_image() -> None: + source: Final = output("git", "rev-parse", "HEAD") + image: Final = output("docker", "image", "inspect", "litellm-docker-database:ci", "--format", "{{.Id}}") + revision: Final = output( + "docker", + "image", + "inspect", + "litellm-docker-database:ci", + "--format", + '{{index .Config.Labels "org.opencontainers.image.revision"}}', + ) + assert re.fullmatch(r"[0-9a-f]{40}", source), "Invalid source revision" + assert revision == source, "Candidate image revision differs from the tested source" + Path("migration-image.json").write_text( + json.dumps( + { + "source_sha": source, + "image_id": image, + "candidate_image": os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), + } + ) + ) + + +def main() -> int: + suite: Final = os.environ["MIGRATION_TEST_SUITE"] + files, expected = SUITES[suite] + metadata: Final = json.loads(Path("migration-image.json").read_text()) + assert metadata["source_sha"] == output("git", "rev-parse", "HEAD"), "Image and test source revisions differ" + assert metadata["image_id"] == output( + "docker", "image", "inspect", os.environ["LITELLM_MIGRATION_TEST_IMAGE"], "--format", "{{.Id}}" + ), "Loaded image differs from the build output" + assert metadata["candidate_image"] == os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), "Wrong release candidate" + destination: Final = Path(os.environ["MIGRATION_TEST_OUTPUT"]) + junit: Final = destination / "junit" / "results.xml" + junit.parent.mkdir(parents=True, exist_ok=True) + result: Final = subprocess.run( + ( + sys.executable, + "-m", + "pytest", + *(f"tests/e2e/migrations/{name}" for name in files), + "-vv", + "--tb=short", + "--durations=10", + f"--junitxml={junit}", + "-o", + "addopts=", + "--reruns=0", + ), + check=False, + timeout=1200, + ) + passed: Final = successful_junit(junit, expected, result.returncode) + (destination / "verdict.json").write_text( + json.dumps( + { + **metadata, + "suite": suite, + "expected_cases": expected, + "passed": passed, + "pytest_exit_code": result.returncode, + "test_revision": metadata["source_sha"], + "workflow_id": os.environ.get("CIRCLE_WORKFLOW_ID", ""), + "job_number": os.environ.get("CIRCLE_BUILD_NUM", ""), + }, + indent=2, + ) + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + if sys.argv[1:] == ["record-image"]: + record_image() + else: + raise SystemExit(main()) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py new file mode 100644 index 00000000000..e4ccbe585a9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -0,0 +1,89 @@ +import random +import time +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.prisma_toolchain import MIGRATION_LOCK_TIMEOUT_ENV_VAR, migration_lock_timeout + +MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big") + +if TYPE_CHECKING: + import psycopg + + +def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]: + database_url: Final = environment.get("DATABASE_URL") + direct_url: Final = environment.get("DIRECT_URL") + if not database_url or not direct_url: + return environment + schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public") + direct: Final = urlsplit(direct_url) + parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema") + return { + **environment, + "DATABASE_URL": urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema))))), + } + + +@dataclass(frozen=True, slots=True) +class _LockResult: + acquired: bool + + +def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = MIGRATION_LOCK_KEY) -> bool: + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(_LockResult)) as cursor: + row: Final = cursor.execute("SELECT pg_try_advisory_xact_lock(%s) AS acquired", (key,)).fetchone() + return row is not None and row.acquired + + +@dataclass(frozen=True, slots=True) +class MigrationCoordinator: + connection: "psycopg.Connection[tuple[object, ...]]" + + def check_connection(self) -> None: + self.connection.execute("SELECT 1") + + def acquire_prisma_lock(self) -> None: + deadline: Final = time.monotonic() + migration_lock_timeout() + while time.monotonic() < deadline: + if _try_lock(self.connection, 72707369): + return + time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) + raise RuntimeError( + "Timed out waiting for Prisma's lock to recover migration history. LiteLLM startup has stopped. " + "Another migration or a pooled database session may still hold the lock. Check the database lock holder. " + "When using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler." + ) + + +@contextmanager +def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]: + import psycopg + + wait_seconds: Final = migration_lock_timeout() + deadline: Final = time.monotonic() + wait_seconds + try: + with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection: + coordinator: Final = MigrationCoordinator(connection) + logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds) + while time.monotonic() < deadline: + with connection.transaction(): + if _try_lock(connection): + logger.info("Acquired the v2 migration coordinator lock") + + yield coordinator + coordinator.check_connection() + return + time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) + except psycopg.Error as exc: + raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc + raise RuntimeError( + f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. " + f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}." + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py new file mode 100644 index 00000000000..9202317c776 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -0,0 +1,158 @@ +import hashlib +import subprocess +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Final +from uuid import uuid4 + +from litellm_proxy_extras import prisma_toolchain +from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.migration_lock import MigrationCoordinator + +if TYPE_CHECKING: + import psycopg + + +@dataclass(frozen=True, slots=True) +class MigrationProgress: + checksum: str + applied_steps_count: int + logs: str + id: str = "" + finished: bool = False + + def confirms_completion(self, script: bytes) -> bool: + return ( + self.applied_steps_count == 1 + and not self.logs.strip() + and self.checksum == hashlib.sha256(script).hexdigest() + ) + + +def _migration_records( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migration: Path +) -> tuple[MigrationProgress, ...]: + from psycopg import sql + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(MigrationProgress)) as cursor: + records: Final = cursor.execute( + sql.SQL( + "SELECT id, checksum, applied_steps_count, coalesce(logs, '') AS logs, " + "finished_at IS NOT NULL AS finished FROM {} " + "WHERE migration_name = %s AND rolled_back_at IS NULL" + ).format(sql.Identifier(schema, "_prisma_migrations")), + (migration.parent.name,), + ).fetchall() + return tuple(records) + + +def recover_completed_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool: + """Finish a proven successful row without erasing its durable completion evidence. + + The caller commits this checkpoint before running another Prisma command. + """ + from psycopg import sql + + coordinator.acquire_prisma_lock() + records: Final = _migration_records(coordinator.connection, schema, migration) + unfinished: Final = tuple(record for record in records if not record.finished) + script: Final = migration.read_bytes() + if not unfinished: + return any(record.checksum == hashlib.sha256(script).hexdigest() for record in records) + if len(unfinished) != 1 or not unfinished[0].confirms_completion(script): + return False + progress: Final = unfinished[0] + result: Final = coordinator.connection.execute( + sql.SQL( + "UPDATE {} SET finished_at = current_timestamp " + "WHERE id = %s AND checksum = %s AND applied_steps_count = 1 " + "AND finished_at IS NULL AND rolled_back_at IS NULL AND coalesce(logs, '') = %s" + ).format(sql.Identifier(schema, "_prisma_migrations")), + (progress.id, progress.checksum, progress.logs), + ) + if result.rowcount != 1: + raise RuntimeError("Could not complete the confirmed migration history row; retry startup.") + logger.info("Completed migration %s using its successful SQL step and matching checksum", migration.parent.name) + return True + + +def migration_files(directory: Path) -> tuple[tuple[str, str], ...]: + return tuple( + (path.parent.name, hashlib.sha256(path.read_bytes()).hexdigest()) + for path in sorted((directory / "migrations").glob("*/migration.sql")) + ) + + +def baseline_current_schema( + coordinator: MigrationCoordinator, + schema: str, + migrations_dir: Path, + prisma_command: str, + prisma_env: Mapping[str, str], +) -> 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() + ): + raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.") + + coordinator.acquire_prisma_lock() + existing: Final = coordinator.connection.execute( + "SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(coordinator.connection),) + ).fetchone() + if existing is not None and existing[0] is not None: + return + try: + prisma_toolchain.run_prisma( + ( + prisma_command, + "migrate", + "diff", + "--from-schema-datasource", + str(migrations_dir / "schema.prisma"), + "--to-schema-datamodel", + str(migrations_dir / "schema.prisma"), + "--exit-code", + ), + timeout=prisma_toolchain.prisma_command_timeout(), + env=prisma_env, + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + "Cannot automatically baseline this database: its schema has not been verified to match this build. " + "Establish the existing migration history before retrying. No schema reconciliation was performed. " + "If using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler. " + f"Schema verification detail: {exc.stderr}" + ) from exc + + coordinator.check_connection() + ledger: Final = sql.Identifier(schema, "_prisma_migrations") + coordinator.connection.execute( + sql.SQL( + "CREATE TABLE {} (id varchar(36) PRIMARY KEY NOT NULL, checksum varchar(64) NOT NULL, " + "finished_at timestamptz, migration_name varchar(255) NOT NULL, logs text, rolled_back_at timestamptz, " + "started_at timestamptz NOT NULL DEFAULT now(), applied_steps_count integer NOT NULL DEFAULT 0)" + ).format(ledger) + ) + with coordinator.connection.cursor() as cursor: + cursor.executemany( + sql.SQL( + "INSERT INTO {} (id, checksum, migration_name, logs, started_at, finished_at) " + "VALUES (%s, %s, %s, '', current_timestamp, current_timestamp)" + ).format(ledger), + tuple((str(uuid4()), checksum, name) for name, checksum in migrations), + ) + logger.warning( + "Legacy migration history was missing. The existing Prisma schema matches this build; " + "adopted %s packaged migrations as a baseline. No schema changes were applied, and " + "historical data backfills were not replayed or verified. Continuing startup; " + "review any feature-specific backfill requirements.", + len(migrations), + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 9cd48fcf11a..07f83f76d2b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -59,6 +59,7 @@ except ImportError: PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT" +MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT" NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 @@ -106,6 +107,10 @@ def prisma_command_timeout() -> float: ) +def migration_lock_timeout() -> float: + return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0) + + def prisma_bootstrap_timeout() -> float: """Seconds the one-time Node toolchain install may run for.""" return _timeout_from_env( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2145f891318..2749db5d754 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -6,6 +6,7 @@ import shutil import subprocess import tempfile import time +from collections.abc import Callable from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Final, Optional @@ -78,15 +79,10 @@ MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @dataclass(frozen=True) class _MigrateAttemptBudget: - """Retries left, and the recoveries already run. - - A recovery that lands something new costs nothing, so a database full of - objects `prisma db push` created works through them one per pass. Anything - that made no progress spends an attempt, so a stuck run still gives up. - """ + """Independent bounds for failed attempts and Prisma lock contention.""" attempts_left: int - recoveries: frozenset[str] = frozenset() + contention_seconds_left: float = 600.0 @property def exhausted(self) -> bool: @@ -99,10 +95,14 @@ class _MigrateAttemptBudget: def spend(self) -> "_MigrateAttemptBudget": return replace(self, attempts_left=self.attempts_left - 1) - def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget": - if recovery in self.recoveries: - return self.spend() - return replace(self, recoveries=self.recoveries | {recovery}) + def after_contention(self, elapsed: float) -> "_MigrateAttemptBudget": + remaining: Final = self.contention_seconds_left - elapsed + if remaining <= 0: + raise RuntimeError( + "Timed out waiting for Prisma's migration advisory lock. Check the running migration " + "or increase LITELLM_MIGRATION_LOCK_TIMEOUT." + ) + return replace(self, contention_seconds_left=remaining) _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) @@ -836,12 +836,47 @@ class ProxyExtrasDBManager: @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: + if not use_migrate: + return ProxyExtrasDBManager._run_database_v2(False) + from litellm_proxy_extras.migration_lock import migration_environment, migration_lock + from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration + + database_url: Final = os.environ.get("DATABASE_URL") + if not database_url: + raise RuntimeError("DATABASE_URL is required for v2 migrations") + lock_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" + + def recover_completed(name: str) -> bool: + if Path(name).name != name or "\\" in name: + return False + migration: Final = Path(os.getcwd()) / "migrations" / name / "migration.sql" + if not migration.is_file(): + return False + with migration_lock(lock_url) as coordinator: + return recover_completed_migration(coordinator, schema, migration) + + def baseline_existing(migrations_dir: str) -> None: + with migration_lock(lock_url) as coordinator: + baseline_current_schema( + coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env()) + ) + + while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): + continue + return True + + @staticmethod + def _run_database_v2( + use_migrate: bool, + recover_completed: Callable[[str], bool] = lambda name: False, + baseline_existing: "Callable[[str], None] | None" = None, + ) -> bool: """ v2 migration resolver (opt-in via --use_v2_migration_resolver). - Runs `prisma migrate deploy` and handles standard recovery paths - (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a - concurrent migrate deploy). Critically, it does + Runs `prisma migrate deploy`, baselines verified existing schemas, + and recovers confirmed SQL completion or reported deadlocks. It does NOT call `_resolve_all_migrations` — the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contended for the same DB during rolling deploys. @@ -850,10 +885,9 @@ class ProxyExtrasDBManager: is logged as a warning, not a fatal error — users whose DBs got into weird shapes from the old thrashing should still be able to start. - The retry budget only counts attempts that made no progress: see - _MigrateAttemptBudget. + False requests a committed recovery checkpoint and another deploy + pass. True means every pending migration is complete. """ - schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" migrations_dir = ProxyExtrasDBManager._get_prisma_dir() if not use_migrate: @@ -886,14 +920,22 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) deploy_timeout = prisma_migrate_deploy_timeout() - budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS) + from litellm_proxy_extras.migration_lock import migration_environment, migration_lock_timeout + + migration_env: Final = migration_environment(_get_prisma_env()) + + budget = _MigrateAttemptBudget( + attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS, + contention_seconds_left=migration_lock_timeout(), + ) try: while not budget.exhausted: + attempt_started = time.monotonic() try: result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=deploy_timeout, - env=_get_prisma_env(), + env=migration_env, ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") return True @@ -909,8 +951,16 @@ class ProxyExtrasDBManager: next_budget = budget.spend() except subprocess.CalledProcessError as e: + if "P3005" in (e.stderr or "") and baseline_existing is not None: + baseline_existing(migrations_dir) + return False + failed_migration = ProxyExtrasDBManager._v2_failed_migration_name(e.stderr or "") + if failed_migration and recover_completed(failed_migration): + return False next_budget = ProxyExtrasDBManager._budget_after_deploy_failure( - e, budget, schema_path + e, + budget, + time.monotonic() - attempt_started, ) if next_budget.attempts_left < budget.attempts_left: @@ -919,19 +969,41 @@ class ProxyExtrasDBManager: raise RuntimeError( f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} " - "attempts that made no progress (timeouts, deadlock retries, or a " - "recovery that had already run once). Check database connectivity, " + "attempts that made no progress (timeouts or deadlock retries). Check database connectivity, " "load, and _prisma_migrations ledger state, and raise " f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) finally: os.chdir(original_dir) + @staticmethod + def _v2_failed_migration_name(stderr: str) -> "str | None": + if "P3009" in stderr: + match = re.search(r"`(\d+_[^`\r\n]+)`", stderr) + return match.group(1) if match else None + if "P3018" in stderr: + match = re.search(r"Migration name: (\d+_[^\r\n]+)", stderr) + return match.group(1) if match else None + return None + + @staticmethod + def _v2_roll_back_migration_best_effort(migration_name: str) -> None: + from litellm_proxy_extras.migration_lock import migration_environment + + try: + prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + timeout=prisma_command_timeout(), + env=migration_environment(_get_prisma_env()), + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + @staticmethod def _budget_after_deploy_failure( error: subprocess.CalledProcessError, budget: "_MigrateAttemptBudget", - schema_path: str, + attempt_seconds: float = 0.0, ) -> "_MigrateAttemptBudget": """Recover from one failed `prisma migrate deploy`, and price the pass. @@ -940,37 +1012,35 @@ class ProxyExtrasDBManager: """ stderr = error.stderr or "" - if "P3005" in stderr and "database schema is not empty" in stderr: - logger.info("Schema exists but no migrations ledger — creating baseline") - if ProxyExtrasDBManager._create_baseline_migration(schema_path): - return budget.after_recovery("baseline") - return budget.spend() - if "P3009" in stderr: - migration_match = re.search(r"`(\d+_\S+?)`", stderr) - if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): - name = migration_match.group(1) - logger.info( - f"Migration {name} failed idempotently — marking applied and retrying" - ) - ProxyExtrasDBManager._mark_migration_applied(name) - return budget.after_recovery(f"resolved:{name}") - if migration_match: - migration_name = migration_match.group(1) + migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) + if migration_name: ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) - if ledger_logs is not None and ( - ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs - ): + if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs: logger.info( "Migration %s failed in a concurrent migrate deploy " "deadlock race, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) return budget.spend() raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" + "Migration completion could not be verified. LiteLLM startup has stopped.\n\n" + f"Prisma migration history (migration name and start time):\n{stderr}\n\n" + "A migration has a start record but no successful completion record. " + "LiteLLM cannot determine whether its SQL committed from this record alone. " + "Startup stopped to avoid repeating or skipping database changes.\n\n" + "Before resolving, stop other migration runners and inspect _prisma_migrations, " + "the named migration.sql from this build, database logs, and the actual database objects and data. " + "Use the same database and this build's schema and migration files for recovery:\n" + "- Only after verifying every migration change is present, run " + "prisma migrate resolve --applied , then retry startup.\n" + "- Only after verifying no migration changes remain (or fully undoing partial changes), run " + "prisma migrate resolve --rolled-back , then retry startup. " + "This command updates history; it does not undo SQL.\n" + "Replace with the reported name. If the outcome remains uncertain, " + "leave migration history unchanged and contact your database administrator. " + "Repeated restarts alone will not resolve this state." ) from error if "P3018" in stderr: @@ -981,25 +1051,13 @@ class ProxyExtrasDBManager: f"and retry.\n\nPrisma error:\n{stderr}" ) from error - migration_match = re.search(r"Migration name: (\d+_\S+)", stderr) - if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): - name = migration_match.group(1) + migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) + if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - f"Migration {name} SQL hit idempotent error — marking applied and retrying" - ) - ProxyExtrasDBManager._mark_migration_applied(name) - return budget.after_recovery(f"resolved:{name}") - - if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: - logger.info( - "Migration %s deadlocked against a concurrent " - "migrate deploy, rolling its ledger row back " - "and retrying", - migration_match.group(1), - ) - ProxyExtrasDBManager._roll_back_migration_best_effort( - migration_match.group(1) + "Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying", + migration_name, ) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) return budget.spend() raise RuntimeError( @@ -1009,19 +1067,17 @@ class ProxyExtrasDBManager: if _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - "prisma migrate deploy attempt %s deadlocked against " - "a concurrent migrate deploy, retrying", + "prisma migrate deploy attempt %s deadlocked against a concurrent migrate deploy, retrying", budget.attempt_number, ) return budget.spend() if "P1002" in stderr and "advisory lock" in stderr: logger.info( - "prisma migrate deploy attempt %s timed out waiting for " - "the advisory lock a concurrent migrate deploy holds, retrying", - budget.attempt_number, + "Waiting for the advisory lock held by another Prisma migration; " + "contention does not spend a migration failure attempt" ) - return budget.spend() + return budget.after_contention(attempt_seconds) raise RuntimeError( "Database migration failed and cannot be auto-recovered. " diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 040d67d25e4..338c571eb4f 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -32,9 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str): def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): """v2: a permission failure during migrate deploy raises RuntimeError.""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -50,9 +48,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): """v2: a non-idempotent migration failure raises (no silent recovery).""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -61,7 +57,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): 'Reason: syntax error at or near "BRKN" LINE 42' ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -176,51 +172,33 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) -def test_v2_resolve_specific_migration_failure_raises_runtime_error( - monkeypatch, tmp_path -): - """If marking a migration as applied fails inside P3009 idempotent - recovery, the subprocess error must be re-raised as RuntimeError so - proxy_cli.py catches it cleanly (instead of leaking CalledProcessError).""" +def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists") monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ProxyExtrasDBManager, + "_v2_roll_back_migration_best_effort", + lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"), ) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") monkeypatch.setattr( - ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"), ) - - # First call: migrate deploy -> P3009 idempotent error. - # Recovery path tries _resolve_specific_migration; that also raises. - def _failing_resolve(*a, **kw): - raise subprocess.CalledProcessError( - returncode=1, - cmd="prisma migrate resolve --applied", - stderr="resolve failed", - output="", - ) - - monkeypatch.setattr( - ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve - ) - - stderr = ( - "Error: P3009\nMigration `20260101000000_some_migration` failed\n" - "relation already exists" - ) - with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises( - RuntimeError, match="Failed to mark migration .* as applied" - ): + stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists" + with patch( + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) + ) as run: + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ) def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -252,9 +230,7 @@ _DEADLOCK_P3018_STDERR = ( def _stub_v2_env(monkeypatch, tmp_path): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") monkeypatch.setattr("time.sleep", lambda _: None) @@ -272,9 +248,7 @@ def _succeed_after(failures: int, stderr: str): return _OkResult() calls["n"] += 1 if calls["n"] <= failures: - raise subprocess.CalledProcessError( - returncode=1, cmd=args[0], stderr=stderr, output="" - ) + raise subprocess.CalledProcessError(returncode=1, cmd=args[0], stderr=stderr, output="") return _OkResult() return _run @@ -288,7 +262,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: rolled_back.append(name), ) monkeypatch.setattr( @@ -306,7 +280,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): """v2: a deadlock on every attempt still fails after the retry budget.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", @@ -335,7 +309,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: rolled_back.append(name), ) monkeypatch.setattr( @@ -350,10 +324,8 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] -def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path): - """v2: empty failed ledger logs mean a concurrent deploy moved it on.""" +def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path): _stub_v2_env(monkeypatch, tmp_path) - stderr = ( "Error: P3009\n" "migrate found failed migrations in the target database\n" @@ -361,22 +333,19 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path "started at 2026-09-01 18:46:13 UTC failed" ) monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") - rolled_back = [] monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", - lambda name: rolled_back.append(name), + "_v2_roll_back_migration_best_effort", + lambda name: pytest.fail("empty logs do not prove rollback is safe"), ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + with patch( + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) + ) as run: + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) - - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) - assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): @@ -392,12 +361,12 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) monkeypatch.setattr( ProxyExtrasDBManager, - "_roll_back_migration", + "_v2_roll_back_migration_best_effort", lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), ) monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -418,7 +387,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index a58c13d6a1c..9ef7cacf1a8 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -6,6 +6,8 @@ Code-style rules for writing tests under `tests/e2e/`. The harness already encod Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family or behavior area. If you add a new folder, you must add a line here describing what kind of tests belong in it, so the layout stays self-describing. `gateway/` is the exception: it holds proxy configuration only and never tests +- `migrations/` - isolated Docker startup, concurrent migration, crash recovery, and legacy database compatibility. The CircleCI migration workflow enables `LITELLM_MIGRATION_TESTS=1`; these tests own their proxy containers and databases, so they do not use the shared proxy preflight or shared database cleanup + - `llm_translation/` - LLM endpoint and provider-translation behavior: passthrough, custom pricing, OCR, and the non-chat inference endpoints (`/v1/responses`, `/v1/messages`, `/embeddings`, `/v1/rerank`, `/v1/audio/speech`, `/v1/images/generations`), each against a deployment the test creates via `/model/new` and deletes on teardown - `access_control/` - the gateway's authorization and error-shape contract: per-key model allow-lists, route-group permissions (`allowed_routes`), and unknown-model validation - `embeddings/` - the `/embeddings` endpoint across providers diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 36569896125..4a5f0aa880f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -64,6 +64,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow") config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", @@ -123,6 +124,11 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: traffic only after the latency-sensitive suites have finished.""" for item in items: attach_result_properties(item) + if os.environ.get("LITELLM_MIGRATION_TESTS") != "1": + deselected = [item for item in items if item.get_closest_marker("migration_startup") is not None] + items[:] = [item for item in items if item.get_closest_marker("migration_startup") is None] + if deselected: + deselected[0].config.hook.pytest_deselected(items=deselected) items.sort(key=lambda item: item.get_closest_marker("load") is not None) @@ -155,7 +161,7 @@ def pytest_runtest_setup(item: pytest.Item) -> None: Unmarked tests (unit coverage of the harness) don't touch the proxy, so they run even when none is up. Never skip for a missing proxy. Replay mode needs the proxy too: only provider-bound traffic replays from the bundle.""" - if item.get_closest_marker("e2e") is None: + if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None: return reason = _proxy_fail_reason() if reason is not None: @@ -168,7 +174,7 @@ def pytest_runtest_call(item: pytest.Item) -> None: guard before truncating the spend-log DB. Tests under `tests/e2e/` without the `e2e` marker (pure unit coverage for the harness itself) never hit the proxy, so they must not arm the destructive DB truncate.""" - if item.get_closest_marker("e2e") is None: + if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None: return item.session.stash[_E2E_TEST_RAN] = True diff --git a/tests/e2e/migrations/__init__.py b/tests/e2e/migrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py new file mode 100644 index 00000000000..b14631ccad8 --- /dev/null +++ b/tests/e2e/migrations/checks.py @@ -0,0 +1,130 @@ +import hashlib +from contextlib import ExitStack +from typing import Final +from uuid import uuid4 + +from psycopg import sql + +from .containers import Containers, Replica, failed, until +from .database import GATE_KEY, Database +from .startup_models import Migration + +COMPLETE_SQL: Final = "CREATE TABLE migration_effect (id int PRIMARY KEY); INSERT INTO migration_effect VALUES (1);" +COMPLETE: Final = Migration("20990101000000_startup_test", COMPLETE_SQL) +NEXT: Final = Migration( + "20990102000000_next_test", + "CREATE TABLE migration_next (id int PRIMARY KEY); INSERT INTO migration_next VALUES (2);", +) +FATAL: Final = Migration(COMPLETE.name, "DO $$ BEGIN RAISE EXCEPTION 'MIGRATION_TEST_FATAL'; END $$;") +GATED: Final = Migration( + COMPLETE.name, f"SELECT pg_advisory_lock({GATE_KEY}); {COMPLETE.script} SELECT pg_advisory_unlock({GATE_KEY});" +) + + +def start_replicas( + stack: ExitStack, containers: Containers, database: Database, migrations: tuple[Migration, ...] = (), count: int = 3 +) -> tuple[Replica, ...]: + return tuple(stack.enter_context(containers.start(database, migrations)) for _ in range(count)) + + +def assert_completed(database: Database, migration: Migration = COMPLETE) -> None: + assert database.query( + "SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (migration.name,), + ) == ((True, True, 1),), "Expected exactly one successful SQL execution" + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + +def confirmed_history(database: Database) -> str: + database.execute(COMPLETE_SQL) + row_id: Final = str(uuid4()) + database.execute( + "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) VALUES (%s, %s, %s, 1)", + (row_id, COMPLETE.name, hashlib.sha256(COMPLETE.script.encode()).hexdigest()), + ) + return row_id + + +def assert_original_proof(database: Database, row_id: str, finished: bool) -> None: + assert database.query( + "SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + ) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof" + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + +def pause_completion(database: Database) -> None: + database.execute( + sql.SQL( + "CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " + "IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN " + "PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; " + "CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()" + ).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY)) + ) + + +def interrupt_owner( + containers: Containers, database: Database, after_commit: bool, *, stop_database_session: bool = True +) -> None: + if after_commit: + pause_completion(database) + with database.lock(): + with containers.start(database, (COMPLETE if after_commit else GATED,)) as owner: + until("migration at the intended crash boundary", lambda: bool(database.blocked())) + assert database.exists("migration_effect") == after_commit + assert database.query( + "SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((True,),) + blocked: Final = database.blocked() + assert len(blocked) == 1 + backend: Final = blocked[0][0] + assert owner.state().Running + owner.kill() + assert owner.state().ExitCode == 137 + if stop_database_session: + database.query("SELECT pg_terminate_backend(%s)", (backend,)) + until( + "terminated migration backend released", + lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)), + ) + assert database.query( + "SELECT finished_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + ) == ((True, int(after_commit)),) + assert database.exists("migration_effect") == after_commit + if not stop_database_session: + until( + "database backend noticed container death", + lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)), + 60, + ) + + +def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None: + failed(replicas, "Migration completion could not be verified") + started: Final = str( + database.query( + "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s", + (COMPLETE.name,), + )[0][0] + ) + for replica in replicas: + assert_guidance(replica.logs(), started) + + +def assert_guidance(log: str, started: str) -> None: + for detail in ( + COMPLETE.name, + started, + "cannot determine whether its SQL committed", + "_prisma_migrations", + "migration.sql", + "Only after verifying every migration change is present", + "prisma migrate resolve --applied ", + "Only after verifying no migration changes remain", + "prisma migrate resolve --rolled-back ", + "leave migration history unchanged", + "Repeated restarts alone", + ): + assert detail in log, f"Missing recovery guidance: {detail}" diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py new file mode 100644 index 00000000000..735adeedbdb --- /dev/null +++ b/tests/e2e/migrations/conftest.py @@ -0,0 +1,62 @@ +import json +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +from _pytest.fixtures import SubRequest + +from .containers import Containers, docker, ready +from .database import Database, Databases + + +@pytest.fixture(scope="session") +def migration_image(tmp_path_factory: pytest.TempPathFactory) -> str: + configured: Final = os.environ.get("LITELLM_MIGRATION_TEST_IMAGE") + assert configured, "LITELLM_MIGRATION_TEST_IMAGE must name the built candidate image" + image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}") + assert image.startswith("sha256:"), "Unable to identify the candidate image" + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) + output.mkdir(parents=True, exist_ok=True) + (output / "image.json").write_text(json.dumps({"requested": configured, "image_id": image})) + return image + + +@pytest.fixture(scope="session") +def databases() -> Databases: + admin: Final = os.environ.get("MIGRATION_TEST_ADMIN_URL", "") + parsed: Final = urlsplit(admin) + assert parsed.hostname in ("127.0.0.1", "localhost"), "Use an isolated loopback PostgreSQL test cluster" + assert parsed.port and parsed.path and not parsed.query, "Supply the test cluster port and admin database" + container_admin: Final = os.environ.get( + "MIGRATION_TEST_CONTAINER_ADMIN_URL", + admin.replace("127.0.0.1", "host.docker.internal").replace("localhost", "host.docker.internal"), + ) + return Databases(admin, container_admin) + + +@pytest.fixture(scope="session") +def migrated_template( + databases: Databases, migration_image: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Database]: + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "seed" + with databases.create() as database: + with Containers(migration_image, output).start(database) as replica: + ready((replica,), database) + yield database + + +@pytest.fixture +def database(databases: Databases, migrated_template: Database) -> Iterator[Database]: + with databases.create(migrated_template) as database: + yield database + + +@pytest.fixture +def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Containers: + configured: Final = os.environ.get("MIGRATION_TEST_OUTPUT") + output: Final = Path(configured) / request.node.name if configured else tmp_path + output.mkdir(parents=True, exist_ok=True) + return Containers(migration_image, output) diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py new file mode 100644 index 00000000000..0f5793b81dd --- /dev/null +++ b/tests/e2e/migrations/containers.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import hashlib +import subprocess +import time +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from uuid import uuid4 + +from e2e_http import NoBody, Success, unwrap +from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse +from transport import HttpTransport + +from .database import Database, prisma_url +from .startup_models import ContainerState, Migration, Observation, Readiness + +MASTER_KEY: Final = "sk-migration-ci-fixture" + + +def docker(*args: str) -> str: + result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90) + assert result.returncode == 0, f"Docker operation failed: {result.stderr}" + return result.stdout.strip() + + +def until(description: str, condition: Callable[[], bool], seconds: float = 150) -> None: + deadline: Final = time.monotonic() + seconds + while time.monotonic() < deadline: + if condition(): + return + time.sleep(0.25) + raise AssertionError(f"Timed out waiting for {description}") + + +@dataclass(frozen=True, slots=True) +class Replica: + name: str + transport: HttpTransport + output: Path + + def state(self) -> ContainerState: + return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name)) + + def observe(self) -> Observation: + state: Final = self.state() + result: Final = self.transport.get( + "/health/readiness", headers=self.transport.master, params=NoBody(), response_type=Readiness, timeout=1 + ) + ready: Final = isinstance(result, Success) and result.data.status == "healthy" and result.data.db == "connected" + return Observation(None if state.Running else state.ExitCode, ready) + + def logs(self) -> str: + result: Final = subprocess.run(("docker", "logs", self.name), capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stderr + return result.stdout + result.stderr + + def kill(self) -> None: + if self.state().Running: + docker("kill", self.name) + + def usable(self, database: Database) -> None: + alias: Final = f"migration-{uuid4().hex}" + key: Final = unwrap( + self.transport.post( + "/key/generate", + headers=self.transport.master, + json=KeyGenerateBody(key_alias=alias), + response_type=KeyGenerateResponse, + ) + ).key + info: Final = unwrap( + self.transport.get( + "/key/info", + headers=self.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ) + assert info.info.key_alias == alias + assert database.query( + 'SELECT key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s', + (hashlib.sha256(key.encode()).hexdigest(),), + ) == ((alias,),) + + +def ready(replicas: tuple[Replica, ...], database: Database) -> None: + def all_ready() -> bool: + observations: Final = tuple(replica.observe() for replica in replicas) + assert all(item.exit_code is None for item in observations), "Replica exited before readiness" + return all(item.ready for item in observations) + + until("every replica ready", all_ready) + for replica in replicas: + replica.usable(database) + + +def failed(replicas: tuple[Replica, ...], marker: str) -> None: + def all_stopped() -> bool: + observations: Final = tuple(replica.observe() for replica in replicas) + assert not any(item.ready for item in observations), "Failed migration exposed a ready proxy" + return all(item.exit_code is not None for item in observations) + + until("every replica to reject startup", all_stopped) + for replica in replicas: + assert replica.state().ExitCode != 0, "Failed startup returned success" + assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}" + + +def waiting(replicas: tuple[Replica, ...], seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while time.monotonic() < deadline: + assert all(item.exit_code is None and not item.ready for item in (replica.observe() for replica in replicas)), ( + "Contending replica exited or served early" + ) + time.sleep(0.25) + + +@dataclass(frozen=True, slots=True) +class Containers: + image: str + output: Path + + @contextmanager + def start( + self, + database: Database, + migrations: tuple[Migration, ...] = (), + *, + v2: bool = True, + disabled: bool = False, + environment: Mapping[str, str] | None = None, + ) -> Generator[Replica]: + name: Final = f"litellm-migration-{uuid4().hex[:16]}" + directory: Final = self.output / name + directory.mkdir(parents=True) + for migration in migrations: + write_migration(directory, migration) + (directory / "config.yaml").write_text( + "model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n" + ) + env: Final = { + "DATABASE_URL": prisma_url(database.container_url, database.schema), + "LITELLM_MASTER_KEY": MASTER_KEY, + "LITELLM_SALT_KEY": MASTER_KEY, + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_TELEMETRY": "False", + "LITELLM_LOG": "INFO", + "DATABASE_CONNECTION_POOL_LIMIT": "2", + "DEFAULT_NUM_WORKERS_LITELLM_PROXY": "1", + "USE_V2_MIGRATION_RESOLVER": str(v2).lower(), + "DISABLE_SCHEMA_UPDATE": str(disabled).lower(), + "LITELLM_MIGRATION_DIR": "/migration-test/prisma", + "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT": "180", + **(environment or {}), + } + try: + docker( + "run", + "-d", + "--name", + name, + "--label", + "litellm-migration-test=true", + "--add-host", + "host.docker.internal:host-gateway", + "-p", + "127.0.0.1::4000", + "-v", + f"{directory}:/migration-test", + *(arg for key, value in env.items() for arg in ("-e", f"{key}={value}")), + self.image, + "--config", + "/migration-test/config.yaml", + "--host", + "0.0.0.0", + "--port", + "4000", + ) + port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1]) + replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory) + yield replica + finally: + try: + state: Final = subprocess.run( + ("docker", "inspect", "--format", "{{json .State}}", name), + capture_output=True, + text=True, + timeout=30, + ) + (directory / "state.json").write_text(state.stdout or state.stderr) + logs: Final = subprocess.run(("docker", "logs", name), capture_output=True, text=True, timeout=30) + (directory / "proxy.log").write_text(logs.stdout + logs.stderr) + finally: + subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True) + + +def write_migration(directory: Path, migration: Migration) -> None: + path: Final = directory / "prisma" / "migrations" / migration.name + path.mkdir(parents=True) + (path / "migration.sql").write_text(migration.script) diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py new file mode 100644 index 00000000000..a370c21ba0b --- /dev/null +++ b/tests/e2e/migrations/database.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Final, LiteralString +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit +from uuid import uuid4 + +import psycopg +from psycopg import sql +from pydantic import TypeAdapter + +Scalar = str | int | bool | None +ROWS: Final = TypeAdapter(tuple[tuple[Scalar, ...], ...]) +GATE_KEY: Final = 39178002 +PRISMA_LOCK: Final = 72707369 +COORDINATOR_LOCK: Final = int.from_bytes(b"llm_mig2", "big") + + +def connect_url(url: str, name: str) -> str: + return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query="")) + + +def prisma_url(url: str, schema: str) -> str: + parsed: Final = urlsplit(url) + query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema") + return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema))))) + + +@dataclass(frozen=True, slots=True) +class Database: + name: str + url: str + container_url: str + schema: str = "public" + + @contextmanager + def connection(self) -> Generator[psycopg.Connection[tuple[object, ...]]]: + with psycopg.connect(self.url, autocommit=True, connect_timeout=5) as connection: + connection.execute(sql.SQL("SET search_path TO {}").format(sql.Identifier(self.schema))) + connection.execute("SET statement_timeout = '15s'") + yield connection + + def execute(self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()) -> None: + with self.connection() as connection: + connection.execute(statement, params or None) + + def query( + self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = () + ) -> tuple[tuple[Scalar, ...], ...]: + with self.connection() as connection: + return ROWS.validate_python(connection.execute(statement, params or None).fetchall()) + + def exists(self, name: str) -> bool: + return self.query("SELECT to_regclass(%s) IS NOT NULL", (name,)) == ((True,),) + + def history(self) -> tuple[tuple[Scalar, ...], ...]: + if not self.exists("_prisma_migrations"): + return () + return self.query( + "SELECT id, migration_name, checksum, started_at::text, finished_at::text, rolled_back_at::text, " + "applied_steps_count, logs FROM _prisma_migrations ORDER BY id" + ) + + def blocked(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]: + return self.query( + "SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND NOT granted " + "AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND classid = %s AND objid = %s ORDER BY pid", + (key >> 32, key & 0xFFFFFFFF), + ) + + @contextmanager + def lock(self, key: int = GATE_KEY) -> Generator[None]: + with self.connection() as connection: + connection.execute("SELECT pg_advisory_lock(%s)", (key,)) + try: + yield + finally: + connection.execute("SELECT pg_advisory_unlock(%s)", (key,)) + + +@dataclass(frozen=True, slots=True) +class Databases: + admin_url: str + container_admin_url: str + + @contextmanager + def create(self, template: Database | None = None, schema: str = "public") -> Generator[Database]: + name: Final = f"litellm_migration_test_{uuid4().hex[:20]}" + database: Final = Database( + name, connect_url(self.admin_url, name), connect_url(self.container_admin_url, name), schema + ) + with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection: + connection.execute( + sql.SQL("CREATE DATABASE {} TEMPLATE {}").format( + sql.Identifier(name), sql.Identifier(template.name if template else "template0") + ) + ) + try: + yield database + finally: + with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection: + connection.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) + + +@contextmanager +def restricted_user(database: Database) -> Generator[Database]: + role: Final = f"migration_reader_{uuid4().hex[:16]}" + password: Final = "migration-test-password" + with database.connection() as connection: + connection.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(sql.Identifier(role), sql.Literal(password)) + ) + try: + database.execute( + sql.SQL("GRANT USAGE ON SCHEMA {} TO {}").format(sql.Identifier(database.schema), sql.Identifier(role)) + ) + database.execute( + sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA {} TO {}").format( + sql.Identifier(database.schema), sql.Identifier(role) + ) + ) + local: Final = urlsplit(database.url) + remote: Final = urlsplit(database.container_url) + yield Database( + database.name, + urlunsplit(local._replace(netloc=f"{role}:{password}@{local.hostname}:{local.port}")), + urlunsplit(remote._replace(netloc=f"{role}:{password}@{remote.hostname}:{remote.port}")), + database.schema, + ) + finally: + database.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) + database.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) diff --git a/tests/e2e/migrations/startup_models.py b/tests/e2e/migrations/startup_models.py new file mode 100644 index 00000000000..03a9b4eda78 --- /dev/null +++ b/tests/e2e/migrations/startup_models.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass + +from pydantic import BaseModel + + +class Readiness(BaseModel): + status: str = "" + db: str = "" + + +class ContainerState(BaseModel): + Running: bool + ExitCode: int + + +@dataclass(frozen=True, slots=True) +class Observation: + exit_code: int | None + ready: bool + + +@dataclass(frozen=True, slots=True) +class Migration: + name: str + script: str diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py new file mode 100644 index 00000000000..7ba73eb82e0 --- /dev/null +++ b/tests/e2e/migrations/test_legacy.py @@ -0,0 +1,84 @@ +from contextlib import ExitStack +from dataclasses import replace +from typing import Final, Literal + +import pytest + +from .checks import COMPLETE, assert_completed, confirmed_history, assert_original_proof, start_replicas +from .containers import Containers, failed, ready +from .database import Database, Databases + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +def adopt_legacy(containers: Containers, database: Database) -> None: + count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0] + existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') + database.execute( + "INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" + ) + database.execute("DROP TABLE _prisma_migrations") + with ExitStack() as stack: + replicas: Final = start_replicas(stack, containers, database) + ready(replicas, database) + logs: Final = "\n".join(replica.logs() for replica in replicas) + for detail in ( + "Legacy migration history was missing", + "historical data backfills were not replayed or verified", + "Continuing startup", + ): + assert detail in logs + assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0" + ) == ((0,),) + assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) + assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == ( + (None,), + ) + + +class TestLegacyMigrations: + def test_matching_schema_warns_and_starts(self, containers: Containers, database: Database) -> None: + adopt_legacy(containers, database) + + @pytest.mark.parametrize("fault", ("schema_drift", "custom_migrations", "empty_ledger")) + def test_unrecognized_legacy_state_is_not_baselined( + self, containers: Containers, database: Database, fault: str + ) -> None: + if fault == "empty_ledger": + database.execute("TRUNCATE _prisma_migrations") + else: + database.execute("DROP TABLE _prisma_migrations") + if fault == "schema_drift": + database.execute('ALTER TABLE "LiteLLM_VerificationToken" DROP COLUMN key_alias CASCADE') + with containers.start(database, (COMPLETE,) if fault == "custom_migrations" else ()) as replica: + failed((replica,), "Cannot automatically baseline" if fault != "empty_ledger" else "migration") + assert not database.exists("migration_effect") + if database.exists("_prisma_migrations"): + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NOT NULL AND applied_steps_count <> 1" + ) == ((0,),) + + @pytest.mark.parametrize("scenario", ("upgrade", "recovery", "legacy")) + def test_non_default_schema( + self, containers: Containers, databases: Databases, scenario: Literal["upgrade", "recovery", "legacy"] + ) -> None: + with databases.create(schema="migration tenant") as database: + with containers.start(database) as seed: + ready((seed,), database) + match scenario: + case "upgrade": + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_completed(database) + case "recovery": + original: Final = confirmed_history(database) + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_original_proof(database, original, True) + case "legacy": + adopt_legacy(containers, database) + public: Final = replace(database, schema="public") + assert not public.exists("_prisma_migrations") + assert not public.exists('"LiteLLM_VerificationToken"') diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py new file mode 100644 index 00000000000..e0a3693b33e --- /dev/null +++ b/tests/e2e/migrations/test_pooling.py @@ -0,0 +1,135 @@ +import subprocess +from collections.abc import Generator +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit +from uuid import uuid4 + +import psycopg +import pytest +from psycopg import sql + +from .checks import COMPLETE, assert_completed +from .containers import Containers, docker, ready, until +from .database import Database, Databases, prisma_url, restricted_user + +POOL_IMAGE: Final = ( + "ghcr.io/cloudnative-pg/pgbouncer@sha256:e6ddfe22d845e603825e235dd8334b21ecd125abea2a2172478f556b8dee2bb8" +) +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +@contextmanager +def application_user(database: Database) -> Generator[Database]: + with restricted_user(database) as application: + role: Final = sql.Identifier(str(urlsplit(application.url).username)) + schema: Final = sql.Identifier(database.schema) + database.execute(sql.SQL("REVOKE CREATE ON SCHEMA {} FROM PUBLIC").format(schema)) + for statement in ( + "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {} TO {}", + "GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {} TO {}", + "ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO {}", + "ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT USAGE, SELECT ON SEQUENCES TO {}", + ): + database.execute(sql.SQL(statement).format(schema, role)) + assert application.query("SELECT has_schema_privilege(current_user, %s, 'CREATE')", (database.schema,)) == ( + (False,), + ) + yield application + + +@contextmanager +def pool(database: Database, output: Path) -> Generator[str]: + name: Final = f"litellm-migration-pool-{uuid4().hex[:12]}" + url: Final = urlsplit(database.container_url) + directory: Final = output / name + directory.mkdir(parents=True) + (directory / "users.txt").write_text(f'"{url.username}" "{url.password}"\n') + (directory / "pgbouncer.ini").write_text( + f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n" + "[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n" + "pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n" + "max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n" + ) + try: + docker( + "run", + "-d", + "--name", + name, + "--label", + "litellm-migration-test=true", + "--add-host", + "host.docker.internal:host-gateway", + "-p", + "0.0.0.0::6432", + "-v", + f"{directory}:/pool:ro", + "--entrypoint", + "/usr/bin/pgbouncer", + POOL_IMAGE, + "/pool/pgbouncer.ini", + ) + port: Final = int(docker("port", name, "6432/tcp").splitlines()[0].rsplit(":", 1)[1]) + local_url: Final = urlunsplit(url._replace(netloc=f"{url.username}:{url.password}@127.0.0.1:{port}")) + + def connected() -> bool: + try: + with psycopg.connect(local_url, autocommit=True, connect_timeout=2) as connection: + return connection.execute("SELECT 1").fetchone() == (1,) + except psycopg.Error: + return False + + until("PgBouncer ready", connected, 30) + yield local_url.replace("127.0.0.1", "host.docker.internal") + "?pgbouncer=true" + finally: + try: + logs: Final = subprocess.run(("docker", "logs", name), text=True, capture_output=True, timeout=30) + (directory / "pool.log").write_text(logs.stdout + logs.stderr) + finally: + subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True) + + +class TestMigrationPooling: + @pytest.mark.parametrize("scenario,replica_count", (("fresh", 3), ("upgrade", 3), ("legacy", 3), ("upgrade", 6))) + def test_direct_migrations_with_one_application_backend( + self, + containers: Containers, + databases: Databases, + migrated_template: Database, + scenario: str, + replica_count: int, + ) -> None: + with databases.create(None if scenario == "fresh" else migrated_template) as database: + if scenario == "legacy": + database.execute("DROP TABLE _prisma_migrations") + with ( + application_user(database) as application, + pool(application, containers.output) as pooled_url, + ExitStack() as stack, + ): + replicas: Final = tuple( + stack.enter_context( + containers.start( + database, + (COMPLETE,) if scenario == "upgrade" else (), + environment={ + "DATABASE_URL": prisma_url(pooled_url, database.schema), + "DIRECT_URL": database.container_url, + }, + ) + ) + for _ in range(replica_count) + ) + ready(replicas, database) + if scenario == "upgrade": + assert_completed(database) + if scenario == "legacy": + assert any( + "historical data backfills were not replayed or verified" in replica.logs() + for replica in replicas + ) + assert database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count <> 0") == ( + (0,), + ) diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py new file mode 100644 index 00000000000..58bf5c348d6 --- /dev/null +++ b/tests/e2e/migrations/test_recovery.py @@ -0,0 +1,183 @@ +from contextlib import ExitStack +from typing import Final, Literal +from uuid import uuid4 + +import pytest + +from .checks import ( + COMPLETE, + FATAL, + GATED, + NEXT, + assert_completed, + confirmed_history, + interrupt_owner, + assert_original_proof, + pause_completion, + start_replicas, + unconfirmed, +) +from .containers import Containers, failed, ready, until, waiting +from .database import COORDINATOR_LOCK, GATE_KEY, Database +from .startup_models import Migration + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestMigrationRecovery: + @pytest.mark.parametrize("after_commit", (False, True)) + def test_container_owner_crash(self, containers: Containers, database: Database, after_commit: bool) -> None: + interrupt_owner(containers, database, after_commit, stop_database_session=False) + history: Final = database.history() + assert database.query( + "SELECT applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((int(after_commit),),) + with ExitStack() as stack: + successors: Final = start_replicas(stack, containers, database, (COMPLETE if after_commit else GATED,)) + if after_commit: + ready(successors, database) + assert_completed(database) + else: + unconfirmed(successors, database) + assert database.history() == history + + @pytest.mark.parametrize("after_commit", (False, True)) + def test_owner_and_database_session_crash( + self, containers: Containers, database: Database, after_commit: bool + ) -> None: + interrupt_owner(containers, database, after_commit) + history: Final = database.history() + with ExitStack() as stack: + successors: Final = start_replicas(stack, containers, database, (COMPLETE,)) + if after_commit: + ready(successors, database) + assert_completed(database) + return + unconfirmed(successors, database) + assert database.history() == history + with containers.start(database, (COMPLETE,)) as restarted: + unconfirmed((restarted,), database) + assert database.history() == history + + @pytest.mark.parametrize("later_failure", (False, True)) + def test_remaining_migrations_after_recovery( + self, containers: Containers, database: Database, later_failure: bool + ) -> None: + original: Final = confirmed_history(database) + next_migration: Final = Migration( + NEXT.name, + f"SELECT pg_advisory_lock({GATE_KEY}); " + + (FATAL.script if later_failure else NEXT.script) + + f" SELECT pg_advisory_unlock({GATE_KEY});", + ) + with ExitStack() as stack: + with database.lock(): + owner: Final = stack.enter_context(containers.start(database, (COMPLETE, next_migration))) + + def pending() -> bool: + observation: Final = owner.observe() + assert observation.exit_code is None and not observation.ready, ( + "Recovered owner served before pending SQL completed" + ) + return bool(database.blocked()) + + until("recovering owner reached the next migration", pending) + assert_original_proof(database, original, True) + assert not database.exists("migration_next") + followers: Final = start_replicas(stack, containers, database, (COMPLETE, next_migration), count=2) + replicas: Final = (owner, *followers) + waiting(replicas, 1) + if later_failure: + failed(replicas, NEXT.name) + assert database.query( + "SELECT finished_at IS NULL, logs LIKE %s FROM _prisma_migrations WHERE migration_name = %s", + ("%MIGRATION_TEST_FATAL%", NEXT.name), + ) == ((True, True),) + else: + ready(replicas, database) + assert database.query("SELECT id FROM migration_next") == ((2,),) + assert_original_proof(database, original, True) + + def test_second_crash_during_recovery_is_atomic(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + pause_completion(database) + with database.lock(): + with containers.start(database, (COMPLETE,)) as recovering: + until("history update blocked before commit", lambda: bool(database.blocked())) + assert_original_proof(database, original, False) + blocked: Final = database.blocked() + assert len(blocked) == 1 + assert database.query("SELECT pg_terminate_backend(%s)", (blocked[0][0],)) == ((True,),) + failed((recovering,), "Lost or could not establish v2 migration coordination") + assert_original_proof(database, original, False) + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_original_proof(database, original, True) + + def test_competing_recovery_rechecks_stale_failures(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + with ExitStack() as stack: + with database.lock(COORDINATOR_LOCK): + replicas: Final = start_replicas(stack, containers, database, (COMPLETE,)) + until( + "all replicas observed the unfinished migration", + lambda: all( + "Waiting for the v2 migration coordinator lock" in replica.logs() for replica in replicas + ), + ) + assert_original_proof(database, original, False) + ready(replicas, database) + assert_original_proof(database, original, True) + + @pytest.mark.parametrize( + "fault", ("no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script") + ) + def test_unproven_history_is_never_repaired( + self, + containers: Containers, + database: Database, + fault: Literal["no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script"], + ) -> None: + confirmed_history(database) + match fault: + case "no_steps": + database.execute( + "UPDATE _prisma_migrations SET applied_steps_count = 0 WHERE migration_name = %s", (COMPLETE.name,) + ) + case "extra_steps": + database.execute( + "UPDATE _prisma_migrations SET applied_steps_count = 2 WHERE migration_name = %s", (COMPLETE.name,) + ) + case "failure_logs": + database.execute( + "UPDATE _prisma_migrations SET logs = 'permission denied' WHERE migration_name = %s", + (COMPLETE.name,), + ) + case "checksum": + database.execute( + "UPDATE _prisma_migrations SET checksum = %s WHERE migration_name = %s", ("0" * 64, COMPLETE.name) + ) + case "duplicate_history": + database.execute( + "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + (str(uuid4()), COMPLETE.name), + ) + case "missing_script": + pass + history: Final = database.history() + with containers.start(database, () if fault == "missing_script" else (COMPLETE,)) as replica: + unconfirmed((replica,), database) + assert database.history() == history + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + def test_coordinator_timeout_preserves_proof(self, containers: Containers, database: Database) -> None: + original: Final = confirmed_history(database) + with database.lock(COORDINATOR_LOCK): + with containers.start( + database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "3"} + ) as replica: + failed((replica,), "Timed out waiting for another v2 migration resolver") + assert_original_proof(database, original, False) + with containers.start(database, (COMPLETE,)) as replica: + ready((replica,), database) + assert_original_proof(database, original, True) diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py new file mode 100644 index 00000000000..a648218cb26 --- /dev/null +++ b/tests/e2e/migrations/test_startup.py @@ -0,0 +1,100 @@ +from contextlib import ExitStack +from typing import Final + +import pytest + +from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas +from .containers import Containers, failed, ready, until, waiting +from .database import PRISMA_LOCK, Database, Databases, restricted_user +from .startup_models import Migration + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestMigrationStartup: + @pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False))) + def test_fresh_database(self, containers: Containers, databases: Databases, replicas: int, v2: bool) -> None: + with databases.create() as database, ExitStack() as stack: + ready(tuple(stack.enter_context(containers.start(database, v2=v2)) for _ in range(replicas)), database) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL AND rolled_back_at IS NULL" + ) == ((0,),) + assert database.query("SELECT count(*) > 0 FROM _prisma_migrations") == ((True,),) + + def test_concurrent_upgrade(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + ready(start_replicas(stack, containers, database, (COMPLETE,)), database) + assert_completed(database) + + def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + with database.lock(): + owner: Final = stack.enter_context(containers.start(database, (GATED,))) + until("owner blocked in migration SQL", lambda: bool(database.blocked())) + followers: Final = start_replicas(stack, containers, database, (GATED,), count=2) + until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2) + waiting((owner, *followers), 120) + ready((owner, *followers), database) + assert_completed(database, GATED) + + def test_lock_deadline_then_restart(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with database.lock(PRISMA_LOCK): + with containers.start( + database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"} + ) as replica: + until("Prisma lock contention", lambda: bool(database.blocked(PRISMA_LOCK))) + failed((replica,), "Timed out waiting for") + assert database.history() == history + assert not database.exists("migration_effect") + with containers.start(database, (COMPLETE,)) as restarted: + ready((restarted,), database) + assert_completed(database) + + def test_fatal_sql(self, containers: Containers, database: Database) -> None: + with ExitStack() as stack: + replicas: Final = start_replicas(stack, containers, database, (FATAL,)) + failed(replicas, COMPLETE.name) + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL", + (COMPLETE.name, "%MIGRATION_TEST_FATAL%"), + ) == ((1,),) + + def test_duplicate_object_does_not_hide_incomplete_sql(self, containers: Containers, database: Database) -> None: + database.execute( + "CREATE TABLE migration_existing (id int PRIMARY KEY); INSERT INTO migration_existing VALUES (42)" + ) + migration: Final = Migration( + COMPLETE.name, "CREATE TABLE migration_existing (id int PRIMARY KEY); " + COMPLETE.script + ) + with ExitStack() as stack: + failed(start_replicas(stack, containers, database, (migration,)), COMPLETE.name) + assert not database.exists("migration_effect") + assert database.query("SELECT id FROM migration_existing") == ((42,),) + assert database.query( + "SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,) + ) == ((True,),) + + @pytest.mark.parametrize("v2", (True, False)) + def test_restart_preserves_history_and_data(self, containers: Containers, database: Database, v2: bool) -> None: + history: Final = database.history() + before: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') + for _ in range(2): + with containers.start(database, v2=v2) as replica: + ready((replica,), database) + assert database.history() == history + assert set(before).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) + + def test_disabled_migrations(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with containers.start(database, (FATAL,), disabled=True) as replica: + ready((replica,), database) + assert database.history() == history + + def test_insufficient_privileges(self, containers: Containers, database: Database) -> None: + history: Final = database.history() + with restricted_user(database) as limited: + with containers.start(limited, (COMPLETE,)) as replica: + failed((replica,), "permission denied") + assert database.history() == history + assert not database.exists("migration_effect") diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index e5826b18668..57133ea95c4 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -728,12 +728,36 @@ ERROR: relation "SomeTable" already exists """ +@pytest.mark.parametrize( + "pooled,direct,expected", + ( + ("postgresql://pool/db?pgbouncer=true", None, "postgresql://pool/db?pgbouncer=true"), + ("postgresql://pool/db?pgbouncer=true", "postgresql://writer/db", "postgresql://writer/db?schema=public"), + ( + "postgresql://pool/db?schema=tenant%20one&pgbouncer=true", + "postgresql://writer/db?sslmode=require&schema=wrong", + "postgresql://writer/db?sslmode=require&schema=tenant+one", + ), + ), +) +def test_v2_migrations_use_the_direct_connection_with_the_runtime_schema(pooled, direct, expected): + from litellm_proxy_extras.migration_lock import migration_environment + + environment = {"DATABASE_URL": pooled, "PRISMA_OFFLINE_MODE": "true"} + configured = {**environment, **({"DIRECT_URL": direct} if direct else {})} + migrated = migration_environment(configured) + + assert migrated["DATABASE_URL"] == expected + assert migrated["PRISMA_OFFLINE_MODE"] == "true" + assert configured["DATABASE_URL"] == pooled + + class _MigrateDeployHarness: """Drives _setup_database_v2 with a scripted sequence of `prisma migrate deploy` outcomes, with every recovery command faked out so nothing touches a database or the packaged migrations directory.""" - def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False): + def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False, confirmed_migrations=()): import subprocess as subprocess_module import litellm_proxy_extras.utils as utils_module @@ -744,16 +768,10 @@ class _MigrateDeployHarness: self._outcomes = list(outcomes) self._repeat_last = repeat_last self._subprocess_module = subprocess_module + self.confirmed_migrations = set(confirmed_migrations) monkeypatch.delenv("DATABASE_URL", raising=False) - monkeypatch.setattr( - ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path)) - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_create_baseline_migration", - staticmethod(self._fake_baseline), - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))) monkeypatch.setattr( ProxyExtrasDBManager, "_roll_back_migration", @@ -765,13 +783,15 @@ class _MigrateDeployHarness: staticmethod(self.resolved.append), ) monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) + monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) self.baseline_succeeds = True def _fake_baseline(self, *args, **kwargs): self.baselines += 1 - return self.baseline_succeeds + if not self.baseline_succeeds: + raise RuntimeError("The existing schema was not verified") def _next_outcome(self): if self._outcomes: @@ -791,79 +811,126 @@ class _MigrateDeployHarness: raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome) def run(self): - return ProxyExtrasDBManager._setup_database_v2(use_migrate=True) + while not ProxyExtrasDBManager._run_database_v2( + use_migrate=True, + recover_completed=self._fake_recovery, + baseline_existing=self._fake_baseline, + ): + continue + return True + + def _fake_recovery(self, name): + if name not in self.confirmed_migrations: + return False + self.confirmed_migrations.remove(name) + self.resolved.append(name) + return True class TestMigrateDeployAttemptAccounting: - """A `prisma db push` database has a full schema and no ledger, so the v2 - resolver baselines it and then works through every migration whose objects - already exist. Those recoveries make progress, so they must not spend the - retry budget, which is there to stop a run that is getting nowhere.""" - - def test_a_push_created_database_finishes_bootstrapping( - self, monkeypatch, tmp_path - ): - already_there = [ - "20250329084805_new_cron_job_table", - "20250806095134_rename_alias_to_server_name_mcp_table", - "20260224203854_add_agent_object_permissions_table", - "20260301120000_fourth_table", - "20260302120000_fifth_table", - "20260303120000_sixth_table", - ] + def test_a_push_created_database_finishes_bootstrapping(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, - [_P3005_STDERR] - + [_p3018_stderr(name) for name in already_there] - + ["ok"], + [_P3005_STDERR, "ok"], ) assert harness.run() is True assert harness.baselines == 1 - assert harness.resolved == already_there - assert len(harness.deploy_calls) == len(already_there) + 2 + assert harness.resolved == [] + assert len(harness.deploy_calls) == 2 - def test_repeated_recovery_of_one_migration_still_gives_up( - self, monkeypatch, tmp_path - ): + def test_repeated_recovery_of_one_migration_still_gives_up(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, [_p3018_stderr("20250329084805_new_cron_job_table")], repeat_last=True, + confirmed_migrations=("20250329084805_new_cron_job_table",), ) with pytest.raises(RuntimeError): harness.run() - assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1 + assert len(harness.deploy_calls) == 2 + assert harness.resolved == ["20250329084805_new_cron_job_table"] def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path): - harness = _MigrateDeployHarness( - monkeypatch, tmp_path, ["timeout"], repeat_last=True - ) + harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["timeout"], repeat_last=True) with pytest.raises(RuntimeError): harness.run() assert len(harness.deploy_calls) == _ATTEMPT_BUDGET - def test_a_baseline_that_never_lands_stops_after_the_budget( - self, monkeypatch, tmp_path - ): - harness = _MigrateDeployHarness( - monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True - ) + def test_an_unverified_baseline_stops_without_replaying_migrations(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True) harness.baseline_succeeds = False with pytest.raises(RuntimeError): harness.run() - assert len(harness.deploy_calls) == _ATTEMPT_BUDGET + assert len(harness.deploy_calls) == 1 + + def test_lock_contention_does_not_spend_the_failure_budget(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + ["Error: P1002\nTimed out waiting for the advisory lock"] * 6 + ["ok"], + ) + assert harness.run() is True + assert len(harness.deploy_calls) == 7 + + def test_duplicate_object_error_without_completion_proof_is_fatal(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_p3018_stderr("20260101000000_x")]) + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + harness.run() + assert harness.resolved == [] + assert len(harness.deploy_calls) == 1 + + @pytest.mark.parametrize("name", ("20260101000000_x", "20260101000000_migration with spaces")) + def test_an_interrupted_migration_with_confirmed_sql_can_finish(self, monkeypatch, tmp_path, name): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [f"Error: P3009\nThe `{name}` migration failed", "ok"], + confirmed_migrations=(name,), + ) + assert harness.run() is True + assert harness.resolved == [name] + + def test_an_interrupted_migration_without_confirmation_stops(self, monkeypatch, tmp_path): + name = "20260101000000_x" + started = "2026-09-12 20:15:06.694553 UTC" + report = f"Error: P3009\nThe `{name}` migration started at {started} failed" + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [report], + ) + with pytest.raises(RuntimeError, match="Migration completion could not be verified") as failure: + harness.run() + message = str(failure.value) + assert name in message + assert started in message + assert "start record but no successful completion record" in message + assert "cannot determine whether its SQL committed" in message + assert "avoid repeating or skipping database changes" in message + assert "_prisma_migrations" in message + assert "migration.sql" in message + assert "same database" in message + assert "Only after verifying every migration change is present" in message + assert "prisma migrate resolve --applied " in message + assert "Only after verifying no migration changes remain" in message + assert "prisma migrate resolve --rolled-back " in message + assert "leave migration history unchanged" in message + assert "Repeated restarts alone" in message + assert report in message + assert len(harness.deploy_calls) == 1 + assert harness.resolved == [] def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness( monkeypatch, tmp_path, - ["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"], + ['Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near "SLECT"\n'], repeat_last=True, ) @@ -873,6 +940,36 @@ class TestMigrateDeployAttemptAccounting: assert harness.resolved == [] +@pytest.mark.parametrize( + "steps,logs,script,expected", + ( + (1, "", b"CREATE TABLE item (id int);", True), + (0, "", b"CREATE TABLE item (id int);", False), + (0, "already exists", b"CREATE TABLE item (id int);", False), + (1, "permission denied", b"CREATE TABLE item (id int);", False), + (1, "", b"CREATE TABLE item (id text);", False), + (2, "", b"CREATE TABLE item (id int);", False), + ), +) +def test_migration_completion_requires_a_matching_successful_script(steps, logs, script, expected): + import hashlib + + from litellm_proxy_extras.migration_recovery import MigrationProgress + + progress = MigrationProgress(hashlib.sha256(b"CREATE TABLE item (id int);").hexdigest(), steps, logs) + assert progress.confirms_completion(script) is expected + + +def test_prisma_lock_waiting_has_its_own_deadline(): + from litellm_proxy_extras.utils import _MigrateAttemptBudget + + budget = _MigrateAttemptBudget(attempts_left=4, contention_seconds_left=2) + waiting = budget.after_contention(1) + assert waiting.attempts_left == 4 + with pytest.raises(RuntimeError, match="advisory lock"): + waiting.after_contention(2) + + class TestJWTKeyMappingCascade: """Regression tests for issue #33702. diff --git a/tests/proxy_migration_tests/test_migration_ci.py b/tests/proxy_migration_tests/test_migration_ci.py new file mode 100644 index 00000000000..30fe7383c07 --- /dev/null +++ b/tests/proxy_migration_tests/test_migration_ci.py @@ -0,0 +1,36 @@ +import importlib.util +from pathlib import Path +from typing import Final + +import pytest + +SCRIPT: Final = Path(__file__).resolve().parents[2] / ".circleci/scripts/run_migration_tests.py" +SPEC: Final = importlib.util.spec_from_file_location("migration_ci", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +MODULE: Final = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +@pytest.mark.parametrize( + "xml,expected,exit_code,passed", + ( + ('', 2, 0, True), + ('', 2, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 1, False), + ("", 1, 5, False), + ("", 1, 0, False), + ("', 2, 0, False), + ), +) +def test_only_a_complete_passing_suite_can_certify_an_image( + tmp_path: Path, xml: str | None, expected: int, exit_code: int, passed: bool +) -> None: + path: Final = tmp_path / "results.xml" + if xml is not None: + path.write_text(xml) + assert MODULE.successful_junit(path, expected, exit_code) is passed From c44757fc010a0f81fe9c86fcabc43e95ecb8dd57 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:33:09 -0700 Subject: [PATCH 002/127] ci: fetch migration test revisions over HTTPS --- .circleci/config.yml | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index f6f31651306..c6e18c40213 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -22,11 +22,12 @@ commands: environment: MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >> command: | - if [ -n "$MIGRATION_SOURCE_SHA" ]; then - [[ "$MIGRATION_SOURCE_SHA" =~ ^[0-9a-f]{40}$ ]] || exit 1 - git fetch origin "$MIGRATION_SOURCE_SHA" - git checkout --detach "$MIGRATION_SOURCE_SHA" - fi + revision="${MIGRATION_SOURCE_SHA:-$CIRCLE_SHA1}" + [[ "$revision" =~ ^[0-9a-f]{40}$ ]] || exit 1 + git init + git remote add origin https://github.com/BerriAI/litellm.git + git fetch --depth 1 origin "$revision" + git checkout --detach FETCH_HEAD skip_if_unrelated_changes: parameters: category: @@ -2869,14 +2870,24 @@ jobs: destination: e2e-server-root-path-playwright-report build_docker_database_image: + parameters: + migration_qualification: + type: boolean + default: false machine: image: ubuntu-2204:2024.04.1 resource_class: large working_directory: ~/project steps: - - checkout - - checkout_migration_source - - skip_if_unrelated_changes + - when: + condition: << parameters.migration_qualification >> + steps: + - checkout_migration_source + - unless: + condition: << parameters.migration_qualification >> + steps: + - checkout + - skip_if_unrelated_changes - run: name: Build Docker image @@ -2923,7 +2934,6 @@ jobs: MIGRATION_TEST_OUTPUT: /tmp/migration-results PYTHONPATH: tests/e2e steps: - - checkout - checkout_migration_source - install_uv - install_rust @@ -3024,7 +3034,8 @@ workflows: migration_startup: when: << pipeline.parameters.run_migration_tests >> jobs: &migration_jobs - - build_docker_database_image + - build_docker_database_image: + migration_qualification: true - migration_startup_tests: name: migration-startup suite: startup From a37f0b4f544513d48001f1b2a86bd1eef59ca2c9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:49:58 -0700 Subject: [PATCH 003/127] test: isolate migration CI selection and exercise resolver boundaries --- .github/e2e-stack/select_tests.py | 2 +- .../litellm_proxy_extras/utils.py | 9 +- .../tests/test_setup_database_fail_fast.py | 156 +++++++----------- .../test_e2e_changed_gate.py | 5 + tests/e2e/conftest.py | 4 +- tests/e2e/migrations/checks.py | 12 +- tests/e2e/migrations/test_legacy.py | 7 +- tests/e2e/migrations/test_pooling.py | 3 +- tests/e2e/migrations/test_recovery.py | 4 +- tests/e2e/migrations/test_startup.py | 3 +- .../test_litellm_proxy_extras_utils.py | 12 +- 11 files changed, 93 insertions(+), 124 deletions(-) diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..9dd880c05cd 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -4,7 +4,7 @@ from typing import Final SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$") UNSUPPORTED: Final = re.compile( - r"^tests/e2e/(ui|claude_code|load)/" + r"^tests/e2e/(ui|claude_code|load|migrations)/" r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2749db5d754..8a83c786e02 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -859,7 +859,11 @@ class ProxyExtrasDBManager: def baseline_existing(migrations_dir: str) -> None: with migration_lock(lock_url) as coordinator: baseline_current_schema( - coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env()) + coordinator, + schema, + Path(migrations_dir), + _get_prisma_command(), + migration_environment(_get_prisma_env()), ) while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): @@ -1054,7 +1058,8 @@ class ProxyExtrasDBManager: migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr: logger.info( - "Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying", + "Migration %s deadlocked against a concurrent migrate deploy, " + "rolling its ledger row back and retrying", migration_name, ) ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 338c571eb4f..832075f6fbe 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -6,7 +6,8 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the """ import subprocess -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch import pytest @@ -31,10 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str): def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): """v2: a permission failure during migrate deploy raises RuntimeError.""" - monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) stderr = ( "Error: P3018\nMigration name: 20250326162113_baseline\n" @@ -47,10 +45,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): """v2: a non-idempotent migration failure raises (no silent recovery).""" - monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) stderr = ( "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n" @@ -131,8 +126,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path): def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path): """v2: a failing `prisma db push` must raise RuntimeError, not leak CalledProcessError past proxy_cli.py's `except RuntimeError`.""" - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) stderr = "db push error" with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): @@ -149,8 +143,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): import psycopg monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) class _FakeConn: def __enter__(self): @@ -173,18 +166,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): - _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists") - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"), - ) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs="relation already exists") stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists" with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) @@ -197,28 +179,15 @@ def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path): def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): - """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + _stub_v2_env(monkeypatch, tmp_path) + run = Mock(side_effect=_succeed_after(0, "")) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) - class FakeResult: - stdout = "Applied migration.\n" - stderr = "" - - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult()) - - resolve_called = {"n": 0} - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_all_migrations", - lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1), + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) is True + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], ) - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) - assert ok is True - assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" - _DEADLOCK_P3018_STDERR = ( "Error: P3018\n" @@ -228,12 +197,34 @@ _DEADLOCK_P3018_STDERR = ( ) -def _stub_v2_env(monkeypatch, tmp_path): +def _stub_v2_env(monkeypatch, tmp_path, ledger_logs=""): + import psycopg + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) monkeypatch.setattr("time.sleep", lambda _: None) + connection = MagicMock() + connection.__enter__.return_value = connection + cursor = connection.cursor.return_value.__enter__.return_value + cursor.execute.return_value = cursor + cursor.fetchone.return_value = SimpleNamespace(acquired=True) + cursor.fetchall.return_value = [] + empty = MagicMock() + empty.fetchall.return_value = [] + empty.fetchone.return_value = None + ledger = MagicMock() + ledger.fetchone.return_value = (ledger_logs,) + + def execute(query, *args, **kwargs): + if "SELECT logs FROM" in str(query): + if ledger_logs is None: + raise psycopg.OperationalError("ledger is unavailable") + return ledger + return empty + + connection.execute.side_effect = execute + monkeypatch.setattr("psycopg.connect", lambda *args, **kwargs: connection) def _succeed_after(failures: int, stderr: str): @@ -259,28 +250,21 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): instance rolls the ledger row back and retries instead of dying.""" _stub_v2_env(monkeypatch, tmp_path) - rolled_back = [] - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: rolled_back.append(name), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), - ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + run = Mock(side_effect=_succeed_after(1, _DEADLOCK_P3018_STDERR)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"], + ["migrate", "deploy"], + ) def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): """v2: a deadlock on every attempt still fails after the retry budget.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", @@ -293,7 +277,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path): """v2: the surviving instance sees the victim's failed ledger row as P3009. When that row's logs show a deadlock, roll it back and retry.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs="ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock") stderr = ( "Error: P3009\n" @@ -301,27 +285,16 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_failed_migration_logs", - lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock", - ) - rolled_back = [] - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: rolled_back.append(name), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - lambda name: pytest.fail("a deadlocked migration must never be marked applied"), - ) - monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) + run = Mock(side_effect=_succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True - assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == ( + ["migrate", "deploy"], + ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"], + ["migrate", "deploy"], + ) def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path): @@ -332,12 +305,6 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("empty logs do not prove rollback is safe"), - ) with patch( "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr) ) as run: @@ -350,7 +317,7 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): """v2: an unreadable ledger cannot establish that P3009 was a deadlock.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs=None) stderr = ( "Error: P3009\n" @@ -358,12 +325,6 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): "The `20260415120000_health_check_latest_per_model_index` migration " "started at 2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_v2_roll_back_migration_best_effort", - lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), - ) monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="Migration completion could not be verified"): @@ -372,7 +333,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" - _stub_v2_env(monkeypatch, tmp_path) + _stub_v2_env(monkeypatch, tmp_path, ledger_logs='ERROR: syntax error at or near "BRKN"') stderr = ( "Error: P3009\n" @@ -380,11 +341,6 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): "The `20260101000000_genuinely_broken` migration started at " "2026-09-01 18:46:13 UTC failed" ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_failed_migration_logs", - lambda name: 'ERROR: syntax error at or near "BRKN"', - ) with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="Migration completion could not be verified"): diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..a628fbb0633 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -112,6 +112,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: ( (("tests/e2e/logging/test_datadog_e2e.py", "litellm/router.py"), ("tests/e2e/logging/test_datadog_e2e.py",)), (("tests/e2e/ui/test_keys.py", "tests/e2e/claude_code/test_cli.py", "tests/e2e/load/test_burst.py"), ()), + (("tests/e2e/migrations/test_startup.py", "tests/e2e/migrations/test_recovery.py"), ()), (("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()), (("tests/e2e/guardrails/test_presidio_masking_e2e.py",), ()), (("tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py",), ()), @@ -157,6 +158,10 @@ def test_a_changed_canary_file_is_selected_once_alongside_a_harness_change() -> assert select_tests((CANARY[1], "tests/e2e/proxy_client.py")) == CANARY +def test_dedicated_migration_tests_do_not_suppress_shared_harness_canaries() -> None: + assert select_tests(("tests/e2e/migrations/test_startup.py", "tests/e2e/conftest.py")) == CANARY + + def test_the_canary_joins_directly_selected_files_in_sorted_order() -> None: assert select_tests(("tests/e2e/logging/test_datadog_e2e.py", ".github/e2e-stack/up.sh")) == ( *CANARY, diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4a5f0aa880f..53d35effdc9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -64,7 +64,9 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow") + config.addinivalue_line( + "markers", "migration_startup: isolated container startup tests run by the migration CI workflow" + ) config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py index b14631ccad8..619ad3b0e6c 100644 --- a/tests/e2e/migrations/checks.py +++ b/tests/e2e/migrations/checks.py @@ -29,7 +29,8 @@ def start_replicas( def assert_completed(database: Database, migration: Migration = COMPLETE) -> None: assert database.query( - "SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + 'SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM ' + '_prisma_migrations WHERE migration_name = %s', (migration.name,), ) == ((True, True, 1),), "Expected exactly one successful SQL execution" assert database.query("SELECT id FROM migration_effect") == ((1,),) @@ -47,7 +48,8 @@ def confirmed_history(database: Database) -> str: def assert_original_proof(database: Database, row_id: str, finished: bool) -> None: assert database.query( - "SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", + 'SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM ' + '_prisma_migrations WHERE migration_name = %s', (COMPLETE.name,), ) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof" assert database.query("SELECT id FROM migration_effect") == ((1,),) @@ -59,7 +61,8 @@ def pause_completion(database: Database) -> None: "CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " "IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN " "PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; " - "CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()" + 'CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW ' + 'EXECUTE FUNCTION migration_pause()' ).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY)) ) @@ -105,7 +108,8 @@ def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None: failed(replicas, "Migration completion could not be verified") started: Final = str( database.query( - "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s", + "SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM " + '_prisma_migrations WHERE migration_name = %s', (COMPLETE.name,), )[0][0] ) diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py index 7ba73eb82e0..ba5e77a3070 100644 --- a/tests/e2e/migrations/test_legacy.py +++ b/tests/e2e/migrations/test_legacy.py @@ -15,7 +15,9 @@ def adopt_legacy(containers: Containers, database: Database) -> None: count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0] existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token') database.execute( - "INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" + 'INSERT INTO "LiteLLM_ShadowEvalJob" (id, group_id, target_id, router_name, judge_model, ' + "shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', " + "'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())" ) database.execute("DROP TABLE _prisma_migrations") with ExitStack() as stack: @@ -30,7 +32,8 @@ def adopt_legacy(containers: Containers, database: Database) -> None: assert detail in logs assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),) assert database.query( - "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0" + 'SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS ' + 'NOT NULL OR applied_steps_count <> 0' ) == ((0,),) assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')) assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == ( diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py index e0a3693b33e..c4015549de3 100644 --- a/tests/e2e/migrations/test_pooling.py +++ b/tests/e2e/migrations/test_pooling.py @@ -50,7 +50,8 @@ def pool(database: Database, output: Path) -> Generator[str]: f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n" "[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n" "pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n" - "max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n" + 'max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = ' + 'extra_float_digits,options\n' ) try: docker( diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py index 58bf5c348d6..80e5747eaac 100644 --- a/tests/e2e/migrations/test_recovery.py +++ b/tests/e2e/migrations/test_recovery.py @@ -159,7 +159,9 @@ class TestMigrationRecovery: ) case "duplicate_history": database.execute( - "INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", + 'INSERT INTO _prisma_migrations (id, migration_name, checksum, ' + 'applied_steps_count) SELECT %s, migration_name, checksum, ' + 'applied_steps_count FROM _prisma_migrations WHERE migration_name = %s', (str(uuid4()), COMPLETE.name), ) case "missing_script": diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index a648218cb26..dc628a8ed7a 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -56,7 +56,8 @@ class TestMigrationStartup: replicas: Final = start_replicas(stack, containers, database, (FATAL,)) failed(replicas, COMPLETE.name) assert database.query( - "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL", + 'SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE ' + '%s AND finished_at IS NULL', (COMPLETE.name, "%MIGRATION_TEST_FATAL%"), ) == ((1,),) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 57133ea95c4..bb329264a11 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -771,17 +771,7 @@ class _MigrateDeployHarness: self.confirmed_migrations = set(confirmed_migrations) monkeypatch.delenv("DATABASE_URL", raising=False) - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_roll_back_migration", - staticmethod(lambda name: None), - ) - monkeypatch.setattr( - ProxyExtrasDBManager, - "_resolve_specific_migration", - staticmethod(self.resolved.append), - ) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) From 20d80b5420508c73391cca91be232b7f74041d1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:45:07 -0700 Subject: [PATCH 004/127] fix(guardrails): scan each choice's tool-call arguments apart on n>1 streams and log why a rewrite was discarded The rebuilt streamed response keyed tool-call fragments by tool index alone, so on n>1 chat streams the two choices' argument fragments were concatenated into one string and post_call guardrails scanned garbled JSON. Fragments are now keyed by (choice index, tool index). When a guardrail's rewrite cannot be written back to the stream (multi-choice streams, a rewrite that adds or drops a tool call, legacy-hook shapes the translation cannot rescan), the pipeline now logs a warning naming the guardrail and the exact reason before releasing the original stream. Also commits the regenerated dashboard API types that make check produced. --- .../streaming_chunk_builder_utils.py | 38 +++++---- .../chat/guardrail_translation/handler.py | 11 ++- .../chat/guardrail_translation/handler.py | 26 +++++-- .../guardrail_translation/handler.py | 11 ++- .../proxy/policy_engine/pipeline_executor.py | 78 ++++++++++++++----- .../test_streaming_chunk_builder_utils.py | 55 +++++++++++++ .../test_openai_guardrail_handler.py | 57 +++++++++++++- .../policy_engine/test_pipeline_executor.py | 42 ++++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 9 files changed, 254 insertions(+), 66 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..5ffe36573d5 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -138,9 +138,13 @@ class _ToolCallDelta(TypedDict, total=False): class _ToolCallChoice(TypedDict, total=False): + index: ReadOnly[int] delta: ReadOnly[_ToolCallDelta] +_ToolCallKey: TypeAlias = tuple[int, int] + + class _ToolCallChunk(TypedDict): choices: ReadOnly[Sequence[_ToolCallChoice]] @@ -416,40 +420,41 @@ class ChunkProcessor: @staticmethod def _iter_tool_call_fragments( tool_call_chunks: Sequence["_ToolCallChunk"], - ) -> Iterator[tuple[int, str, str]]: + ) -> Iterator[tuple[_ToolCallKey, str, str]]: for chunk in tool_call_chunks: for choice in chunk["choices"]: delta = choice.get("delta") if not delta: continue + choice_index = choice.get("index", 0) for tool_call in delta.get("tool_calls", ()): if not tool_call: continue if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + key = (choice_index, tool_call.get("index", 0)) function = tool_call.get("function") if isinstance(function, dict): if fragment_arguments := function.get("arguments"): - yield index, "arguments", fragment_arguments + yield key, "arguments", fragment_arguments elif function_arguments := getattr(function, "arguments", None): - yield index, "arguments", function_arguments + yield key, "arguments", function_arguments custom = tool_call.get("custom") if isinstance(custom, dict) and (custom_input := custom.get("input")): - yield index, "custom_input", custom_input + yield key, "custom_input", custom_input else: - index = getattr(tool_call, "index", 0) + key = (choice_index, getattr(tool_call, "index", 0)) function = getattr(tool_call, "function", None) if object_arguments := getattr(function, "arguments", None): - yield index, "arguments", object_arguments + yield key, "arguments", object_arguments custom = getattr(tool_call, "custom", None) if object_custom_input := getattr(custom, "input", None): - yield index, "custom_input", object_custom_input + yield key, "custom_input", object_custom_input @staticmethod - def _join_fragments_by_index_and_field( - fragment_records: Iterator[tuple[int, str, str]], - ) -> Mapping[tuple[int, str], str]: - def group_key(record: tuple[int, str, str]) -> tuple[int, str]: + def _join_fragments_by_key_and_field( + fragment_records: Iterator[tuple[_ToolCallKey, str, str]], + ) -> Mapping[tuple[_ToolCallKey, str], str]: + def group_key(record: tuple[_ToolCallKey, str, str]) -> tuple[_ToolCallKey, str]: return record[0], record[1] return MappingProxyType( @@ -467,13 +472,14 @@ class ChunkProcessor: tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[int, dict[str, Any]]] = {} # Map to store tool calls by index + tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} # Map to store tool calls by choice and index for chunk in tool_call_chunks: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) tool_calls = delta.get("tool_calls", []) + choice_index = choice.get("index", 0) for tool_call in tool_calls: # Handle both dict and object formats @@ -495,9 +501,9 @@ class ChunkProcessor: # Get index (handle both dict and object) if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + index = (choice_index, tool_call.get("index", 0)) else: - index = getattr(tool_call, "index", 0) + index = (choice_index, getattr(tool_call, "index", 0)) if index not in tool_call_map: tool_call_map[index] = { @@ -572,7 +578,7 @@ class ChunkProcessor: if isinstance(provider_fields, dict): merged_provider_fields.update(provider_fields) - joined_fragments: Final = self._join_fragments_by_index_and_field( + joined_fragments: Final = self._join_fragments_by_key_and_field( self._iter_tool_call_fragments(tool_call_chunks) ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..c7ba5daec56 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1172,7 +1172,10 @@ class AnthropicMessagesHandler(BaseTranslation): if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream never reported a stop_reason, so the text rewrite has no assembled response to land on", + ) return responses_so_far def _prepare_request_data( @@ -1318,7 +1321,11 @@ class AnthropicMessagesHandler(BaseTranslation): if len(block_indices) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(block_indices)} tool_use blocks", + ) rewrites_by_block: Final = MappingProxyType( { index: after diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..e4943690639 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1041,13 +1041,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) if len(stream_choice_indices) != 1: - # stream_chunk_builder collapses every choice into one index-0 - # choice, so a rewrite of the rebuilt response cannot be attributed - # back to a single choice on an n>1 stream: report it undeliverable - # rather than deliver the rewrite on the wrong choice from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and the rebuilt response's text rewrite " + "cannot be attributed to one of them", + ) target_choice_index: Final = next(iter(stream_choice_indices)) await self._apply_guardrail_responses_to_output_streaming( responses=responses_so_far, @@ -1105,10 +1105,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) - if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + if len(stream_choice_indices) != 1: from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " + "back on single-choice streams", + ) + if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(fragments_by_tool_call)}", + ) for before, (name, arguments), fragments in zip( pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call ): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..2be36a826f7 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -957,7 +957,10 @@ class OpenAIResponsesHandler(BaseTranslation): if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream carried no terminal response envelope to write the text rewrite back into", + ) return responses_so_far @staticmethod @@ -1070,7 +1073,11 @@ class OpenAIResponsesHandler(BaseTranslation): ): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls and the stream's " + f"{len(tool_call_items)} function_call items could not be lined up with them by call_id", + ) for output_item, rewrite in ( (output_item, rewrites_by_call_id[call_id]) for output_item, call_id in zip(tool_call_items, call_ids) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ad45781d5d2..26dd806b3e5 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -50,12 +50,13 @@ except ImportError: class UndeliverableStreamRewrite(Exception): - def __init__(self, guardrail_name: str) -> None: + def __init__(self, guardrail_name: str, reason: str) -> None: super().__init__( - f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's " - "streaming pipeline cannot deliver" + f"Guardrail '{guardrail_name}' rewrote the streamed response but the rewrite cannot be written " + f"back to the stream: {reason}" ) self.guardrail_name: Final = guardrail_name + self.reason: Final = reason class UnappliableRequestRewrite(Exception): @@ -91,8 +92,22 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent -def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: - return sent is not None and returned is not None and len(returned) != len(sent) +def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> tuple[int, int] | None: + if sent is None or returned is None or len(returned) == len(sent): + return None + return (len(sent), len(returned)) + + +def _tool_call_mismatch_reason( + sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None +) -> str | None: + if sent == returned: + return None + sent_count: Final = len(sent or ()) + returned_count: Final = len(returned or ()) + if sent_count == returned_count: + return "the legacy hook changed a tool call's name or arguments, which this path cannot write back" + return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}" _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -119,7 +134,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False - self.changed_tool_call_count = False + self.tool_call_count_change: tuple[int, int] | None = None def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -140,11 +155,22 @@ class _StreamRewriteObserver(CustomGuardrail): returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) - self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + self.tool_call_count_change = self.tool_call_count_change or _count_change( sent_tool_shapes, returned_tool_shapes ) return outputs + def discard_reason(self, deliver_rewrites: bool) -> str | None: + if self.tool_call_count_change is not None: + sent, returned = self.tool_call_count_change + return ( + f"the guardrail returned {returned} tool calls for a stream that carried {sent}, and a rewrite " + "that drops or adds a tool call cannot be written back" + ) + if not deliver_rewrites and (self.rewrote_texts or self.rewrote_tool_calls): + return "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + return None + class _ScannedTextRecorder(CustomGuardrail): def __init__(self, guardrail_name: str) -> None: @@ -209,13 +235,24 @@ class _LegacyHookStreamAdapter(CustomGuardrail): if rewrite is None: return inputs rescanned: Final = await self._rescan(rewrite, logging_obj) + guardrail_name: Final = self.guardrail_name or "unknown" if rescanned is None: - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_name, "the legacy hook's response could not be rescanned by this endpoint's translation" + ) rewritten: Final = rescanned.get("texts") - if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") - if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + returned_text_count: Final = len(_scanned_texts(rewritten)) + sent_text_count: Final = len(_scanned_texts(inputs.get("texts"))) + if returned_text_count != sent_text_count: + raise UndeliverableStreamRewrite( + guardrail_name, + f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}", + ) + tool_call_mismatch: Final = _tool_call_mismatch_reason( + _tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls")) + ) + if tool_call_mismatch is not None: + raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch) if not rewritten: return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} @@ -262,14 +299,16 @@ def _prepare_hook_input( def _release_original_chunks( guardrail_name: str, + reason: str, streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place originals: Sequence[object], ) -> None: streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives verbose_proxy_logger.warning( - "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming " - "pipeline cannot deliver yet; the rewrite was discarded and the original stream released", + "Pipeline: guardrail '%s' rewrote the streamed response but the rewrite could not be written back to " + "the stream: %s. The whole rewrite, text rewrites included, was discarded and the original stream released", guardrail_name, + reason, ) @@ -442,13 +481,12 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, request_data=hook_input, ) - except UndeliverableStreamRewrite: - _release_original_chunks(step.guardrail, streaming_chunks, originals) + except UndeliverableStreamRewrite as undeliverable: + _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return - if observer.changed_tool_call_count or ( - not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) - ): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + discard_reason: Final = observer.discard_reason(deliver_rewrites) + if discard_reason is not None: + _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index efe4209c1c9..2266258bf20 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1288,6 +1288,61 @@ def _tool_call_delta_chunk(tool_call: dict[str, object] | ChatCompletionDeltaToo return {"choices": [{"delta": {"tool_calls": [tool_call]}}]} +def _choice_tool_call_delta_chunk(choice_index: int, tool_call: dict[str, object]) -> dict[str, object]: + return {"choices": [{"index": choice_index, "delta": {"tool_calls": [tool_call]}}]} + + +def test_get_combined_tool_content_keeps_each_choices_arguments_apart_when_choices_share_a_tool_index(): + processor = ChunkProcessor.__new__(ChunkProcessor) + chunks = [ + _choice_tool_call_delta_chunk(0, {"index": 0, "id": "call_a", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "id": "call_b", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": '{"fruit": "pers'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": '{"fruit": "dur'}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": 'immon"}'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": 'ian"}'}}), + ] + + combined = processor.get_combined_tool_content(chunks) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in combined] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + +def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart(): + def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + object="chat.completion.chunk", + created=1234567890, + model="gpt-4.1-mini", + choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]), finish_reason=None)], + ) + + def fragment(arguments: str, name: str | None = None, call_id: str | None = None) -> ChatCompletionDeltaToolCall: + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + response = stream_chunk_builder( + chunks=[ + chunk(0, fragment("", name="lookup_fruit", call_id="call_a")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_b")), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), + ] + ) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in response.choices[0].message.tool_calls] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order(): processor = ChunkProcessor.__new__(ChunkProcessor) first_fragments = [f"a{i};" for i in range(300)] diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5a29a96829f..60a5752e83a 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1267,7 +1267,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=self._world_masking_guardrail(), @@ -1275,6 +1275,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test-mask" + assert raised.value.reason == ( + "the stream carries 2 choices and the rebuilt response's text rewrite cannot be attributed to one of them" + ) + @staticmethod def _two_choice_tool_call_stream_chunks() -> list: from litellm.types.utils import ( @@ -1310,12 +1315,51 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return [ chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), - chunk(0, fragment('{"fruit": "persimmon"}')), - chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), chunk(0, None, finish_reason="tool_calls"), chunk(1, None, finish_reason="tool_calls"), ] + @staticmethod + def _recording_guardrail() -> CustomGuardrail: + class Recorder(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="recorder") + self.seen_inputs: list[GenericGuardrailAPIInputs] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_inputs.append(inputs) + return inputs + + return Recorder() + + @pytest.mark.asyncio + async def test_ended_multi_choice_stream_scans_each_choices_tool_call_arguments_apart(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + guardrail = self._recording_guardrail() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert [ + (tool_call["id"], tool_call["function"]["arguments"]) + for tool_call in guardrail.seen_inputs[-1]["tool_calls"] + ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1323,7 +1367,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_tool_call_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=MockGuardrail(guardrail_name="test"), @@ -1331,6 +1375,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test" + assert raised.value.reason == ( + "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 0a2641082dc..e7689cc7d0c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1122,7 +1122,7 @@ class _RefusingTranslation: deliver_ended_stream_rewrites=False, ): responses_so_far[0]["text"] = "half-written" - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name) + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name, "the translation refused it") def _chunk(): @@ -1143,10 +1143,22 @@ async def _run_streaming_step(translation, streaming_chunks=None): ) -def _assert_passed_with_discard_warning(result, caplog): +NO_WRITE_BACK_REASON = "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + + +def _assert_passed_with_discard_warning(result, caplog, reason): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] - assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + discard_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + and "'masker'" in record.getMessage() + and "discarded" in record.getMessage() + ] + assert len(discard_warnings) == 1 + assert reason in discard_warnings[0] + assert "text rewrites included" in discard_warnings[0] assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @@ -1159,7 +1171,7 @@ async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(translation, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] assert translation.seen_guardrail_names == ["masker"] @@ -1196,7 +1208,7 @@ async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1258,7 +1270,9 @@ async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the guardrail returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1270,7 +1284,7 @@ async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_ with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1326,7 +1340,7 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_RefusingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the translation refused it") assert chunks == [_chunk()] @@ -1561,7 +1575,7 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook returned 2 texts for a stream that carried 1") assert chunks == [_chunk()] @@ -1576,7 +1590,7 @@ async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_chunk()] @@ -1588,7 +1602,9 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the legacy hook returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1620,7 +1636,7 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() ) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_tool_only_chunk()] @@ -1658,7 +1674,7 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook's response could not be rescanned") assert chunks == [_chunk()] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 65160a97c54da63f24bf674f4f03551b22ac97c3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:55:03 -0700 Subject: [PATCH 005/127] fix(guardrails): keep tool calls carried by a later choice of a packed multi-choice chunk The rebuild's tool-call selection and its text-only fast path only looked at choice 0 of each chunk, so a chunk that packs several choices (Gemini with candidateCount above 1) lost a tool call carried by a later candidate, and a chunk whose later choice had no tool calls at all made the rebuild raise. Both now consider every choice in the chunk. --- .../streaming_chunk_builder_utils.py | 4 +- litellm/main.py | 66 +++++++++++-------- tests/test_litellm/test_main.py | 62 +++++++++++++++++ 3 files changed, 104 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 5ffe36573d5..f5b723755aa 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -427,7 +427,7 @@ class ChunkProcessor: if not delta: continue choice_index = choice.get("index", 0) - for tool_call in delta.get("tool_calls", ()): + for tool_call in delta.get("tool_calls") or (): if not tool_call: continue if isinstance(tool_call, dict): @@ -478,7 +478,7 @@ class ChunkProcessor: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + tool_calls = delta.get("tool_calls") or () choice_index = choice.get("index", 0) for tool_call in tool_calls: diff --git a/litellm/main.py b/litellm/main.py index 17edafcdfca..a4a648acd4e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8749,6 +8749,39 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o setattr(usage, "cost", computed_cost) +_NON_TEXT_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "annotations", + "audio", + "images", + "provider_specific_fields", +) + + +def _stream_choice_delta(choice: object) -> Mapping[str, object]: + delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) + if isinstance(delta, Mapping): + return delta + if isinstance(delta, BaseModel): + return delta.model_dump() + return {} + + +def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool: + return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS) + + +def _simple_text_part(choices: Sequence[object]) -> str | None: + deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices) + if any(_delta_carries_more_than_text(delta) for delta in deltas): + return None + content: Final = deltas[0].get("content") + return content if isinstance(content, str) else "" + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -8793,31 +8826,11 @@ def stream_chunk_builder( if not chunk.get("choices"): continue - choice = chunk["choices"][0] - delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) - if isinstance(delta_obj, dict): - delta = delta_obj - elif hasattr(delta_obj, "model_dump"): - delta = cast(dict[str, Any], delta_obj.model_dump()) - else: - delta = {} - - if ( - delta.get("tool_calls") is not None - or delta.get("function_call") is not None - or delta.get("reasoning_content") is not None - or delta.get("thinking_blocks") is not None - or delta.get("annotations") is not None - or delta.get("audio") is not None - or delta.get("images") is not None - or delta.get("provider_specific_fields") is not None - ): + if (part := _simple_text_part(chunk["choices"])) is None: is_simple_text_stream = False break - - content = delta.get("content") - if isinstance(content, str) and content: - simple_content_parts.append(content) + if part: + simple_content_parts.append(part) if is_simple_text_stream: if simple_content_parts: @@ -8854,9 +8867,10 @@ def stream_chunk_builder( tool_call_chunks: Final = [ chunk for chunk in chunks - if chunk.get("choices") - and "tool_calls" in chunk["choices"][0]["delta"] - and chunk["choices"][0]["delta"]["tool_calls"] is not None + if any( + "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None + for choice in chunk.get("choices") or () + ) ] if len(tool_call_chunks) > 0: diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4f7a51eb531..42599f45ade 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1659,6 +1659,68 @@ async def test_async_mock_delay(): assert delay >= 0.01 +def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): + from litellm import stream_chunk_builder + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-multi-choice", + created=1751934860, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=choices, + ) + + chunks = [ + chunk( + [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), + StreamingChoices( + index=1, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + index=0, + type="function", + function=Function(name="lookup_fruit", arguments='{"fruit":'), + ) + ], + ), + ), + ] + ), + chunk( + [ + StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), + StreamingChoices( + index=1, + delta=Delta( + tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] + ), + finish_reason="tool_calls", + ), + ] + ), + ] + + response = stream_chunk_builder(chunks=chunks) + + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ + ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') + ] + + def test_stream_chunk_builder_thinking_blocks(): from litellm import stream_chunk_builder from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices From 0ac0362b42890d5d40698dbd5c206065b0015e5f Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 21:01:16 +0000 Subject: [PATCH 006/127] fix(proxy): enforce virtual key budgets for JEV test routing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 2 +- .../test_auto_router_endpoints.py | 76 ++++++++++++++++++- 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 200ed6c3bf3..f79425d2e97 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -319,7 +319,7 @@ async def _authorize_models_this_test_can_call( its calls through the proxy. Team and member budgets are already enforced on every route. """ models: Final = _models_this_test_can_call(config) - if not models: + if not models and config.classifier_type != "jev": return from litellm.proxy.proxy_server import proxy_logging_obj diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 067f30c2fd7..36130137c64 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -3,23 +3,33 @@ Unit tests for auto router management endpoints """ from collections.abc import Mapping, Sequence +from functools import partial from pathlib import Path from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.management_endpoints import auto_router_endpoints from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router +from litellm.router_strategy.complexity_router import ComplexityRouter +from litellm.router_strategy.complexity_router.jev_classifier import ( + JevChoiceAnswer, + JevClassifierClient, + JevSystemOneResponse, +) from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, @@ -422,8 +432,67 @@ async def test_a_key_over_its_budget_cannot_run_a_classifier_config(monkeypatch: assert calls == [] +@pytest.mark.parametrize( + "max_budget, spend, denied", + ( + pytest.param(0.0, 0.0, True, id="zero-budget"), + pytest.param(1.0, 1.0, True, id="budget-reached"), + pytest.param(1.0, 2.0, True, id="budget-exceeded"), + pytest.param(1.0, 0.5, False, id="budget-remaining"), + pytest.param(None, 2.0, False, id="unlimited"), + ), +) @pytest.mark.asyncio -async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.MonkeyPatch): +async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-budget-test", + user_id="admin", + models=["cheap-model"], + max_budget=max_budget, + spend=spend, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + assert exc_info.value.param is None + assert "Budget has been exceeded!" in exc_info.value.message + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routed_model == "cheap-model" + assert response.routing_decision["cause"] == "jev_classifier" + assert response.routing_decision["classifier_model"] == "typesafe/jev-test" + client.evaluate.assert_awaited_once() + + +@pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) +@pytest.mark.asyncio +async def test_a_heuristic_config_does_not_need_a_budget( + monkeypatch: pytest.MonkeyPatch, max_budget: float, spend: float +): import litellm.proxy.proxy_server as proxy_server monkeypatch.setattr(proxy_server, "llm_router", _router()) @@ -435,8 +504,8 @@ async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.Mon user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-broke", user_id="admin", - max_budget=1.0, - spend=2.0, + max_budget=max_budget, + spend=spend, models=["cheap-model"], ), ) @@ -851,7 +920,6 @@ class TestAutoRouterBenchmarks: # --------------------------------------------------------------------------- from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock from litellm.proxy.management_endpoints.auto_router_endpoints import ( get_shadow_eval_job, From b9e5bb3abb0f2f0ddc06dcaf9edb63563cac2a2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 22:03:37 +0000 Subject: [PATCH 007/127] test(proxy): allow JEV dependency in budget fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_auto_router_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 36130137c64..a5c93c41a84 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -459,7 +459,7 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-jev-budget-test", user_id="admin", - models=["cheap-model"], + models=["cheap-model", "typesafe/jev-test"], max_budget=max_budget, spend=spend, ) From 2ee6fb5dc45d02a3a4df78b836bd371ed71ab375 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 22:08:21 -0700 Subject: [PATCH 008/127] ci: skip the integration matrix during migration-qualification pipelines --- .circleci/config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index d44850234a7..4e65b5cddcd 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3151,6 +3151,7 @@ workflows: only: litellm_internal_staging jobs: *migration_jobs integration: + unless: << pipeline.parameters.run_migration_tests >> jobs: - integration_contracts: name: integration-<< matrix.suite >> From 2863559ba8a422df87981e108c486316927ea859 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 08:36:22 -0700 Subject: [PATCH 009/127] fix(proxy): wait for the spend-log table before creating startup views On a fresh database where the migrations run in a separate job while the proxy boots with DISABLE_SCHEMA_UPDATE=true, the startup view check ran as a fire-and-forget task, used up its three 10 second retries before LiteLLM_SpendLogs existed, and died with an unretrieved exception. The spend views were never created, so the /global/spend routes returned 500 until the pod was restarted PrismaClient now holds a view setup task. It polls to_regclass for the spend-log table every 5 seconds, creates the views and loads the spend log row count once the table is there, keeps polling if an attempt raises while the schema is still settling, and logs an ERROR with the last failure if nothing worked after 15 minutes. Proxy shutdown cancels the task The spend route e2e tests for the five view-backed routes are no longer skipped and wait for the views through the harness convergence helper --- litellm/proxy/proxy_server.py | 15 +- litellm/proxy/utils.py | 83 ++++++ .../spend_tracking/spend_e2e_client.py | 13 + .../spend_tracking/test_spend_routes.py | 31 +-- tests/test_litellm/proxy/test_proxy_server.py | 63 ++++- .../test_prisma_client_lifecycle.py | 239 ++++++++++++++++++ 6 files changed, 418 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36fdea605c2..a4d99d9859e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1464,6 +1464,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): + try: + await prisma_client.stop_view_setup_task() + except Exception as e: + verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) + await _drain_spend_event_producer_on_shutdown() await flush_spend_counters_on_shutdown() @@ -10635,14 +10641,7 @@ class ProxyStartupEvent: if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): await prisma_client.db.start_token_refresh_task() - ## Add necessary views to proxy ## - asyncio.create_task( - prisma_client.check_view_exists() - ) # check if all necessary views exist. Don't block execution - - asyncio.create_task( - prisma_client._set_spend_logs_row_count_in_proxy_state() - ) # set the spend logs row count in proxy state. Don't block execution + prisma_client.start_view_setup_task() if hasattr(prisma_client, "start_db_health_watchdog_task"): await prisma_client.start_db_health_watchdog_task() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b078a65759e..7f5609b2e39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,6 +38,7 @@ from typing import ( Literal, Optional, Protocol, + TypeAlias, TypeVar, Union, cast, @@ -268,6 +269,15 @@ class _RelTuplesRow(TypedDict): reltuples: ReadOnly[int] +_VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 +_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 +_VIEW_SETUP_GATE_TABLE: Final = "LiteLLM_SpendLogs" +_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) + +_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] +_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception + + class _EndUserBatchTable(Protocol): def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... @@ -4284,6 +4294,7 @@ class PrismaClient: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: asyncio.Task | None = None + self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None self._db_last_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) self._db_read_only_recreate_ts: float = 0.0 @@ -6330,6 +6341,78 @@ class PrismaClient: self._db_health_watchdog_task = None verbose_proxy_logger.info("Stopped Prisma DB health watchdog") + def start_view_setup_task(self) -> None: + if self._view_setup_task is not None: + return + self._view_setup_task = asyncio.create_task(self._run_view_setup()) + + async def stop_view_setup_task(self) -> None: + if self._view_setup_task is None: + return + self._view_setup_task.cancel() + try: + await self._view_setup_task + except asyncio.CancelledError: + pass + self._view_setup_task = None + + async def _run_view_setup( + self, + poll_interval_seconds: float = _VIEW_SETUP_POLL_INTERVAL_SECONDS, + deadline_seconds: float = _VIEW_SETUP_DEADLINE_SECONDS, + ) -> _ViewSetupOutcome: + deadline: Final = time.monotonic() + deadline_seconds + while True: + if (attempt := await self._attempt_view_setup()) == "ready": + return "ready" + if time.monotonic() >= deadline: + self._log_view_setup_timeout(attempt, deadline_seconds) + return "timed_out" + await asyncio.sleep(poll_interval_seconds) + + async def _attempt_view_setup(self) -> _ViewSetupAttempt: + try: + if not await self._view_setup_gate_table_present(): + verbose_proxy_logger.debug( + "Waiting for table %s before creating the spend views", self._view_setup_gate_table() + ) + return "table_missing" + await self.check_view_exists() + await self._set_spend_logs_row_count_in_proxy_state() + return "ready" + except Exception as e: + verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e) + return e + + def _log_view_setup_timeout( + self, last_attempt: Literal["table_missing"] | Exception, deadline_seconds: float + ) -> None: + if isinstance(last_attempt, Exception): + verbose_proxy_logger.error( + "Gave up creating the spend views after %ss; the last attempt failed with: %s. " + "Fix that error and restart the proxy.", + deadline_seconds, + last_attempt, + ) + return + verbose_proxy_logger.error( + "Gave up creating the spend views: table %s did not appear within %ss. " + "Run the database migrations against this database and restart the proxy.", + self._view_setup_gate_table(), + deadline_seconds, + ) + + async def _view_setup_gate_table_present(self) -> bool: + rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( + await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", self._view_setup_gate_table()) + ) + return rows[0]["present"] + + @staticmethod + def _view_setup_gate_table() -> str: + pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") + return f'"{pg_schema}"."{_VIEW_SETUP_GATE_TABLE}"' + async def _db_health_watchdog_loop(self) -> None: while True: try: diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 9ac97f57f47..233aa81c2af 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -359,6 +359,19 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: + """Re-probe a route that depends on startup work the proxy finishes after it + starts serving, such as the spend views it creates once migrations land.""" + outcome: Final = await_converged( + lambda: self.probe(path, params=params), + converged=lambda result: result.healthy, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: return unwrap( self.proxy.transport.post( diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py index 8cb3e3927f0..67fd88bc84d 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py @@ -17,9 +17,11 @@ fast: no batch-write wait, no provider calls. """ from datetime import datetime, timedelta, timezone +from typing import Final import pytest +from e2e_http import ProbeResult from models import DateRangeParams from spend_e2e_client import SpendClient @@ -72,15 +74,10 @@ SPEND_ROUTES = ( _SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") -_MISSING_VIEW_SKIP = pytest.mark.skip( - reason=( - "LIT-5211: on a fresh database the proxy's startup view creation can lose the race " - "against schema migrations, leaving MonthlyGlobalSpend/DailyTagSpend/Last30d* views " - "missing and these routes 500ing until the views exist" - ) -) - -_VIEW_BACKED_ROUTES = frozenset( +# Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the +# proxy creates in the background once the schema migrations have landed, so on a +# fresh database they can 500 for a while after the proxy starts serving. +_VIEW_BACKED_ROUTES: Final = frozenset( ( "/global/spend", "/global/spend/keys", @@ -98,15 +95,15 @@ def _date_range() -> DateRangeParams: return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) -@pytest.mark.parametrize( - "route", - tuple( - pytest.param(route, marks=_MISSING_VIEW_SKIP) if route in _VIEW_BACKED_ROUTES else route - for route in SPEND_ROUTES - ), -) +def _probe(client: SpendClient, route: str) -> ProbeResult: + if route in _VIEW_BACKED_ROUTES: + return client.probe_until_healthy(route, params=_date_range()) + return client.probe(route, params=_date_range()) + + +@pytest.mark.parametrize("route", SPEND_ROUTES) def test_spend_route_responsive(client: SpendClient, route: str) -> None: - result = client.probe(route, params=_date_range()) + result = _probe(client, route) print(f"{route} -> {result.status_code}\n{result.body[:600]}") assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e2e2045e826..7ab142c5fd0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1628,6 +1628,40 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key +class _ShutdownAwarePrisma(MockPrisma): + def __init__(self): + super().__init__() + self.stop_view_setup_task = AsyncMock() + + +@pytest.mark.asyncio +async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): + """The view setup task keeps polling for the spend-log table while migrations + run, so a shutdown inside that window has to cancel it rather than leave it + to die with the event loop.""" + import yaml + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + fake_prisma = _ShutdownAwarePrisma() + config_path = tmp_path / "config.yaml" + with open(config_path, "w") as f: + yaml.dump({"general_settings": {"master_key": "sk-12345"}}, f) + monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path)) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) + + async with proxy_startup_event(FastAPI()): + stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + + actual = { + "stopped_while_serving": stopped_while_serving, + "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + } + assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} + + def test_team_info_masking(): """ Test that sensitive team information is properly masked @@ -13307,6 +13341,7 @@ def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + client.start_view_setup_task = MagicMock() client.start_db_health_watchdog_task = AsyncMock() client.health_check = AsyncMock(side_effect=health_check_error) return client @@ -13368,13 +13403,39 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) call_order = MagicMock() + call_order.attach_mock(mock_client.start_view_setup_task, "view_setup") call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog") call_order.attach_mock(mock_client.health_check, "health_check") await _run_setup_prisma_client(mock_client) assert mock_client.start_db_health_watchdog_task.await_count == 1 - assert [call[0] for call in call_order.mock_calls] == ["watchdog", "health_check"] + assert [call[0] for call in call_order.mock_calls] == ["view_setup", "watchdog", "health_check"] + + +@pytest.mark.asyncio +async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): + """View creation used to be two fire-and-forget ``asyncio.create_task`` calls + that raised and vanished when the migrations Job had not created + ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's + held task, which waits for the table, and must not call the two coroutines directly.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + + mock_client = _mock_startup_prisma_client() + result = await _run_setup_prisma_client(mock_client) + + actual = { + "result": result, + "view_setup_started": mock_client.start_view_setup_task.call_count, + "direct_view_creation": mock_client.check_view_exists.await_count, + "direct_row_count": mock_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "result": mock_client, + "view_setup_started": 1, + "direct_view_creation": 0, + "direct_row_count": 0, + } @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 18b02ac7772..bb34486771c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -5,11 +5,15 @@ Symbols pinned here: - ``PrismaClient.writer_db`` - ``PrismaClient.connect`` - ``PrismaClient.disconnect`` + - ``PrismaClient.start_view_setup_task`` + - ``PrismaClient.stop_view_setup_task`` + - ``PrismaClient._run_view_setup`` """ from __future__ import annotations import asyncio +import logging from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -17,6 +21,27 @@ import pytest from litellm.proxy.utils import PrismaClient +_PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" + + +def _absent() -> list[dict[str, bool]]: + return [{"present": False}] + + +def _present() -> list[dict[str, bool]]: + return [{"present": True}] + + +def _wire_view_setup(prisma_client: PrismaClient, probe: AsyncMock) -> MagicMock: + prisma_client.db.query_raw = probe + prisma_client.check_view_exists = AsyncMock() + prisma_client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + call_order = MagicMock() + call_order.attach_mock(probe, "probe") + call_order.attach_mock(prisma_client.check_view_exists, "views") + call_order.attach_mock(prisma_client._set_spend_logs_row_count_in_proxy_state, "row_count") + return call_order + @pytest.mark.asyncio async def test_prismaclient_init_wires_default_config( @@ -205,3 +230,217 @@ async def test_disconnect_raises_when_underlying_fails( prisma_client.db.disconnect = AsyncMock(side_effect=RuntimeError("disconnect boom")) with pytest.raises(RuntimeError, match="disconnect boom"): await prisma_client.disconnect() + + +@pytest.mark.asyncio +async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """On a fresh database the migrations Job can still be running when the proxy + boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the + table exists raised inside a fire-and-forget task and the views never appeared.""" + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + "probe_args": probe.await_args.args, + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "probe", "views", "row_count"], + "probe_args": (_PROBE_SQL, '"public"."LiteLLM_SpendLogs"'), + } + + +@pytest.mark.asyncio +async def test_view_setup_probes_the_configured_database_schema( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant") + probe = AsyncMock(return_value=_present()) + _wire_view_setup(prisma_client, probe) + + await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + assert probe.await_args.args == (_PROBE_SQL, '"litellm_tenant"."LiteLLM_SpendLogs"') + + +@pytest.mark.asyncio +async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "kept_polling": probe.await_count > 1, + "views_attempted": prisma_client.check_view_exists.await_count, + "row_count_attempted": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "outcome": "timed_out", + "kept_polling": True, + "views_attempted": 0, + "row_count_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: + """``LiteLLM_SpendLogs`` lands early in the migration set while + ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, + so the first attempt after the table appears can still fail.""" + probe = AsyncMock(return_value=_present()) + call_order = _wire_view_setup(prisma_client, probe) + prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None] + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "views", "probe", "views", "row_count"], + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_client: PrismaClient) -> None: + probe = AsyncMock(side_effect=[RuntimeError("connection reset"), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "views", "row_count"], + } + + +@pytest.mark.asyncio +async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_table": '"public"."LiteLLM_SpendLogs"' in errors[0], + "tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_table": True, + "tells_operator_to_migrate": True, + } + + +@pytest.mark.asyncio +async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + """A database role without CREATE on the schema fails every attempt even though + the table is there, so the timeout must blame that error, not missing migrations.""" + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_the_error": "permission denied for schema public" in errors[0], + "blames_missing_migrations": "did not appear" in errors[0], + "tells_operator_to_restart": "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_the_error": True, + "blames_missing_migrations": False, + "tells_operator_to_restart": True, + } + + +@pytest.mark.asyncio +async def test_run_view_setup_stays_quiet_when_views_are_ready( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + actual = { + "outcome": outcome, + "errors": [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR], + } + assert actual == {"outcome": "ready", "errors": []} + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_cancels_a_task_parked_between_polls(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + prisma_client.start_view_setup_task() + task = prisma_client._view_setup_task + await asyncio.sleep(0) + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "probed_before_parking": probe.await_count, + "task_cancelled": task is not None and task.cancelled(), + "reference_cleared": prisma_client._view_setup_task, + "views_attempted": prisma_client.check_view_exists.await_count, + } + assert actual == { + "probed_before_parking": 1, + "task_cancelled": True, + "reference_cleared": None, + "views_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_is_a_noop_without_a_task(prisma_client: PrismaClient) -> None: + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + assert prisma_client._view_setup_task is None + + +@pytest.mark.asyncio +async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: PrismaClient) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + prisma_client.start_view_setup_task() + first = prisma_client._view_setup_task + prisma_client.start_view_setup_task() + second = prisma_client._view_setup_task + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "first_is_task": isinstance(first, asyncio.Task), + "second_is_first": second is first, + } + assert actual == {"first_is_task": True, "second_is_first": True} From 5f54f87d9887c13b75c42ffca5142ef2b7bcf2f1 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:40:10 +0000 Subject: [PATCH 010/127] feat(fal_ai): add Seedance video generation via fal queue API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/__init__.py | 3 + litellm/llms/fal_ai/videos/transformation.py | 512 ++++++++++++++++++ ...odel_prices_and_context_window_backup.json | 121 +++++ litellm/utils.py | 4 + model_prices_and_context_window.json | 121 +++++ .../test_fal_ai_video_transformation.py | 231 ++++++++ 6 files changed, 992 insertions(+) create mode 100644 litellm/llms/fal_ai/videos/__init__.py create mode 100644 litellm/llms/fal_ai/videos/transformation.py create mode 100644 tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py diff --git a/litellm/llms/fal_ai/videos/__init__.py b/litellm/llms/fal_ai/videos/__init__.py new file mode 100644 index 00000000000..c7e8f76c75b --- /dev/null +++ b/litellm/llms/fal_ai/videos/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + +__all__ = ("FalAIVideoConfig",) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py new file mode 100644 index 00000000000..f8ebf828d68 --- /dev/null +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -0,0 +1,512 @@ +import math +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import FileContent, RequestFiles +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import ( + CharacterObject, + VideoCreateOptionalRequestParams, + VideoObject, +) +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_video_id_with_provider, +) + + +class FalAIVideoError(BaseLLMException): + pass + + +_ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}) +_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"}) +_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( + (480, "480p"), + (720, "720p"), + (1080, "1080p"), +) +_FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value + + +def _queue_request_base_path(model: str) -> str: + segments: Final[tuple[str, ...]] = tuple(model.split("/")) + segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + return "/".join(segments[:segment_count]) + + +def _duration_value(value: object) -> str | None: + if isinstance(value, str) and value == "auto": + return value + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + return str(int(float(value))) + except (TypeError, ValueError): + return None + + +def _resolution_for_height(height: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") + + +def _size_params(size: object) -> Mapping[str, str]: + if not isinstance(size, str): + return MappingProxyType({}) + if size in _ALLOWED_RESOLUTIONS: + return MappingProxyType({"resolution": size}) + if size.count("x") != 1: + return MappingProxyType({}) + width_text, height_text = size.split("x") + if not (width_text.isdigit() and height_text.isdigit()): + return MappingProxyType({}) + width: Final[int] = int(width_text) + height: Final[int] = int(height_text) + if width <= 0 or height <= 0: + return MappingProxyType({}) + reduced_gcd: Final[int] = math.gcd(width, height) + aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" + resolution: Final[str] = _resolution_for_height(height) + if aspect_ratio in _ALLOWED_ASPECT_RATIOS: + return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) + return MappingProxyType({"resolution": resolution}) + + +def _numeric_duration(value: object) -> float | None: + duration: Final[str | None] = _duration_value(value) + if duration is None or duration == "auto": + return None + return float(duration) + + +def _response_data(raw_response: httpx.Response) -> Mapping[str, object]: + return TypeAdapter(Mapping[str, object]).validate_python(raw_response.json()) + + +def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str: + value: Final[object] = response_data.get(key) + return value if isinstance(value, str) else default + + +class FalAIVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list + return [ # mutable-ok: API contract requires a list + "model", + "prompt", + "input_reference", + "seconds", + "size", + "user", + "extra_headers", + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) + input_reference: Final[object] = video_create_optional_params.get("input_reference") + input_reference_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "input_reference" not in video_create_optional_params + else ( + MappingProxyType({"image_url": input_reference}) + if isinstance(input_reference, str) + else self._invalid_input_reference() + ) + ) + duration_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "seconds" not in video_create_optional_params + else self._duration_params(video_create_optional_params["seconds"]) + ) + size_params: Final[Mapping[str, str]] = ( + self._size_params(video_create_optional_params["size"]) + if "size" in video_create_optional_params + else MappingProxyType({}) + ) + user_params: Final[Mapping[str, str]] = ( + MappingProxyType({"end_user_id": user}) + if isinstance(user := video_create_optional_params.get("user"), str) + else MappingProxyType({}) + ) + return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping + MappingProxyType( + { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: dynamic passthrough fields require a mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + ) + ) # mutable-ok: BaseVideoConfig requires a mutable mapping + + @staticmethod + def _invalid_input_reference() -> Mapping[str, str]: + raise ValueError("fal.ai needs a public image URL for input_reference") + + @staticmethod + def _duration_params(seconds: object) -> Mapping[str, str]: + duration: Final[str | None] = _duration_value(seconds) + if duration is None: + raise ValueError("fal.ai seconds must be a numeric value") + return MappingProxyType({"duration": duration}) + + @staticmethod + def _size_params(size: object) -> Mapping[str, str]: + return _size_params(size) + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + final_api_key: Final[str | None] = ( + api_key + or (litellm_params.api_key if litellm_params is not None else None) + or get_secret_str("FAL_AI_API_KEY") + or get_secret_str("FAL_KEY") + ) + if not final_api_key: + raise ValueError("fal.ai API key is required") + return dict( # mutable-ok: BaseVideoConfig requires mutable headers + MappingProxyType( + { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + ) + ) # mutable-ok: BaseVideoConfig requires mutable headers + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + ) -> str: + return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters + str, object + ], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data + MappingProxyType( + { + "prompt": prompt, + **{ # mutable-ok: dynamic request fields require a mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } + ) + ) + return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({}) + request_id: Final[str] = _response_string(response_data, "request_id") + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + duration: Final[float | None] = _numeric_duration(request_params.get("duration")) + resolution: Final[object] = request_params.get("resolution") + seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None + size: Final[str | None] = resolution if isinstance(resolution, str) else None + usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping + MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), + ) + if value is not None + } + ) + ) # mutable-ok: VideoObject requires a mutable usage mapping + video_object: Final[VideoObject] = VideoObject( + id=encode_video_id_with_provider(request_id, provider, model), + object="video", + status="queued", + created_at=int(time.time()), + model=model, + seconds=seconds, + size=size, + ) + video_object.usage = usage + return video_object + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") + status: Final[str] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } + ).get(raw_status, "queued") + error_value: Final[object] = response_data.get("error") + error: Final[str | None] = error_value if isinstance(error_value, str) else None + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + return VideoObject( + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + object="video", + status="failed" if error else status, + created_at=0, + error=( + {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict + ), # mutable-ok: VideoObject requires a dict + ) + + @staticmethod + def _decode_video_id(video_id: str) -> tuple[str, str]: + decoded: Final = decode_video_id_with_provider(video_id) + request_id: Final[str] = decoded.get("video_id", video_id) + model_id: Final[str | None] = decoded.get("model_id") + if not model_id: + raise ValueError("fal.ai video ids must be created through litellm with a model") + return request_id, model_id + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + variant: str | None = None, + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + @staticmethod + def _extract_video_url(response_data: Mapping[str, object]) -> str: + raw_video_data: Final[object] = response_data.get("video") + video_data: Final[Mapping[str, object] | None] = ( + TypeAdapter(Mapping[str, object]).validate_python(raw_video_data) + if isinstance(raw_video_data, Mapping) + else None + ) + if video_data is not None: + video_url: Final[object] = video_data.get("url") + if isinstance(video_url, str) and video_url: + return video_url + error_message: Final[str | None] = next( + (value for key in ("error", "detail") if isinstance(value := response_data.get(key), str)), + None, + ) + if error_message: + raise ValueError(f"fal.ai video result did not include a video URL: {error_message}") + raise ValueError("fal.ai video result did not include a video URL") + + def transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + httpx_client: Final[HTTPHandler] = _get_httpx_client() + video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + async def async_transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI) + video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_create_character_request( + self, + name: str, + video: object, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_create_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_get_character_request( + self, + character_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_get_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + video_file: FileContent | None = None, + extra_body: Mapping[str, object] | None = None, + prefetched_source_data: Mapping[str, object] | None = None, + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_edit_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video extension is not supported for fal.ai") + + def transform_video_extension_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video extension is not supported for fal.ai") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + ) -> BaseLLMException: + return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4dbf0337894..324eb6b2d66 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/litellm/utils.py b/litellm/utils.py index 48d13bc16af..3991cecdac6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9403,6 +9403,10 @@ class ProviderConfigManager: from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig return RunwayMLVideoConfig() + elif LlmProviders.FAL_AI == provider: + from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + + return FalAIVideoConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4dbf0337894..324eb6b2d66 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py new file mode 100644 index 00000000000..d4b058ddcf6 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -0,0 +1,231 @@ +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +import litellm.llms.fal_ai.videos.transformation as fal_video_module +from litellm.cost_calculator import default_video_cost_calculator +from litellm.llms.fal_ai.videos.transformation import ( + FalAIVideoConfig, + FalAIVideoError, + _queue_request_base_path, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.utils import decode_video_id_with_provider +from litellm.utils import ProviderConfigManager + +MODEL = "bytedance/seedance-2.5/text-to-video" + + +class TestFalAIVideoTransformation: + def setup_method(self): + self.config = FalAIVideoConfig() + self.logging_obj = Mock() + + def test_map_openai_params(self): + mapped = self.config.map_openai_params( + { + "seconds": "5", + "size": "1280x720", + "input_reference": "https://example.com/image.png", + "user": "user-123", + "generate_audio": False, + }, + MODEL, + False, + ) + + assert mapped == { + "duration": "5", + "resolution": "720p", + "aspect_ratio": "16:9", + "image_url": "https://example.com/image.png", + "end_user_id": "user-123", + "generate_audio": False, + } + + assert self.config.map_openai_params({"size": "1080x1080"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "1:1", + } + assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + + def test_map_openai_params_rejects_non_url_input_reference(self): + with pytest.raises(ValueError, match="public image URL"): + self.config.map_openai_params({"input_reference": b"image"}, MODEL, False) + + def test_transform_video_create_request(self): + body, files, url = self.config.transform_video_create_request( + model=MODEL, + prompt="A quiet ocean at sunrise", + api_base="https://queue.fal.run", + video_create_optional_request_params={ + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + "model": MODEL, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == f"https://queue.fal.run/{MODEL}" + assert files == [] + assert body == { + "prompt": "A quiet ocean at sunrise", + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + } + assert "model" not in body + + def test_transform_video_create_response_encodes_model_and_usage(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "5", "resolution": "480p"}, + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["custom_llm_provider"] == "fal_ai" + assert decoded["model_id"] == MODEL + assert decoded["video_id"] == "abc" + assert video.status == "queued" + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "480p"} + + auto_video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "auto"}, + ) + assert auto_video.usage == {"video_resolution": "720p"} + assert auto_video.seconds is None + assert auto_video.size is None + + def test_status_request_uses_queue_base_path(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={}, + ) + + url, params = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert params == {} + assert _queue_request_base_path("workflows/owner/app/x") == "workflows/owner/app" + assert _queue_request_base_path("comfy/owner/app/x") == "comfy/owner/app" + + def test_status_request_rejects_unencoded_video_id(self): + with pytest.raises(ValueError, match="must be created through litellm"): + self.config.transform_video_status_retrieve_request( + video_id="abc", + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + @pytest.mark.parametrize( + ("response_data", "expected_status"), + [ + ({"request_id": "abc", "status": "IN_QUEUE"}, "queued"), + ({"request_id": "abc", "status": "IN_PROGRESS"}, "in_progress"), + ({"request_id": "abc", "status": "COMPLETED"}, "completed"), + ], + ) + def test_status_response_mapping(self, response_data, expected_status): + response = Mock(spec=httpx.Response) + response.json.return_value = response_data + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == expected_status + assert video.created_at == 0 + + def test_status_response_error(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "COMPLETED", + "error": "generation failed", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "failed" + assert video.error == {"code": "fal_error", "message": "generation failed"} + + def test_content_response_downloads_video_url(self, monkeypatch): + content_response = httpx.Response( + 200, + content=b"video-bytes", + request=httpx.Request("GET", "https://cdn.example.com/video.mp4"), + ) + + class FakeHTTPClient: + def get(self, url): + assert url == "https://cdn.example.com/video.mp4" + return content_response + + monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: FakeHTTPClient()) + response = Mock(spec=httpx.Response) + response.json.return_value = {"video": {"url": "https://cdn.example.com/video.mp4"}} + + assert self.config.transform_video_content_response(response, self.logging_obj) == b"video-bytes" + + def test_content_response_rejects_missing_video(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"error": "generation failed"} + + with pytest.raises(ValueError, match="generation failed"): + self.config.transform_video_content_response(response, self.logging_obj) + + def test_provider_config_and_error_class(self): + provider_config = ProviderConfigManager.get_provider_video_config( + model=MODEL, + provider=LlmProviders.FAL_AI, + ) + assert isinstance(provider_config, FalAIVideoConfig) + assert isinstance(self.config.get_error_class("bad key", 401, {}), FalAIVideoError) + + def test_video_cost_uses_tiered_rows(self): + rows = { + model: row + for model, row in litellm.model_cost.items() + if row.get("litellm_provider") == "fal_ai" and row.get("mode") == "video_generation" + } + assert rows + for model, row in rows.items(): + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="480p") == ( + 5 * row["output_cost_per_second_480p"] + ) + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="720p") == ( + 5 * row["output_cost_per_second"] + ) From 0f4ce95492cde57b1f2eb29a1ea119dda0b13e24 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:45:53 +0000 Subject: [PATCH 011/127] refactor(fal_ai): simplify video config mappings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 178 +++++++++---------- 1 file changed, 80 insertions(+), 98 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f8ebf828d68..f0142ddc3de 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -2,7 +2,7 @@ import math import time from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import httpx from httpx._types import FileContent, RequestFiles @@ -42,12 +42,25 @@ _RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( (720, "720p"), (1080, "1080p"), ) +_QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy")) +_STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } +) _FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value +_SupportedParams: TypeAlias = list[str] +_VideoParams: TypeAlias = dict[str, object] +_VideoHeaders: TypeAlias = dict[str, str] +_VideoStringParams: TypeAlias = dict[str, str] +_VideoFiles: TypeAlias = list[object] def _queue_request_base_path(model: str) -> str: segments: Final[tuple[str, ...]] = tuple(model.split("/")) - segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + segment_count: Final[int] = 3 if segments and segments[0] in _QUEUE_NAMESPACES else 2 return "/".join(segments[:segment_count]) @@ -105,8 +118,8 @@ def _response_string(response_data: Mapping[str, object], key: str, default: str class FalAIVideoConfig(BaseVideoConfig): - def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list - return [ # mutable-ok: API contract requires a list + def get_supported_openai_params(self, model: str) -> _SupportedParams: + supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list "model", "prompt", "input_reference", @@ -115,23 +128,22 @@ class FalAIVideoConfig(BaseVideoConfig): "user", "extra_headers", ] + return supported_params def map_openai_params( self, video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + ) -> _VideoParams: supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) input_reference: Final[object] = video_create_optional_params.get("input_reference") + if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): + raise ValueError("fal.ai needs a public image URL for input_reference") input_reference_params: Final[Mapping[str, str]] = ( MappingProxyType({}) - if "input_reference" not in video_create_optional_params - else ( - MappingProxyType({"image_url": input_reference}) - if isinstance(input_reference, str) - else self._invalid_input_reference() - ) + if not isinstance(input_reference, str) + else MappingProxyType({"image_url": input_reference}) ) duration_params: Final[Mapping[str, str]] = ( MappingProxyType({}) @@ -139,7 +151,7 @@ class FalAIVideoConfig(BaseVideoConfig): else self._duration_params(video_create_optional_params["seconds"]) ) size_params: Final[Mapping[str, str]] = ( - self._size_params(video_create_optional_params["size"]) + _size_params(video_create_optional_params["size"]) if "size" in video_create_optional_params else MappingProxyType({}) ) @@ -148,23 +160,16 @@ class FalAIVideoConfig(BaseVideoConfig): if isinstance(user := video_create_optional_params.get("user"), str) else MappingProxyType({}) ) - return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping - MappingProxyType( - { - **input_reference_params, - **duration_params, - **size_params, - **user_params, - **{ # mutable-ok: dynamic passthrough fields require a mapping - key: value for key, value in video_create_optional_params.items() if key not in supported_params - }, - } - ) - ) # mutable-ok: BaseVideoConfig requires a mutable mapping - - @staticmethod - def _invalid_input_reference() -> Mapping[str, str]: - raise ValueError("fal.ai needs a public image URL for input_reference") + mapped_params: Final[_VideoParams] = { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: BaseVideoConfig requires a mutable parameter mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + return mapped_params @staticmethod def _duration_params(seconds: object) -> Mapping[str, str]: @@ -173,17 +178,13 @@ class FalAIVideoConfig(BaseVideoConfig): raise ValueError("fal.ai seconds must be a numeric value") return MappingProxyType({"duration": duration}) - @staticmethod - def _size_params(size: object) -> Mapping[str, str]: - return _size_params(size) - def validate_environment( self, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, model: str, api_key: str | None = None, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + ) -> _VideoHeaders: final_api_key: Final[str | None] = ( api_key or (litellm_params.api_key if litellm_params is not None else None) @@ -192,21 +193,18 @@ class FalAIVideoConfig(BaseVideoConfig): ) if not final_api_key: raise ValueError("fal.ai API key is required") - return dict( # mutable-ok: BaseVideoConfig requires mutable headers - MappingProxyType( - { - **headers, - "Authorization": f"Key {final_api_key}", - "Content-Type": "application/json", - } - ) - ) # mutable-ok: BaseVideoConfig requires mutable headers + validated_headers: Final[_VideoHeaders] = { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + return validated_headers def get_complete_url( self, model: str, api_base: str | None, - litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: _VideoParams, ) -> str: return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") @@ -215,22 +213,16 @@ class FalAIVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters - str, object - ], # mutable-ok: BaseVideoConfig requires mutable parameters + video_create_optional_request_params: _VideoParams, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings - request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data - MappingProxyType( - { - "prompt": prompt, - **{ # mutable-ok: dynamic request fields require a mapping - key: value for key, value in video_create_optional_request_params.items() if key != "model" - }, - } - ) - ) + headers: _VideoHeaders, + ) -> tuple[_VideoParams, RequestFiles, str]: + request_data: Final[_VideoParams] = { + "prompt": prompt, + **{ # mutable-ok: HTTP JSON payload requires a mutable mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list def transform_video_create_response( @@ -249,18 +241,14 @@ class FalAIVideoConfig(BaseVideoConfig): resolution: Final[object] = request_params.get("resolution") seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None size: Final[str | None] = resolution if isinstance(resolution, str) else None - usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping - MappingProxyType( - { - key: value - for key, value in ( - ("duration_seconds", duration), - ("video_resolution", resolution if isinstance(resolution, str) else "720p"), - ) - if value is not None - } + usage: Final[_VideoParams] = { # mutable-ok: VideoObject requires a mutable usage mapping + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), ) - ) # mutable-ok: VideoObject requires a mutable usage mapping + if value is not None + } video_object: Final[VideoObject] = VideoObject( id=encode_video_id_with_provider(request_id, provider, model), object="video", @@ -278,8 +266,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -295,13 +283,7 @@ class FalAIVideoConfig(BaseVideoConfig): ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") - status: Final[str] = MappingProxyType( - { - "IN_QUEUE": "queued", - "IN_PROGRESS": "in_progress", - "COMPLETED": "completed", - } - ).get(raw_status, "queued") + status: Final[str] = _STATUS_MAP.get(raw_status, "queued") error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER @@ -312,7 +294,7 @@ class FalAIVideoConfig(BaseVideoConfig): created_at=0, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict - ), # mutable-ok: VideoObject requires a dict + ), ) @staticmethod @@ -329,9 +311,9 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, variant: str | None = None, - ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoStringParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -383,9 +365,9 @@ class FalAIVideoConfig(BaseVideoConfig): prompt: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video remix is not supported for fal.ai") def transform_video_remix_response( @@ -400,12 +382,12 @@ class FalAIVideoConfig(BaseVideoConfig): self, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, after: str | None = None, limit: int | None = None, order: str | None = None, extra_query: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_list_response( @@ -413,7 +395,7 @@ class FalAIVideoConfig(BaseVideoConfig): raw_response: httpx.Response, logging_obj: object, custom_llm_provider: str | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> _VideoStringParams: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_delete_request( @@ -421,8 +403,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video delete is not supported for fal.ai") def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: @@ -434,8 +416,8 @@ class FalAIVideoConfig(BaseVideoConfig): video: object, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + headers: _VideoHeaders, + ) -> tuple[str, _VideoFiles]: raise NotImplementedError("video character creation is not supported for fal.ai") def transform_video_create_character_response( @@ -450,8 +432,8 @@ class FalAIVideoConfig(BaseVideoConfig): character_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video character retrieval is not supported for fal.ai") def transform_video_get_character_response( @@ -467,7 +449,7 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, video_file: FileContent | None = None, extra_body: Mapping[str, object] | None = None, prefetched_source_data: Mapping[str, object] | None = None, @@ -490,9 +472,9 @@ class FalAIVideoConfig(BaseVideoConfig): seconds: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video extension is not supported for fal.ai") def transform_video_extension_response( @@ -507,6 +489,6 @@ class FalAIVideoConfig(BaseVideoConfig): self, error_message: str, status_code: int, - headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + headers: _VideoHeaders | httpx.Headers, ) -> BaseLLMException: return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) From 141548dcf3ec7588e86c2e304fca33be6c4cdc7e Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:54:11 +0000 Subject: [PATCH 012/127] fix(fal_ai): keep status ids pollable and size resolution by the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 17 +++++++-- .../test_fal_ai_video_transformation.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f0142ddc3de..4ca1f918c90 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -75,8 +75,16 @@ def _duration_value(value: object) -> str | None: return None -def _resolution_for_height(height: int) -> str: - return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") +def _resolution_for_short_side(short_side: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") + + +def _model_path_from_queue_url(url: object) -> str | None: + if not isinstance(url, str) or not url: + return None + path: Final[str] = httpx.URL(url).path.strip("/") + model_path, separator, _ = path.partition("/requests/") + return model_path if separator and model_path else None def _size_params(size: object) -> Mapping[str, str]: @@ -95,7 +103,7 @@ def _size_params(size: object) -> Mapping[str, str]: return MappingProxyType({}) reduced_gcd: Final[int] = math.gcd(width, height) aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" - resolution: Final[str] = _resolution_for_height(height) + resolution: Final[str] = _resolution_for_short_side(min(width, height)) if aspect_ratio in _ALLOWED_ASPECT_RATIOS: return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) return MappingProxyType({"resolution": resolution}) @@ -287,8 +295,9 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), object="video", status="failed" if error else status, created_at=0, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index d4b058ddcf6..f367fa331d5 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -51,6 +51,14 @@ class TestFalAIVideoTransformation: "aspect_ratio": "1:1", } assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + assert self.config.map_openai_params({"size": "720x1280"}, MODEL, False) == { + "resolution": "720p", + "aspect_ratio": "9:16", + } + assert self.config.map_openai_params({"size": "1080x1920"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "9:16", + } def test_map_openai_params_rejects_non_url_input_reference(self): with pytest.raises(ValueError, match="public image URL"): @@ -165,6 +173,35 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + def test_status_response_id_stays_pollable(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "IN_PROGRESS", + "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + status_url, _ = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + content_url, _ = self.config.transform_video_content_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + def test_status_response_error(self): response = Mock(spec=httpx.Response) response.json.return_value = { From c359ef763eae44523cace811bcb1999992250bde Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:01:09 +0000 Subject: [PATCH 013/127] fix(fal_ai): keep model in polled video ids and pick resolution from the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 28 ++++++-- .../test_fal_ai_video_transformation.py | 64 +++++++++++-------- 2 files changed, 58 insertions(+), 34 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 4ca1f918c90..766316a5b18 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -79,12 +79,22 @@ def _resolution_for_short_side(short_side: int) -> str: return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") -def _model_path_from_queue_url(url: object) -> str | None: - if not isinstance(url, str) or not url: +def _model_path_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: return None - path: Final[str] = httpx.URL(url).path.strip("/") - model_path, separator, _ = path.partition("/requests/") - return model_path if separator and model_path else None + model_segments: Final[tuple[str, ...]] = segments[: segments.index("requests")] + segment_count: Final[int] = 3 if len(model_segments) >= 3 and model_segments[-3] in _QUEUE_NAMESPACES else 2 + return "/".join(model_segments[-segment_count:]) if len(model_segments) >= segment_count else None + + +def _request_id_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: + return None + request_index: Final[int] = segments.index("requests") + request_id_index: Final[int] = request_index + 1 + return segments[request_id_index] if len(segments) > request_id_index else None def _size_params(size: object) -> Mapping[str, str]: @@ -295,12 +305,16 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER - model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) + model_path: Final[str | None] = _model_path_from_request_url(raw_response) + request_id: Final[str] = _response_string(response_data, "request_id") or ( + _request_id_from_request_url(raw_response) or "" + ) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), + id=encode_video_id_with_provider(request_id, provider, model_path), object="video", status="failed" if error else status, created_at=0, + model=model_path, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict ), diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index f367fa331d5..8e0c68e30bb 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -161,8 +161,8 @@ class TestFalAIVideoTransformation: ], ) def test_status_response_mapping(self, response_data, expected_status): - response = Mock(spec=httpx.Response) - response.json.return_value = response_data + status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + response = httpx.Response(200, json=response_data, request=httpx.Request("GET", status_url)) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -172,43 +172,32 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "bytedance/seedance-2.5" + assert decoded["video_id"] == "abc" - def test_status_response_id_stays_pollable(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { - "request_id": "abc", - "status": "IN_PROGRESS", - "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", - } - - video = self.config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=self.logging_obj, - custom_llm_provider="fal_ai", - ) - - status_url, _ = self.config.transform_video_status_retrieve_request( + poll_url, _ = self.config.transform_video_status_retrieve_request( video_id=video.id, api_base="https://queue.fal.run", litellm_params=GenericLiteLLMParams(), headers={}, ) - content_url, _ = self.config.transform_video_content_request( - video_id=video.id, - api_base="https://queue.fal.run", - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" - assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + assert poll_url == status_url def test_status_response_error(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { + response_data = { "request_id": "abc", "status": "COMPLETED", "error": "generation failed", } + response = httpx.Response( + 200, + json=response_data, + request=httpx.Request( + "GET", + "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status", + ), + ) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -219,6 +208,27 @@ class TestFalAIVideoTransformation: assert video.status == "failed" assert video.error == {"code": "fal_error", "message": "generation failed"} + def test_status_response_uses_namespaced_request_url(self): + response = httpx.Response( + 200, + json={"status": "IN_PROGRESS"}, + request=httpx.Request( + "GET", + "https://example.com/proxy/workflows/owner/app/requests/xyz/status", + ), + ) + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "workflows/owner/app" + assert decoded["video_id"] == "xyz" + assert video.model == "workflows/owner/app" + def test_content_response_downloads_video_url(self, monkeypatch): content_response = httpx.Response( 200, From aa5f0858f75b3e074264e0266f87b70a6cb70391 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:13:41 +0000 Subject: [PATCH 014/127] test(pricing): allow video endpoint and rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b40c10de428..d694c08510a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -940,6 +940,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/audio/transcriptions", "/v1/audio/speech", "/v1/ocr", + "/v1/videos", "/vertex_ai/live", "/v1/listen", "/v1beta/interactions", @@ -1069,6 +1070,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): # Add any model IDs that should be exempt from the cost validation # Example: "expensive-model-id", "runwayml/seedance2", # 4K output is 150 credits/second = $1.50/second + "fal_ai/bytedance/seedance-2.0/text-to-video", + "fal_ai/bytedance/seedance-2.0/image-to-video", + "fal_ai/bytedance/seedance-2.0/reference-to-video", ] is_valid, violations = validate_model_cost_values(actual_json, exceptions) From e0b455e94e83dfedad0364aedc6b4cdfcb59acb0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:27:44 +0000 Subject: [PATCH 015/127] fix(fal_ai): read only the documented FAL_AI_API_KEY env var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 5 ++--- .../test_fal_ai_video_transformation.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 766316a5b18..98528c82f6b 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -207,10 +207,9 @@ class FalAIVideoConfig(BaseVideoConfig): api_key or (litellm_params.api_key if litellm_params is not None else None) or get_secret_str("FAL_AI_API_KEY") - or get_secret_str("FAL_KEY") ) if not final_api_key: - raise ValueError("fal.ai API key is required") + raise ValueError("FAL_AI_API_KEY is not set") validated_headers: Final[_VideoHeaders] = { **headers, "Authorization": f"Key {final_api_key}", @@ -224,7 +223,7 @@ class FalAIVideoConfig(BaseVideoConfig): api_base: str | None, litellm_params: _VideoParams, ) -> str: - return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + return (api_base or "https://queue.fal.run").rstrip("/") def transform_video_create_request( self, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index 8e0c68e30bb..5e2e4532265 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -91,6 +91,26 @@ class TestFalAIVideoTransformation: } assert "model" not in body + def test_get_complete_url_respects_api_base_override(self): + url = self.config.get_complete_url( + model=MODEL, + api_base="https://proxy.internal/", + litellm_params={}, + ) + + assert url == "https://proxy.internal" + + def test_validate_environment_requires_fal_ai_api_key(self, monkeypatch): + monkeypatch.setattr(fal_video_module, "get_secret_str", lambda _: None) + + with pytest.raises(ValueError, match="FAL_AI_API_KEY is not set"): + self.config.validate_environment( + headers={}, + model=MODEL, + api_key=None, + litellm_params=GenericLiteLLMParams(), + ) + def test_transform_video_create_response_encodes_model_and_usage(self): response = Mock(spec=httpx.Response) response.json.return_value = {"request_id": "abc"} From 3b0d32ec6f14979687d4bb76f51db5f9427a131a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 19 Sep 2026 18:11:50 +0000 Subject: [PATCH 016/127] fix(proxy): reject throttled exhausted budgets in JEV previews Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 8 +++ .../test_auto_router_endpoints.py | 49 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f79425d2e97..c5a10cf5c80 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -345,6 +345,14 @@ async def _authorize_models_this_test_can_call( code=status.HTTP_400_BAD_REQUEST, ) from e + if config.classifier_type == "jev" and user_api_key_dict.budget_throttle_pct is not None: + raise ProxyException( + message="Budget has been exceeded! JEV Test Routing requires available budget.", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=status.HTTP_400_BAD_REQUEST, + ) + @router.post( "/auto_router/validate_complexity_router_config", diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index a5c93c41a84..eb9076a9d2d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -12,6 +12,7 @@ import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +import litellm from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, @@ -488,6 +489,54 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( client.evaluate.assert_awaited_once() +@pytest.mark.parametrize( + "max_budget, spend, denied", + ((0.0, 0.0, True), (1.0, 2.0, True), (1.0, 0.5, False), (None, 2.0, False)), +) +@pytest.mark.asyncio +async def test_jev_test_routing_hard_blocks_exhausted_throttle_enabled_keys( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-throttle-test", + user_id="admin", + models=["cheap-model", "typesafe/jev-test"], + max_budget=max_budget, + spend=spend, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routing_decision["cause"] == "jev_classifier" + client.evaluate.assert_awaited_once() + + @pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) @pytest.mark.asyncio async def test_a_heuristic_config_does_not_need_a_budget( From 85a6a8e2063ef0932dc40b19c97eb0120a0c4235 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 18:22:58 +0000 Subject: [PATCH 017/127] test(e2e): cover fal Seedance video create, poll and download Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 + tests/e2e/coverage_registry/schema.py | 2 + .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 + tests/e2e/llm_translation/endpoints_client.py | 36 +++++++++- .../test_video_generation_e2e.py | 69 +++++++++++++++++++ 5 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 50f9b9808b2..6970567b6f0 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,6 +80,7 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} +- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fa6dad90126..3ae17432863 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,6 +44,7 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", + "videos", ] LlmRoute = Literal[ @@ -53,6 +54,7 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", + "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 44d6e79122e..178af054f2b 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,6 +48,7 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | +| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -61,6 +62,7 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | +| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 4d2c73e7078..165a83e76c0 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, Result, StreamingResponse +from e2e_http import BinaryStream, NoBody, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,6 +26,8 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", + "VideoObject", + "VideoRequest", ] @@ -127,6 +129,13 @@ class ImageRequest(BaseModel): size: str = "1024x1024" +class VideoRequest(BaseModel): + model: str + prompt: str + seconds: str = "4" + size: str = "1280x720" + + class ImageEditForm(BaseModel): model: str prompt: str @@ -267,6 +276,12 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] +class VideoObject(BaseModel): + id: str + status: str + model: str | None = None + + class TranscriptionResult(BaseModel): text: str = "" @@ -440,6 +455,25 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) + def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: + return self._send( + "/v1/videos", key, VideoRequest(model=model, prompt=prompt) + ) + + def video_status(self, key: str, video_id: str) -> Result[VideoObject]: + return self.proxy.transport.get( + f"/v1/videos/{video_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=VideoObject, + ) + + def video_content(self, key: str, video_id: str) -> StreamingResponse: + return self.proxy.transport.download( + f"/v1/videos/{video_id}/content", + headers=self.proxy.transport.bearer(key), + ) + def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py new file mode 100644 index 00000000000..b65529aa260 --- /dev/null +++ b/tests/e2e/llm_translation/test_video_generation_e2e.py @@ -0,0 +1,69 @@ +"""Live e2e: POST /v1/videos creates a video and serves its content. + +Registers a fal.ai Seedance deployment at runtime, polls the queued video until it +completes, and asserts the generated content is returned as binary data. +""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call, unwrap +from endpoints_client import EndpointsClient, VideoObject +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +_POLL_INTERVAL_SECONDS: Final[float] = 5.0 +_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 + + +def _wait_for_completion( + endpoints_client: EndpointsClient, key: str, created: VideoObject +) -> VideoObject: + deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS + while time.monotonic() < deadline: + status = unwrap(endpoints_client.video_status(key, created.id)) + assert status.id == created.id + if status.status == "completed": + return status + if status.status == "failed": + pytest.fail(f"fal.ai video generation failed: {status}") + time.sleep(_POLL_INTERVAL_SECONDS) + pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") + + +class TestVideoGeneration: + @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") + def test_fal_seedance_video_completes_and_downloads( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-fal-video-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="fal_ai/bytedance/seedance-2.5/text-to-video", + api_key="os.environ/FAL_AI_API_KEY", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.videos( + key, model, "a red fox running through snow at dawn" + ) + require_successful_call(result) + created = VideoObject.model_validate_json(result.body) + assert created.id + assert created.model + + _wait_for_completion(endpoints_client, key, created) + + content = endpoints_client.video_content(key, created.id) + require_successful_call(content) + assert len(content.body) > 0 + assert not (content.content_type or "").startswith("application/json") From 9a63e06c63e60b3b30126aa7461115a2ccc041fe Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:28:51 +0000 Subject: [PATCH 018/127] test(integration): cover fal Seedance video queue wire contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 - tests/e2e/coverage_registry/schema.py | 2 - .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 - tests/e2e/llm_translation/endpoints_client.py | 36 +--------- .../test_video_generation_e2e.py | 69 ------------------ tests/integration/contracts.json | 3 + .../providers/test_fal_ai_video_wire.py | 72 +++++++++++++++++++ 7 files changed, 76 insertions(+), 109 deletions(-) delete mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py create mode 100644 tests/integration/providers/test_fal_ai_video_wire.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 6970567b6f0..50f9b9808b2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,7 +80,6 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} -- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 3ae17432863..fa6dad90126 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,7 +44,6 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", - "videos", ] LlmRoute = Literal[ @@ -54,7 +53,6 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", - "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 178af054f2b..44d6e79122e 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,7 +48,6 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | -| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -62,7 +61,6 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | -| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 165a83e76c0..4d2c73e7078 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, NoBody, Result, StreamingResponse +from e2e_http import BinaryStream, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,8 +26,6 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", - "VideoObject", - "VideoRequest", ] @@ -129,13 +127,6 @@ class ImageRequest(BaseModel): size: str = "1024x1024" -class VideoRequest(BaseModel): - model: str - prompt: str - seconds: str = "4" - size: str = "1280x720" - - class ImageEditForm(BaseModel): model: str prompt: str @@ -276,12 +267,6 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] -class VideoObject(BaseModel): - id: str - status: str - model: str | None = None - - class TranscriptionResult(BaseModel): text: str = "" @@ -455,25 +440,6 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) - def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/videos", key, VideoRequest(model=model, prompt=prompt) - ) - - def video_status(self, key: str, video_id: str) -> Result[VideoObject]: - return self.proxy.transport.get( - f"/v1/videos/{video_id}", - headers=self.proxy.transport.bearer(key), - params=NoBody(), - response_type=VideoObject, - ) - - def video_content(self, key: str, video_id: str) -> StreamingResponse: - return self.proxy.transport.download( - f"/v1/videos/{video_id}/content", - headers=self.proxy.transport.bearer(key), - ) - def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py deleted file mode 100644 index b65529aa260..00000000000 --- a/tests/e2e/llm_translation/test_video_generation_e2e.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Live e2e: POST /v1/videos creates a video and serves its content. - -Registers a fal.ai Seedance deployment at runtime, polls the queued video until it -completes, and asserts the generated content is returned as binary data. -""" - -from __future__ import annotations - -import time -from typing import Final - -import pytest -from e2e_config import unique_marker -from e2e_http import require_successful_call, unwrap -from endpoints_client import EndpointsClient, VideoObject -from lifecycle import ResourceManager -from models import LiteLLMParamsBody - -pytestmark = pytest.mark.e2e - -_POLL_INTERVAL_SECONDS: Final[float] = 5.0 -_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 - - -def _wait_for_completion( - endpoints_client: EndpointsClient, key: str, created: VideoObject -) -> VideoObject: - deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS - while time.monotonic() < deadline: - status = unwrap(endpoints_client.video_status(key, created.id)) - assert status.id == created.id - if status.status == "completed": - return status - if status.status == "failed": - pytest.fail(f"fal.ai video generation failed: {status}") - time.sleep(_POLL_INTERVAL_SECONDS) - pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") - - -class TestVideoGeneration: - @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") - def test_fal_seedance_video_completes_and_downloads( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-fal-video-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="fal_ai/bytedance/seedance-2.5/text-to-video", - api_key="os.environ/FAL_AI_API_KEY", - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.videos( - key, model, "a red fox running through snow at dawn" - ) - require_successful_call(result) - created = VideoObject.model_validate_json(result.body) - assert created.id - assert created.model - - _wait_for_completion(endpoints_client, key, created) - - content = endpoints_client.video_content(key, created.id) - require_successful_call(content) - assert len(content.body) > 0 - assert not (content.content_type or "").startswith("application/json") diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 6958ade50f7..01f7af6e8fe 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -160,6 +160,9 @@ "other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_video_create_status_and_content_follow_queue_wire_contract": [ + "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [ "mcp.call_tool.saved_headers.reach_actual_transport" ], diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py new file mode 100644 index 00000000000..c1a2655f0aa --- /dev/null +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -0,0 +1,72 @@ +import json +import sys +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "bytedance/seedance-2.5/text-to-video" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4 + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download") +def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: Gateway) -> None: + request_id: Final = "fal-req-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.target == f"/files/{request_id}.mp4": + assert request.method == "GET" + return Reply(body=_MP4, content_type="video/mp4") + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_MODEL}" + assert json.loads(request.body) == { + "prompt": "a cat playing volleyball on a beach", + "duration": "4", + "resolution": "720p", + "aspect_ratio": "16:9", + } + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/bytedance/seedance-2.5/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/bytedance/seedance-2.5/requests/{request_id}" + return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": "4", + "size": "1280x720", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + assert isinstance(video_id, str) and video_id + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "completed" + status_id_matches_created_id: Final = status["id"] == video_id + sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.headers["content-type"].startswith("video/mp4") + assert content.content == _MP4 + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_MODEL}"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}/status"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), + ("GET", f"/files/{request_id}.mp4"), + ] From 21a2ed62448ebda3ab9de1245b2550aa0bf164e0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:29:25 +0000 Subject: [PATCH 019/127] test(integration): drop id diagnostic from fal video wire test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/providers/test_fal_ai_video_wire.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index c1a2655f0aa..8c72810ffb6 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -1,5 +1,4 @@ import json -import sys import uuid from typing import Final @@ -58,8 +57,6 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: assert isinstance(video_id, str) and video_id status: Final = gateway.get(f"/v1/videos/{video_id}") assert status["status"] == "completed" - status_id_matches_created_id: Final = status["id"] == video_id - sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") assert content.status_code == 200, content.text assert content.headers["content-type"].startswith("video/mp4") From 1b8f704035a358d31962ca24d8ceda77d3dde935 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:17:18 -0700 Subject: [PATCH 020/127] fix(proxy): await the cancelled view setup task quietly and assert it starts at boot Use contextlib.suppress for the cancelled task in stop_view_setup_task, make the legacy prisma setup test inject a plain mock for the synchronous start_view_setup_task and assert it is called, and drop the docstrings the branch added to tests --- litellm/proxy/utils.py | 4 +--- .../spend_tracking/spend_e2e_client.py | 2 -- tests/proxy_unit_tests/test_proxy_server.py | 17 +++++------------ tests/test_litellm/proxy/test_proxy_server.py | 7 ------- .../test_prisma_client_lifecycle.py | 8 -------- 5 files changed, 6 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7f5609b2e39..37c5c8acec8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6350,10 +6350,8 @@ class PrismaClient: if self._view_setup_task is None: return self._view_setup_task.cancel() - try: + with contextlib.suppress(asyncio.CancelledError): await self._view_setup_task - except asyncio.CancelledError: - pass self._view_setup_task = None async def _run_view_setup( diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 233aa81c2af..b7f59fe5f89 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -360,8 +360,6 @@ class SpendClient: return self.proxy.transport.probe(path, params=params) def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: - """Re-probe a route that depends on startup work the proxy finishes after it - starts serving, such as the spend views it creates once migrations land.""" outcome: Final = await_converged( lambda: self.probe(path, params=params), converged=lambda result: result.healthy, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 47792b90b08..ed0380058a5 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2382,7 +2382,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py @pytest.mark.asyncio async def test_proxy_server_prisma_setup(): - from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_state + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging from litellm.caching import DualCache @@ -2393,35 +2393,28 @@ async def test_proxy_server_prisma_setup(): ) as mock_prisma_client: mock_client = mock_prisma_client.return_value # This is the mocked instance mock_client.connect = AsyncMock() # Mock the connect method - mock_client.check_view_exists = AsyncMock() # Mock the check_view_exists method + mock_client.start_view_setup_task = MagicMock() mock_client.health_check = AsyncMock() # Mock the health_check method - mock_client._set_spend_logs_row_count_in_proxy_state = ( - AsyncMock() - ) # Mock the _set_spend_logs_row_count_in_proxy_state method mock_client.start_db_health_watchdog_task = AsyncMock() # Mock the db attribute with start_token_refresh_task for RDS IAM token refresh mock_db = MagicMock() mock_db.start_token_refresh_task = AsyncMock() mock_client.db = mock_db - await ProxyStartupEvent._setup_prisma_client( + prisma_client = await ProxyStartupEvent._setup_prisma_client( database_url=os.getenv("DATABASE_URL"), proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) - # Verify our mocked methods were called + assert prisma_client is mock_client mock_client.connect.assert_called_once() - mock_client.check_view_exists.assert_called_once() + mock_client.start_view_setup_task.assert_called_once() # Note: This is REALLY IMPORTANT to check that the health check is called # This is how we ensure the DB is ready before proceeding mock_client.health_check.assert_called_once() - # check that the spend logs row count is set in proxy state - mock_client._set_spend_logs_row_count_in_proxy_state.assert_called_once() - assert proxy_state.get_proxy_state_variable("spend_logs_row_count") is not None - @pytest.mark.asyncio async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7ab142c5fd0..67f386ee457 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1636,9 +1636,6 @@ class _ShutdownAwarePrisma(MockPrisma): @pytest.mark.asyncio async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): - """The view setup task keeps polling for the spend-log table while migrations - run, so a shutdown inside that window has to cancel it rather than leave it - to die with the event loop.""" import yaml from fastapi import FastAPI @@ -13415,10 +13412,6 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch @pytest.mark.asyncio async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): - """View creation used to be two fire-and-forget ``asyncio.create_task`` calls - that raised and vanished when the migrations Job had not created - ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's - held task, which waits for the table, and must not call the two coroutines directly.""" monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") mock_client = _mock_startup_prisma_client() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index bb34486771c..c31713d5802 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -236,9 +236,6 @@ async def test_disconnect_raises_when_underlying_fails( async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: - """On a fresh database the migrations Job can still be running when the proxy - boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the - table exists raised inside a fire-and-forget task and the views never appeared.""" monkeypatch.delenv("DATABASE_SCHEMA", raising=False) probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) call_order = _wire_view_setup(prisma_client, probe) @@ -293,9 +290,6 @@ async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: P @pytest.mark.asyncio async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: - """``LiteLLM_SpendLogs`` lands early in the migration set while - ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, - so the first attempt after the table appears can still fail.""" probe = AsyncMock(return_value=_present()) call_order = _wire_view_setup(prisma_client, probe) prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None] @@ -358,8 +352,6 @@ async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table( prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture ) -> None: - """A database role without CREATE on the schema fails every attempt even though - the table is there, so the timeout must blame that error, not missing migrations.""" _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") From a114af2a26ed20f2a0dbb3fc452bc30f7d078d46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:54:20 -0700 Subject: [PATCH 021/127] fix(proxy): resolve the view setup gate through the search_path and set the row count before the views --- litellm/proxy/utils.py | 15 +++----- .../test_prisma_client_lifecycle.py | 38 ++++++++++++------- 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 37c5c8acec8..b45bcc3bdc3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -271,7 +271,7 @@ class _RelTuplesRow(TypedDict): _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 -_VIEW_SETUP_GATE_TABLE: Final = "LiteLLM_SpendLogs" +_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] @@ -6372,11 +6372,11 @@ class PrismaClient: try: if not await self._view_setup_gate_table_present(): verbose_proxy_logger.debug( - "Waiting for table %s before creating the spend views", self._view_setup_gate_table() + "Waiting for table %s before creating the spend views", _VIEW_SETUP_GATE_TABLE ) return "table_missing" - await self.check_view_exists() await self._set_spend_logs_row_count_in_proxy_state() + await self.check_view_exists() return "ready" except Exception as e: verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e) @@ -6396,21 +6396,16 @@ class PrismaClient: verbose_proxy_logger.error( "Gave up creating the spend views: table %s did not appear within %ss. " "Run the database migrations against this database and restart the proxy.", - self._view_setup_gate_table(), + _VIEW_SETUP_GATE_TABLE, deadline_seconds, ) async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", self._view_setup_gate_table()) + await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] - @staticmethod - def _view_setup_gate_table() -> str: - pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") - return f'"{pg_schema}"."{_VIEW_SETUP_GATE_TABLE}"' - async def _db_health_watchdog_loop(self) -> None: while True: try: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index c31713d5802..9aa57c7a19b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -233,10 +233,7 @@ async def test_disconnect_raises_when_underlying_fails( @pytest.mark.asyncio -async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( - prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("DATABASE_SCHEMA", raising=False) +async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views(prisma_client: PrismaClient) -> None: probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) call_order = _wire_view_setup(prisma_client, probe) @@ -249,13 +246,13 @@ async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( } assert actual == { "outcome": "ready", - "calls": ["probe", "probe", "probe", "views", "row_count"], - "probe_args": (_PROBE_SQL, '"public"."LiteLLM_SpendLogs"'), + "calls": ["probe", "probe", "probe", "row_count", "views"], + "probe_args": (_PROBE_SQL, '"LiteLLM_SpendLogs"'), } @pytest.mark.asyncio -async def test_view_setup_probes_the_configured_database_schema( +async def test_view_setup_probe_resolves_through_the_connection_search_path( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant") @@ -264,7 +261,23 @@ async def test_view_setup_probes_the_configured_database_schema( await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) - assert probe.await_args.args == (_PROBE_SQL, '"litellm_tenant"."LiteLLM_SpendLogs"') + assert probe.await_args.args == (_PROBE_SQL, '"LiteLLM_SpendLogs"') + + +@pytest.mark.asyncio +async def test_view_setup_sets_the_row_count_even_when_view_creation_keeps_failing( + prisma_client: PrismaClient, +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "row_count_set": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count >= 1, + } + assert actual == {"outcome": "timed_out", "row_count_set": True} @pytest.mark.asyncio @@ -302,7 +315,7 @@ async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_ } assert actual == { "outcome": "ready", - "calls": ["probe", "views", "probe", "views", "row_count"], + "calls": ["probe", "row_count", "views", "probe", "row_count", "views"], } @@ -319,15 +332,14 @@ async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_clien } assert actual == { "outcome": "ready", - "calls": ["probe", "probe", "views", "row_count"], + "calls": ["probe", "probe", "row_count", "views"], } @pytest.mark.asyncio async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( - prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture ) -> None: - monkeypatch.delenv("DATABASE_SCHEMA", raising=False) _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): @@ -337,7 +349,7 @@ async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( actual = { "outcome": outcome, "error_count": len(errors), - "names_table": '"public"."LiteLLM_SpendLogs"' in errors[0], + "names_table": '"LiteLLM_SpendLogs"' in errors[0], "tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0], } assert actual == { From 2a35dc5217f7fc6697337985702a1401344519dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:00:03 -0700 Subject: [PATCH 022/127] fix(guardrails): keep the undeliverable rewrite reason through copies and name the responses mismatch --- .../guardrail_translation/handler.py | 50 ++++++++++++++----- .../proxy/policy_engine/pipeline_executor.py | 11 ++-- ...test_openai_responses_guardrail_handler.py | 30 +++++++++-- .../policy_engine/test_pipeline_executor.py | 13 +++++ 4 files changed, 83 insertions(+), 21 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index ddfd2199822..1ef1011591e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -167,6 +167,34 @@ def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCa return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments) +def _undeliverable_tool_call_rewrite_reason( + call_ids: Sequence[str], + tool_call_item_count: int, + post_guardrail_tool_call_count: int, + unresolved_argument_event: bool, + rewritten_call_ids: frozenset[str], + event_call_ids: frozenset[str], +) -> str | None: + if len(call_ids) != tool_call_item_count: + return ( + f"{tool_call_item_count - len(call_ids)} of the stream's {tool_call_item_count} tool call items " + "carry no call_id" + ) + if len(frozenset(call_ids)) != len(call_ids): + return "the stream's tool call items repeat a call_id" + if len(call_ids) != post_guardrail_tool_call_count: + return ( + f"the guardrail returned {post_guardrail_tool_call_count} tool calls for the stream's " + f"{len(call_ids)} tool call items" + ) + if unresolved_argument_event: + return "a tool call argument event names an item_id that no output_item event introduced" + missing_call_ids: Final = sorted(rewritten_call_ids - event_call_ids) + if missing_call_ids: + return f"no stream event carries the rewritten call_id {', '.join(missing_call_ids)}" + return None + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -1110,20 +1138,18 @@ class OpenAIResponsesHandler(BaseTranslation): call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES for event, call_id in zip(stream_events, event_call_ids) ) - if ( - len(call_ids) != len(tool_call_items) - or len(frozenset(call_ids)) != len(call_ids) - or len(call_ids) != len(post_guardrail_tool_calls) - or unresolved_argument_event - or not rewrites_by_call_id.keys() <= frozenset(event_call_ids) - ): + undeliverable_reason: Final = _undeliverable_tool_call_rewrite_reason( + call_ids=call_ids, + tool_call_item_count=len(tool_call_items), + post_guardrail_tool_call_count=len(post_guardrail_tool_calls), + unresolved_argument_event=unresolved_argument_event, + rewritten_call_ids=frozenset(rewrites_by_call_id), + event_call_ids=frozenset(call_id for call_id in event_call_ids if call_id is not None), + ) + if undeliverable_reason is not None: from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite( - guardrail_name, - f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls and the stream's " - f"{len(tool_call_items)} function_call items could not be lined up with them by call_id", - ) + raise UndeliverableStreamRewrite(guardrail_name, undeliverable_reason) for output_item, rewrite in ( (output_item, rewrites_by_call_id[call_id]) for output_item, call_id in zip(tool_call_items, call_ids) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index cfb0dc4b617..0b81e7af84d 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -51,13 +51,16 @@ except ImportError: class UndeliverableStreamRewrite(Exception): def __init__(self, guardrail_name: str, reason: str) -> None: - super().__init__( - f"Guardrail '{guardrail_name}' rewrote the streamed response but the rewrite cannot be written " - f"back to the stream: {reason}" - ) + super().__init__(guardrail_name, reason) self.guardrail_name: Final = guardrail_name self.reason: Final = reason + def __str__(self) -> str: + return ( + f"Guardrail '{self.guardrail_name}' rewrote the streamed response but the rewrite cannot be written " + f"back to the stream: {self.reason}" + ) + def _tool_call_shape(tool_call: object) -> tuple[object, object]: plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 81adb283dcc..b45cd2ec299 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1610,13 +1610,14 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: events = self._ended_custom_tool_call_stream_events() events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}] - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite) as undeliverable: await handler.process_output_streaming_response( responses_so_far=events, guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), litellm_logging_obj=None, deliver_ended_stream_rewrites=True, ) + assert undeliverable.value.reason == "no stream event carries the rewritten call_id call_999" @staticmethod def _bridged_function_call_stream_events() -> List[dict]: @@ -1688,8 +1689,21 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []} @pytest.mark.asyncio - @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"]) - async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch): + @pytest.mark.parametrize( + ("mismatch", "expected_reason"), + [ + ("orphan_call_id", "no stream event carries the rewritten call_id call_999"), + ("duplicate_call_id", "the stream's tool call items repeat a call_id"), + ("missing_call_id", "1 of the stream's 1 tool call items carry no call_id"), + ( + "unknown_argument_item_id", + "a tool call argument event names an item_id that no output_item event introduced", + ), + ], + ) + async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed( + self, mismatch, expected_reason + ): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite handler = OpenAIResponsesHandler() @@ -1697,16 +1711,22 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: envelope_item = events[5]["response"]["output"][0] if mismatch == "orphan_call_id": events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}] - else: + elif mismatch == "duplicate_call_id": events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)] + elif mismatch == "missing_call_id": + events[5]["response"]["output"] = [{key: value for key, value in envelope_item.items() if key != "call_id"}] + else: + events[1]["item_id"] = "fc_unknown" - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite) as undeliverable: await handler.process_output_streaming_response( responses_so_far=events, guardrail_to_apply=self._argument_masking_guardrail(), litellm_logging_obj=None, deliver_ended_stream_rewrites=True, ) + assert undeliverable.value.reason == expected_reason + assert str(undeliverable.value).endswith(f"cannot be written back to the stream: {expected_reason}") @pytest.mark.asyncio async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 6af5e9573f2..6aa7eca0f15 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -6,6 +6,7 @@ Uses mock guardrails to validate pipeline execution without external services. import copy import logging +import pickle from typing import Literal from unittest.mock import MagicMock @@ -1727,3 +1728,15 @@ async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(mon assert chunks[0]["text"] == "[REWRITTEN] hello world" assert [call["response"] for call in masker.calls] == [_native("hello world")] assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")] + + +@pytest.mark.parametrize("clone", [copy.deepcopy, lambda exc: pickle.loads(pickle.dumps(exc))], ids=["deepcopy", "pickle"]) +def test_undeliverable_stream_rewrite_keeps_its_reason_through_a_copy(clone): + original = UndeliverableStreamRewrite("masker", "the translation refused it") + + copied = clone(original) + + assert copied.guardrail_name == "masker" + assert copied.reason == "the translation refused it" + assert str(copied) == str(original) + assert str(copied).endswith("cannot be written back to the stream: the translation refused it") From a82f0a0bd2d31a435755640953ccd43476ac961d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:10:49 -0700 Subject: [PATCH 023/127] fix(auth): reject deactivated JWT users and invalidate cached status --- litellm/proxy/auth/user_api_key_auth.py | 10 +++ .../internal_user_endpoints.py | 2 +- .../proxy/auth/test_user_api_key_auth.py | 67 ++++++++++++++++--- .../test_internal_user_endpoints.py | 45 +++++++++++++ 4 files changed, 115 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index de0131772bc..49aab20e461 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1730,6 +1730,16 @@ async def _user_api_key_auth_builder( ) return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) + if ( + user_object is not None + and isinstance(user_object.metadata, dict) + and user_object.metadata.get("scim_active") is False + ): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", + ) + valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 4832c2f4c21..029a968e156 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1571,7 +1571,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if "model_max_budget" in non_default_values: + if "model_max_budget" in non_default_values or "metadata" in data_json: await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 8593be751fa..dcbc0713404 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2091,7 +2091,8 @@ async def test_auto_register_binds_api_key_to_token_hash(): @pytest.mark.asyncio -async def test_auto_register_first_request_propagates_user_email(): +@pytest.mark.parametrize("active", [True, False]) +async def test_auto_register_first_request_propagates_user_email(active: bool) -> None: """ The first auto-registered JWT request must also carry user_email (resolved from the validated LiteLLM_UserTable), so attribution is consistent with the @@ -2120,6 +2121,7 @@ async def test_auto_register_first_request_propagates_user_email(): user_id="validated-user", user_email="validated@example.com", user_role="internal_user", + metadata={"scim_active": active}, ) mock_jwt_result = { "is_proxy_admin": False, @@ -2150,7 +2152,7 @@ async def test_auto_register_first_request_propagates_user_email(): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))), patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), patch( "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", @@ -2170,8 +2172,22 @@ async def test_auto_register_first_request_propagates_user_email(): "litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping", new_callable=AsyncMock, return_value=auto_registered_key, - ), + ) as auto_register, ): + if not active: + with pytest.raises(ProxyException, match="deactivated via SCIM") as exc: + await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + assert int(exc.value.code) == 401 + auto_register.assert_not_awaited() + return result = await _user_api_key_auth_builder( request=mock_request, api_key=jwt_token, @@ -7315,15 +7331,15 @@ class TestJWTAuthUserEmail: the Prometheus `user_email` label and `user_api_key_user_email` in StandardLogging/SpendLogs metadata, which were always None for JWT traffic.""" - def _jwt_request(self, jwt_token): + def _jwt_request(self, jwt_token, route="/v1/chat/completions"): mock_request = MagicMock() - mock_request.url.path = "/v1/chat/completions" - mock_request.method = "POST" + mock_request.url.path = route + mock_request.method = "GET" if route.endswith("/list") else "POST" mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} return mock_request - async def _run_jwt_auth(self, mock_jwt_result, jwt_token): + async def _run_jwt_auth(self, mock_jwt_result, jwt_token, route="/v1/chat/completions"): with ( patch( "litellm.proxy.proxy_server.general_settings", @@ -7344,7 +7360,7 @@ class TestJWTAuthUserEmail: litellm_jwtauth=LiteLLM_JWTAuth(), ) return await user_api_key_auth( - request=self._jwt_request(jwt_token), + request=self._jwt_request(jwt_token, route), api_key=f"Bearer {jwt_token}", ) @@ -7376,6 +7392,41 @@ class TestJWTAuthUserEmail: assert result.user_id == "jwt-human-user" assert result.user_email == "resolved@example.com" + @pytest.mark.asyncio + @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions"]) + @pytest.mark.parametrize("active", [False, True, None, "false", 0]) + async def test_jwt_auth_rejects_deactivated_user(self, route: str, active: bool | str | int | None) -> None: + from typing import Final + + jwt_token: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + result: Final = { + "is_proxy_admin": False, + "team_object": None, + "user_object": LiteLLM_UserTable( + user_id="jwt-human-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + metadata={} if active is None else {"scim_active": active}, + ), + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": None, + "user_id": "jwt-human-user", + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1"}, + } + + if active is False: + with pytest.raises(ProxyException, match="deactivated via SCIM") as exc: + await self._run_jwt_auth(result, jwt_token, route) + assert int(exc.value.code) == 401 + else: + token: Final = await self._run_jwt_auth(result, jwt_token, route) + assert token.user_id == "jwt-human-user" + @pytest.mark.asyncio async def test_jwt_auth_populates_user_email_on_proxy_admin(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 3f2ba365a04..b64f7c8fb7c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2228,6 +2228,51 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize("by_email", [False, True]) +@pytest.mark.parametrize("active", [False, True, None]) +async def test_user_status_update_refreshes_cached_user( + mocker: MockerFixture, by_email: bool, active: bool | None +) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper + + saved_user: Final = LiteLLM_UserTable( + user_id="user-spruce", + user_email="spruce@example.test", + metadata={"scim_active": False if active is None else not active, "department": "engineering"}, + ) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id=None if by_email else saved_user.user_id, + user_email=saved_user.user_email if by_email else None, + metadata={"department": "engineering"} if active is None else {"scim_active": active}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert prisma_client.update_data.call_args.kwargs["user_id"] == saved_user.user_id + assert prisma_client.update_data.call_args.kwargs["data"]["metadata"] == ( + {"department": "engineering"} if active is None else {"scim_active": active} + ) + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + @pytest.mark.asyncio async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None: from litellm.proxy._types import LiteLLM_UserTable From 810acdad97f66635aacc3b67303f1d3890dd8250 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:35:11 -0700 Subject: [PATCH 024/127] chore(streaming): drop the redundant tool-call map comment and restore the OpenAPI snapshot --- litellm/litellm_core_utils/streaming_chunk_builder_utils.py | 2 +- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 903edfbc9c2..fcd55c844c6 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -473,7 +473,7 @@ class ChunkProcessor: tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} # Map to store tool calls by choice and index + tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} for chunk in tool_call_chunks: choices = chunk["choices"] diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 391f0042ed0..06e157498aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19632,7 +19632,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From f8204724880405680649eaa75c937529abfc649b Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 20 Sep 2026 01:44:10 +0000 Subject: [PATCH 025/127] chore: remove the dead telemetry flag from the SDK, proxy CLI and configs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- cookbook/livekit_agent_sdk/config.example.yaml | 1 - litellm/__init__.py | 1 - litellm/proxy/dev_config.yaml | 1 - litellm/proxy/example_config_yaml/oai_misc_config.yaml | 1 - litellm/proxy/proxy_cli.py | 10 ---------- litellm/proxy/proxy_server.py | 4 ---- litellm/proxy/wildcard_config.yaml | 1 - proxy_server_config.yaml | 1 - scripts/benchmark_anthropic_messages_perf.py | 1 - scripts/benchmark_chat_completions_perf.py | 1 - tests/integration/_support/process.py | 2 -- .../test_bedrock_knowledgebase_hook.py | 1 - tests/proxy_unit_tests/test_proxy_utils.py | 1 - .../llms/wandb/test_wandb_chat_transformation.py | 1 - .../test_litellm/proxy/proxy_server/test_lifecycle.py | 6 +++--- 15 files changed, 3 insertions(+), 30 deletions(-) diff --git a/cookbook/livekit_agent_sdk/config.example.yaml b/cookbook/livekit_agent_sdk/config.example.yaml index 1361f36af34..072625018a7 100644 --- a/cookbook/livekit_agent_sdk/config.example.yaml +++ b/cookbook/livekit_agent_sdk/config.example.yaml @@ -15,7 +15,6 @@ model_list: litellm_settings: drop_params: True - telemetry: False general_settings: master_key: sk-1234 # Change this to a secure key diff --git a/litellm/__init__.py b/litellm/__init__.py index e17ab613dac..738dd0cac76 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -240,7 +240,6 @@ email: Optional[str] = ( token: Optional[str] = ( None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 ) -telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = drop_params_env_flag(os.environ, verbose_logger) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) diff --git a/litellm/proxy/dev_config.yaml b/litellm/proxy/dev_config.yaml index f78431f694b..2eb503144e1 100644 --- a/litellm/proxy/dev_config.yaml +++ b/litellm/proxy/dev_config.yaml @@ -212,7 +212,6 @@ sandbox_tools: litellm_settings: drop_params: True - telemetry: False code_interpreter_interception_params: enabled: true sandbox_tool_name: e2b_sandbox diff --git a/litellm/proxy/example_config_yaml/oai_misc_config.yaml b/litellm/proxy/example_config_yaml/oai_misc_config.yaml index 16cc69c19a5..b6b2d0f71f5 100644 --- a/litellm/proxy/example_config_yaml/oai_misc_config.yaml +++ b/litellm/proxy/example_config_yaml/oai_misc_config.yaml @@ -38,7 +38,6 @@ litellm_settings: # budget_duration: 30d num_retries: 5 request_timeout: 600 - telemetry: False context_window_fallbacks: [{"gpt-5-mini": ["gpt-5.5"]}] default_team_settings: - team_id: team-1 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0477b6c62e9..ddb02508c1f 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -56,8 +56,6 @@ if litellm_mode == "DEV": load_dotenv() from enum import Enum -telemetry: Final = None - class LiteLLMDatabaseConnectionPool(Enum): database_connection_pool_limit = 10 @@ -756,12 +754,6 @@ class ProxyInitializationHelpers: type=float, help="Set max budget for API calls - works for hosted models like OpenAI, TogetherAI, Anthropic, etc.`", ) -@click.option( - "--telemetry", - default=True, - type=bool, - help="Helps us know if people are using this feature. Turn this off by doing `--telemetry False`", -) @click.option( "--log_config", default=None, @@ -977,7 +969,6 @@ def run_server( add_function_to_prompt, config, max_budget, - telemetry, test, local, num_workers, @@ -1082,7 +1073,6 @@ def run_server( max_tokens=max_tokens, request_timeout=request_timeout, max_budget=max_budget, - telemetry=telemetry, drop_params=drop_params, add_function_to_prompt=add_function_to_prompt, headers=headers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af25d418a63..df326615764 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2374,7 +2374,6 @@ user_debug = False user_max_tokens = None user_request_timeout = None user_temperature = None -user_telemetry = True user_config: Final = None user_headers = None user_config_file_path: str | None = None @@ -8580,7 +8579,6 @@ async def initialize( max_tokens=None, request_timeout=600, max_budget=None, - telemetry=False, drop_params=True, add_function_to_prompt=True, headers=None, @@ -8596,7 +8594,6 @@ async def initialize( user_user_max_tokens, \ user_request_timeout, \ user_temperature, \ - user_telemetry, \ user_headers, \ experimental, \ llm_model_list, \ @@ -8703,7 +8700,6 @@ async def initialize( dynamic_config["general"]["max_budget"] = litellm.max_budget if experimental: pass - user_telemetry = telemetry # for streaming diff --git a/litellm/proxy/wildcard_config.yaml b/litellm/proxy/wildcard_config.yaml index 7c178690836..e1bae7abe65 100644 --- a/litellm/proxy/wildcard_config.yaml +++ b/litellm/proxy/wildcard_config.yaml @@ -49,4 +49,3 @@ general_settings: litellm_settings: drop_params: True - telemetry: False diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 703d56bc0cd..05df5e31e49 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -173,7 +173,6 @@ litellm_settings: # budget_duration: 30d num_retries: 5 request_timeout: 600 - telemetry: False context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}] default_team_settings: - team_id: team-1 diff --git a/scripts/benchmark_anthropic_messages_perf.py b/scripts/benchmark_anthropic_messages_perf.py index 3c8a22f0cc2..3e4b4b25b6a 100644 --- a/scripts/benchmark_anthropic_messages_perf.py +++ b/scripts/benchmark_anthropic_messages_perf.py @@ -256,7 +256,6 @@ general_settings: master_key: {api_key} litellm_settings: - telemetry: false """, encoding="utf-8", ) diff --git a/scripts/benchmark_chat_completions_perf.py b/scripts/benchmark_chat_completions_perf.py index 2c211f674fe..c9f025a145f 100644 --- a/scripts/benchmark_chat_completions_perf.py +++ b/scripts/benchmark_chat_completions_perf.py @@ -228,7 +228,6 @@ general_settings: litellm_settings: drop_params: true - telemetry: false """, encoding="utf-8", ) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 84ee2ad1b79..e0923c9d055 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -74,8 +74,6 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str], str(port), "--num_workers", "1", - "--telemetry", - "False", "--use_prisma_db_push", "--enforce_prisma_migration_check", ], diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 9faaaf492e8..5571c15beff 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -883,7 +883,6 @@ async def test_provider_specific_fields_in_proxy_http_response( max_tokens=None, request_timeout=600, max_budget=None, - telemetry=False, drop_params=True, add_function_to_prompt=False, headers=None, diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 160753e3442..7cdd7365209 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1112,7 +1112,6 @@ def test_settings_store_preserves_yaml_team_configuration_when_db_value_is_null( }, "param_name": "litellm_settings", "db_param_value": { - "telemetry": False, "drop_params": True, "num_retries": 5, "request_timeout": 600, diff --git a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py index dd0d1bdbb9d..a25ed585ecc 100644 --- a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py +++ b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py @@ -47,7 +47,6 @@ WANDB_REASONING_MODELS: Final = ( @pytest.fixture def wandb_test_config(local_model_cost_map, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(litellm, "telemetry", False) monkeypatch.setattr(litellm, "drop_params", False) diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 6121608b658..ba560cebd52 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -376,7 +376,7 @@ def _lit4152_worker_config_dict(): "master_key": _LIT4152_SECRETS[0], "database_url": _LIT4152_SECRETS[3], "api_key": _LIT4152_SECRETS[2], - "telemetry": True, + "drop_params": True, } @@ -394,7 +394,7 @@ def test__redact_worker_config_for_logging_dict_masks_all_secret_shapes(): assert secret not in rendered, f"leak: {secret} in {rendered!r}" assert isinstance(redacted, dict) assert redacted["model"] == "openai/gpt-4o-mini" - assert redacted["telemetry"] is True + assert redacted["drop_params"] is True def test__redact_worker_config_for_logging_json_string_round_trips_masked(): @@ -501,7 +501,7 @@ def test__redact_worker_config_for_logging_masks_nested_secret_fields(): def test_initialize_signature_is_async_with_expected_params(): sig = inspect.signature(initialize) # Hard-coded so a signature change (param added/removed) trips the gate. - expected_param_count = 17 + expected_param_count = 16 observed = { "is_async": inspect.iscoroutinefunction(initialize), "param_count": len(sig.parameters), From 2e83871d546b38509ecc172f70759422934f5ea9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:52:28 -0700 Subject: [PATCH 026/127] test(guardrails): type the recorder hook's logging_obj as object --- .../chat/guardrail_translation/test_openai_guardrail_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index e60d26cda59..fe97d597ca1 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1392,7 +1392,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: object = None, ) -> GenericGuardrailAPIInputs: self.seen_inputs.append(inputs) return inputs From 6ea74d70f1b67ccf49311271228d1fd9bcd86b7c Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:33:47 -0700 Subject: [PATCH 027/127] fix(auth): enforce SCIM status for admin JWTs and refresh SCIM caches --- litellm/proxy/auth/handle_jwt.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 20 +++---- .../management_endpoints/scim/scim_v2.py | 7 +++ .../proxy/auth/test_handle_jwt.py | 43 +++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 11 ++-- .../scim/test_scim_key_deactivation.py | 60 +++++++++++++++++++ 6 files changed, 128 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 803093ff93a..d621bc634e6 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2463,7 +2463,7 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - if provisioning is None: + if provisioning is None or prisma_client is not None: identity: Final = await JWTAuthManager._resolve_claim_identity( jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 49aab20e461..4371ce4fda8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1714,6 +1714,16 @@ async def _user_api_key_auth_builder( jwt_claims = result.get("jwt_claims", None) agent_id: Final[str | None] = result.get("agent_id") + if ( + user_object is not None + and isinstance(user_object.metadata, dict) + and user_object.metadata.get("scim_active") is False + ): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", + ) + if is_proxy_admin: # Proxy admins authenticate via auth_builder (full # access), not via a mapped virtual key. If @@ -1730,16 +1740,6 @@ async def _user_api_key_auth_builder( ) return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) - if ( - user_object is not None - and isinstance(user_object.metadata, dict) - and user_object.metadata.get("scim_active") is False - ): - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", - ) - valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 2b74dc1e838..b676c0ddb82 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.scim.scim_transformations import ( @@ -1804,6 +1805,9 @@ async def update_user( where={"user_id": user_id}, data=update_data, ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) if client_set_active: new_active: Final = _scim_active_value(metadata) @@ -2375,6 +2379,9 @@ async def patch_user( where={"user_id": user_id}, data=update_data, ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) if new_active is not None and new_active != (True if prev_active is None else prev_active): await _set_user_keys_blocked(user_id=user_id, blocked=not new_active) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 15defb196af..232884d89fa 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -7144,3 +7144,46 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc else: create_team.assert_not_awaited() assert result["team_id"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_user", [False, True]) +@pytest.mark.parametrize("warm_cache", [False, True]) +async def test_scope_admin_admission_resolves_existing_user_without_provisioning( + monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool +) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + private_key, jwk = _get_rsa_key_and_jwk("admin-status") + cache: Final = UserApiKeyCache() + cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk]) + user_id: Final = f"admin-status-{existing_user}-{warm_cache}" + user: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": False}, organization_memberships=[]) + if existing_user and warm_cache: + cache.set_cache(user_id, user) + database: Final = MagicMock() + users: Final = database.db.litellm_usertable + users.find_unique = AsyncMock(return_value=user if existing_user else None) + users.find_first = AsyncMock(return_value=None) + users.create = AsyncMock() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", user_id_upsert=True), + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://admin.example/jwks") + monkeypatch.setenv("JWT_ISSUER", "https://admin.example") + monkeypatch.setenv("JWT_AUDIENCE", "gateway") + token: Final = _encode_rsa_jwt( + private_key, "https://admin.example", "gateway", "admin-status", + {"sub": user_id, "scope": "litellm_proxy_admin"}, + ) + result: Final = await JWTAuthManager.auth_builder( + api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache, + parent_otel_span=None, proxy_logging_obj=MagicMock(), request_data={}, general_settings={}, route="/user/info", + ) + assert result["is_proxy_admin"] is True + assert result["user_id"] == user_id + assert result["user_object"] == (user if existing_user else None) + users.create.assert_not_awaited() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index dcbc0713404..da36071a5b4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7393,18 +7393,21 @@ class TestJWTAuthUserEmail: assert result.user_email == "resolved@example.com" @pytest.mark.asyncio - @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions"]) + @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"]) @pytest.mark.parametrize("active", [False, True, None, "false", 0]) - async def test_jwt_auth_rejects_deactivated_user(self, route: str, active: bool | str | int | None) -> None: + @pytest.mark.parametrize("is_admin", [False, True]) + async def test_jwt_auth_rejects_deactivated_user( + self, route: str, active: bool | str | int | None, is_admin: bool + ) -> None: from typing import Final jwt_token: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" result: Final = { - "is_proxy_admin": False, + "is_proxy_admin": is_admin, "team_object": None, "user_object": LiteLLM_UserTable( user_id="jwt-human-user", - user_role=LitellmUserRoles.INTERNAL_USER.value, + user_role=LitellmUserRoles.PROXY_ADMIN.value if is_admin else LitellmUserRoles.INTERNAL_USER.value, metadata={} if active is None else {"scim_active": active}, ), "end_user_object": None, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py index 0a9cf8b84cb..2cfbfbbb3cb 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py @@ -541,3 +541,63 @@ async def test_scim_put_user_explicit_active_false_blocks_keys(): assert update_kwargs["where"] == {"token": "hash-block-me"} assert update_kwargs["data"]["blocked"] is True assert '"scim_blocked": true' in update_kwargs["data"]["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["PUT", "PATCH"]) +@pytest.mark.parametrize("active", [False, True]) +@pytest.mark.parametrize("failure", [None, "write", "keys"]) +@pytest.mark.parametrize("status_change", [False, True]) +async def test_scim_status_write_refreshes_user_cache( + method: str, active: bool, failure: str | None, status_change: bool +) -> None: + import json + from typing import Final + + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + user_id: Final = "scim-cache-user" + saved: Final = LiteLLM_UserTable( + user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": not active if status_change else active}, + ) + updated: Final = LiteLLM_UserTable( + user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": active}, + ) + client, db = _build_prisma_with_keys([], mock_user=saved.model_copy(deep=True), updated_user=updated) + if failure == "write": + db.litellm_usertable.update.side_effect = RuntimeError("status write failed") + if failure == "keys": + db.litellm_verificationtoken.find_many.side_effect = RuntimeError("key update failed") + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=user_id, value=saved, model_type=LiteLLM_UserTable) + with ( + patch("litellm.proxy.proxy_server.prisma_client", client), # test-quality-ok: substitute the database dependency + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: exercise a real isolated cache + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: isolate the logging dependency + patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=AsyncMock, + ) as broadcast, + ): + request: Final = ( + update_user(user_id=user_id, user=SCIMUser.model_validate(_build_put_user_payload(user_id, active=active))) + if method == "PUT" else + patch_user(user_id=user_id, patch_ops=SCIMPatchOp( + Operations=[SCIMPatchOperation(op="replace", path="active", value=active)] + )) + ) + if failure == "write" or (failure == "keys" and status_change): + with pytest.raises(ProxyException, match="status write failed" if failure == "write" else "key update failed"): + await request + else: + response: Final = await request + assert response.active is active + assert json.loads(db.litellm_usertable.update.await_args.kwargs["data"]["metadata"])["scim_active"] is active + cached: Final = await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) + if failure == "write": + assert cached == saved + broadcast.assert_not_awaited() + else: + assert cached is None + broadcast.assert_awaited_once_with(cache_key=user_id) From 8e530cf819443540192c393957956032ffbe19ea Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 20 Sep 2026 02:36:59 +0000 Subject: [PATCH 028/127] fix: keep --telemetry as a hidden no-op so existing start commands still parse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_cli.py | 8 ++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ddb02508c1f..464d1141f8d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -754,6 +754,14 @@ class ProxyInitializationHelpers: type=float, help="Set max budget for API calls - works for hosted models like OpenAI, TogetherAI, Anthropic, etc.`", ) +@click.option( + "--telemetry", + default=None, + type=bool, + hidden=True, + expose_value=False, + help="Deprecated no-op kept so existing start commands still parse", +) @click.option( "--log_config", default=None, diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c806725d594..8095930bbef 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -617,6 +617,16 @@ class TestProxyInitializationHelpers: assert "Skipping server startup" in result.output mock_uvicorn_run.assert_not_called() + # --- legacy --telemetry flag is accepted, ignored and hidden --- + result = runner.invoke( + run_server, ["--local", "--skip_server_startup", "--telemetry", "False"] + ) + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + assert "Skipping server startup" in result.output + assert "telemetry" not in runner.invoke(run_server, ["--help"]).output + # --- normal startup --- mock_uvicorn_run.reset_mock() From fdd91a347aa9ece3778c6a8e62ffc65e38afc6df Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 20 Sep 2026 02:44:52 +0000 Subject: [PATCH 029/127] test: drop narration comment from telemetry flag test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_cli.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 8095930bbef..8cbae859b5c 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -617,7 +617,6 @@ class TestProxyInitializationHelpers: assert "Skipping server startup" in result.output mock_uvicorn_run.assert_not_called() - # --- legacy --telemetry flag is accepted, ignored and hidden --- result = runner.invoke( run_server, ["--local", "--skip_server_startup", "--telemetry", "False"] ) From 124196cbaa68832894e1491e5d0f28e0b38a86dd Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:48:53 -0700 Subject: [PATCH 030/127] fix(auth): preserve scope-admin email policy during status lookup --- litellm/proxy/auth/handle_jwt.py | 19 +++++++++++++++++-- .../proxy/auth/test_handle_jwt.py | 16 +++++++++++----- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index d621bc634e6..996911cdfaa 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2463,12 +2463,27 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - if provisioning is None or prisma_client is not None: + if provisioning is None: identity: Final = await JWTAuthManager._resolve_claim_identity( jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) return {**admin_result, "user_object": identity.user_object} - return admin_result + if prisma_client is None: + return admin_result + try: + admin_user: Final = await get_user_object( + user_id=user_id, + user_email=user_email, + sso_user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return admin_result + return {**admin_result, "user_object": admin_user} # Get team with model access ## Check if team_id is specified via x-litellm-team-id header diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 232884d89fa..e768139f04a 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -7149,16 +7149,17 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc @pytest.mark.asyncio @pytest.mark.parametrize("existing_user", [False, True]) @pytest.mark.parametrize("warm_cache", [False, True]) +@pytest.mark.parametrize("email", [None, "admin@external.example", "admin@allowed.example"]) async def test_scope_admin_admission_resolves_existing_user_without_provisioning( - monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool + monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool, email: str | None ) -> None: from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache private_key, jwk = _get_rsa_key_and_jwk("admin-status") cache: Final = UserApiKeyCache() cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk]) - user_id: Final = f"admin-status-{existing_user}-{warm_cache}" - user: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": False}, organization_memberships=[]) + user_id: Final = f"admin-status-{existing_user}-{warm_cache}-{email}" + user: Final = LiteLLM_UserTable(user_id=user_id, user_email="admin@allowed.example", metadata={"scim_active": False}, organization_memberships=[]) if existing_user and warm_cache: cache.set_cache(user_id, user) database: Final = MagicMock() @@ -7170,14 +7171,17 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning handler.update_environment( prisma_client=database, user_api_key_cache=cache, - litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", user_id_upsert=True), + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", user_id_upsert=True, user_email_jwt_field="email", + user_allowed_email_domain="allowed.example", + ), ) monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://admin.example/jwks") monkeypatch.setenv("JWT_ISSUER", "https://admin.example") monkeypatch.setenv("JWT_AUDIENCE", "gateway") token: Final = _encode_rsa_jwt( private_key, "https://admin.example", "gateway", "admin-status", - {"sub": user_id, "scope": "litellm_proxy_admin"}, + {"sub": user_id, "scope": "litellm_proxy_admin", **({"email": email} if email else {})}, ) result: Final = await JWTAuthManager.auth_builder( api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache, @@ -7187,3 +7191,5 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning assert result["user_id"] == user_id assert result["user_object"] == (user if existing_user else None) users.create.assert_not_awaited() + if existing_user: + assert users.find_unique.await_count == (0 if warm_cache else 1) From 4fda0092d3aa21df1da30689c9c7660ac0bdcd80 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:52:13 -0700 Subject: [PATCH 031/127] fix(mcp): explain missing public client dependencies --- litellm/experimental_mcp_client/Readme.md | 17 +++++++++ litellm/experimental_mcp_client/__init__.py | 7 +++- .../base_sdk_tests/check_base_sdk_install.py | 12 ++++++ .../test_mcp_client.py | 37 +++++++++++++++++++ 4 files changed, 72 insertions(+), 1 deletion(-) diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md index 14decce0256..3385a37cf69 100644 --- a/litellm/experimental_mcp_client/Readme.md +++ b/litellm/experimental_mcp_client/Readme.md @@ -2,6 +2,17 @@ LiteLLM MCP Client allows you to use MCP tools with LiteLLM +Install the optional dependencies with `pip install 'litellm[mcp]'`, then use the existing public imports: + +```python +from litellm.experimental_mcp_client import call_openai_tool, load_mcp_tools +from litellm.experimental_mcp_client.client import MCPClient + +client = MCPClient(server_url="https://mcp.example.com/mcp") +``` + +Core `import litellm` works without the MCP extra. Importing the experimental MCP client without its MCP or HTTPX2 dependency raises an error with this installation command + ## MCP Python SDK compatibility The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP @@ -16,6 +27,12 @@ The shared unit-test workflow runs the MCP integration suite once, with SDK2 in See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes +## Custom HTTP clients and authentication + +MCP HTTP and SSE transports now use `httpx2`. Custom authentication passed through `aws_auth` or `resolved_auth` must implement `httpx2.Auth`. Integrations that override the client's HTTP client factory or customize its event hooks must use `httpx2.AsyncClient`, request, response, timeout and transport types + +HTTPX1 clients, auth objects and hooks are not adapted by a compatibility shim. Migrate those integrations to HTTPX2 before upgrading. Ordinary `MCPClient` construction and LiteLLM's existing helper imports remain supported; this does not restore SDK1 Python imports or camelCase SDK model attributes in the shared Python environment + ## HTTP redirects For streamable HTTP POST requests, the MCP SDK follows method-preserving redirects such as HTTP 307/308 within the configured endpoint's origin. Redirects to another path on the same scheme, host and port work. The SDK also permits an HTTP-to-HTTPS upgrade on the same host using the default ports diff --git a/litellm/experimental_mcp_client/__init__.py b/litellm/experimental_mcp_client/__init__.py index 5399968ff74..7a3917a50ea 100644 --- a/litellm/experimental_mcp_client/__init__.py +++ b/litellm/experimental_mcp_client/__init__.py @@ -1,3 +1,8 @@ -from .tools import call_openai_tool, load_mcp_tools +try: + from .tools import call_openai_tool, load_mcp_tools +except ModuleNotFoundError as exc: + if exc.name not in ("mcp", "httpx2"): + raise + raise ImportError("MCP client dependencies are missing. Install them with: pip install 'litellm[mcp]'") from exc __all__ = ["call_openai_tool", "load_mcp_tools"] diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 190a900faf9..f680ba92645 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -50,6 +50,17 @@ def check_completion() -> str: return "mock completion round-trips" +def check_mcp_install_guidance() -> str: + try: + import litellm.experimental_mcp_client + except ImportError as error: + _require("pip install 'litellm[mcp]'" in str(error), f"missing MCP installation guidance: {error}") + _require(isinstance(error.__cause__, ModuleNotFoundError), "original missing-dependency cause was lost") + _require(error.__cause__.name == "mcp", f"unexpected missing dependency: {error.__cause__}") + return "optional MCP client explains how to install litellm[mcp]" + raise AssertionError("MCP client imported without the MCP extra") + + def check_embedding() -> str: import litellm @@ -109,6 +120,7 @@ def check_bedrock_credential_resolution() -> str: CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("environment is base-only", check_environment_is_base_only), ("import litellm", check_import), + ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), ("embedding", check_embedding), ("bundled model metadata", check_bundled_model_metadata), diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 7e4598c2e58..4b698f1258d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,10 +1,12 @@ import asyncio import base64 +import importlib import json import os import sys from collections.abc import AsyncIterator from pathlib import Path +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2160,3 +2162,38 @@ async def test_404_before_session_initialization_preserves_method_not_found() -> ) assert caught.value.error.code == METHOD_NOT_FOUND assert caught.value.error.message == "Not Found" + + +@pytest.mark.parametrize("missing_module", ("mcp", "httpx2", "mcp.types", "openai.types.chat")) +def test_public_mcp_import_missing_dependency(missing_module: str) -> None: + with patch.dict(sys.modules): + for name in tuple(sys.modules): + if name.startswith(("litellm.experimental_mcp_client", "mcp.", "mcp_types.")) or name == "mcp": + del sys.modules[name] + with patch.dict(sys.modules, {missing_module: None}): + with pytest.raises(ImportError) as caught: + importlib.import_module("litellm.experimental_mcp_client.client") + + if missing_module in ("mcp", "httpx2"): + assert "pip install 'litellm[mcp]'" in str(caught.value) + assert isinstance(caught.value.__cause__, ModuleNotFoundError) + assert caught.value.__cause__.name == missing_module + else: + assert isinstance(caught.value, ModuleNotFoundError) + assert caught.value.name == missing_module + assert caught.value.__cause__ is None + assert "litellm[mcp]" not in str(caught.value) + + +def test_public_mcp_import_preserves_incompatible_sdk_error() -> None: + with patch.dict(sys.modules): + for name in tuple(sys.modules): + if name.startswith("litellm.experimental_mcp_client"): + del sys.modules[name] + with patch.dict(sys.modules, {"mcp": ModuleType("mcp")}): + with pytest.raises(ImportError, match="cannot import name 'ClientSession'") as caught: + importlib.import_module("litellm.experimental_mcp_client.client") + + assert not isinstance(caught.value, ModuleNotFoundError) + assert caught.value.__cause__ is None + assert "litellm[mcp]" not in str(caught.value) From b2e123da4332344677fba1b819dc423864cee89c Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 20 Sep 2026 03:14:11 +0000 Subject: [PATCH 032/127] fix(proxy): drop legacy telemetry key from persisted WORKER_CONFIG before initialize Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/template.yaml | 2 +- cookbook/misc/config.yaml | 1 - litellm/proxy/proxy_server.py | 11 +++++++++-- .../test_litellm/proxy/proxy_server/test_lifecycle.py | 11 +++++++++++ 4 files changed, 21 insertions(+), 4 deletions(-) diff --git a/.github/template.yaml b/.github/template.yaml index d4db2c2ac1f..c77e578e53c 100644 --- a/.github/template.yaml +++ b/.github/template.yaml @@ -21,7 +21,7 @@ Parameters: WorkerConfigParameter: Type: String Description: Sample environment variable - Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "telemetry": true, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}' + Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}' Resources: MyUrlFunctionPermissions: diff --git a/cookbook/misc/config.yaml b/cookbook/misc/config.yaml index d1d06eb5842..27a6332a882 100644 --- a/cookbook/misc/config.yaml +++ b/cookbook/misc/config.yaml @@ -55,7 +55,6 @@ litellm_settings: # budget_duration: 30d num_retries: 5 request_timeout: 600 - telemetry: False context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}] general_settings: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index df326615764..9c4ea7e8599 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1208,12 +1208,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: general_settings, ) = await proxy_config.load_config(router=llm_router, config_file_path=worker_config) elif isinstance(worker_config, dict): - await initialize(**worker_config) + await initialize_from_worker_config(worker_config) else: # if not, assume it's a json string worker_config = json.loads(worker_config) if isinstance(worker_config, dict): - await initialize(**worker_config) + await initialize_from_worker_config(worker_config) # check if DATABASE_URL in environment - load from there if prisma_client is None: @@ -8568,6 +8568,13 @@ def save_worker_config(**data): os.environ["WORKER_CONFIG"] = json.dumps(data) +LEGACY_WORKER_CONFIG_KEYS: Final = frozenset({"telemetry"}) + + +async def initialize_from_worker_config(worker_config: dict[str, object]) -> None: + await initialize(**{k: v for k, v in worker_config.items() if k not in LEGACY_WORKER_CONFIG_KEYS}) + + async def initialize( model=None, alias=None, diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index ba560cebd52..e9695685ca5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -43,6 +43,7 @@ from litellm.proxy.proxy_server import ( cost_tracking, get_litellm_model_info, initialize, + initialize_from_worker_config, load_from_azure_key_vault, proxy_shutdown_event, proxy_startup_event, @@ -522,6 +523,16 @@ async def test_initialize_invalid_unexpected_kwarg_raises_type_error(): await initialize(this_is_not_a_real_kwarg=True) +@pytest.mark.asyncio +async def test_initialize_from_worker_config_drops_legacy_telemetry_key(): + with pytest.raises(TypeError): + await initialize(telemetry=True) + await initialize_from_worker_config({"telemetry": True, "request_timeout": 77}) + assert ps.user_request_timeout == 77 + with pytest.raises(TypeError): + await initialize_from_worker_config({"this_is_not_a_real_kwarg": True}) + + # --------------------------------------------------------------------------- # load_from_azure_key_vault # --------------------------------------------------------------------------- From f7756d01cf038097b35fbcbb79246971761b1ed2 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 20 Sep 2026 03:21:01 +0000 Subject: [PATCH 033/127] refactor(proxy): freeze the filtered worker config before initialize Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c4ea7e8599..72bfb2b53a9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8571,8 +8571,9 @@ def save_worker_config(**data): LEGACY_WORKER_CONFIG_KEYS: Final = frozenset({"telemetry"}) -async def initialize_from_worker_config(worker_config: dict[str, object]) -> None: - await initialize(**{k: v for k, v in worker_config.items() if k not in LEGACY_WORKER_CONFIG_KEYS}) +async def initialize_from_worker_config(worker_config: Mapping[str, object]) -> None: + supported: Final = MappingProxyType({k: v for k, v in worker_config.items() if k not in LEGACY_WORKER_CONFIG_KEYS}) + await initialize(**supported) async def initialize( From 19d77e244260c7c0546173f8ef75f34eab812372 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:21:37 -0700 Subject: [PATCH 034/127] test(guardrails): type the recorder hook's request_data as a Mapping --- .../guardrail_translation/test_openai_guardrail_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index fe97d597ca1..258226ae22c 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -6,6 +6,7 @@ with guardrail transformations, including tool calls. """ import json +from collections.abc import Mapping from typing import Any, Literal, Optional import pytest @@ -1390,7 +1391,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: object = None, ) -> GenericGuardrailAPIInputs: From 37204356389bbdc8986066cd84a5588f332f6f4b Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 04:00:57 +0000 Subject: [PATCH 035/127] chore(prices): sync OpenRouter prices: 1 model openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78848d7eea8..6f8db4d2215 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -67138,9 +67138,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.696e-08, - "output_cost_per_token": 7.392e-08, - "cache_read_input_token_cost": 7.392e-09, + "input_cost_per_token": 3.668e-08, + "output_cost_per_token": 7.336e-08, + "cache_read_input_token_cost": 7.336e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78848d7eea8..6f8db4d2215 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -67138,9 +67138,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.696e-08, - "output_cost_per_token": 7.392e-08, - "cache_read_input_token_cost": 7.392e-09, + "input_cost_per_token": 3.668e-08, + "output_cost_per_token": 7.336e-08, + "cache_read_input_token_cost": 7.336e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From 3619142a529ed1a0bf36d21c739be76cb3589081 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 04:30:59 +0000 Subject: [PATCH 036/127] chore(prices): sync OpenRouter prices: 1 model openrouter/z-ai/glm-5.2: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6f8db4d2215..11da46f8eb3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66776,9 +66776,9 @@ "supports_web_search": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 5.544e-07, - "output_cost_per_token": 1.7424e-06, - "cache_read_input_token_cost": 1.0296e-07, + "input_cost_per_token": 6.496e-07, + "output_cost_per_token": 2.0416e-06, + "cache_read_input_token_cost": 1.2064e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6f8db4d2215..11da46f8eb3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66776,9 +66776,9 @@ "supports_web_search": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 5.544e-07, - "output_cost_per_token": 1.7424e-06, - "cache_read_input_token_cost": 1.0296e-07, + "input_cost_per_token": 6.496e-07, + "output_cost_per_token": 2.0416e-06, + "cache_read_input_token_cost": 1.2064e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, From 3f4fe7db82df189d1c4a978f7c0dbd666f759ce4 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 06:24:06 +0000 Subject: [PATCH 037/127] docs(tests): define the tier contract for unit, integration and e2e Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/AGENTS.md | 38 +++++++++++++++++++++++++++++++++++++ tests/e2e/AGENTS.md | 24 +++++++++++++++++++++++ tests/integration/AGENTS.md | 25 ++++++++++++++++++++++++ tests/unit/AGENTS.md | 37 ++++++++++++++++++++++++++++++++++++ 4 files changed, 124 insertions(+) create mode 100644 tests/AGENTS.md create mode 100644 tests/integration/AGENTS.md create mode 100644 tests/unit/AGENTS.md diff --git a/tests/AGENTS.md b/tests/AGENTS.md new file mode 100644 index 00000000000..ad2b8d95eaf --- /dev/null +++ b/tests/AGENTS.md @@ -0,0 +1,38 @@ +# Tests + +Nothing on the other side of the call: `tests/unit`. A proxy we start with an upstream we script: +`tests/integration`. Someone else's service with real credentials: `tests/e2e`. Two fit, split it + +## What good looks like + +Red when the claim in the name is broken. Prove it: mutate the behaviour, red; restore, green. Put the +mutation in the PR body + +```python +def test_custom_price_is_reported_and_charged(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + response = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "price control"}]}) + assert response.status_code == 200, response.text + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002) +``` + +Rates in the test, expected computed by hand, one call, `response.text` in the assert + +Assert the whole value. Iterating `expected_body.items()` (`test_responses_api_request_body.py`) cannot +see an extra key; that is the shape of `stream_options.include_usage` (#19777, #28553) + +The linter catches no-assert, mock-echo, credential skips and patched internals. It cannot see an assert +behind an `if` (a poll that ends in `pytest.fail` is fine), `except Exception` around the call +(`test_router.py`: `except Exception as e: print(f"FAILED TEST")`), or blanket `--reruns` + +## Where it goes + +What the assertion depends on goes in the test; everything else in conftest. A rate in a fixture three +directories up makes a failed assertion unreadable. Extend the file that already covers the behaviour + +## Writing it so a human can read it + +Name says what broke: `test_send_batched_with_valid_data` says nothing. Build, one call, assert, on one +screen. Helpers named for what they return, `_pii_prompt(marker, email)`, not `_setup()`. Context in the +assert message, not a comment diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 9b662e511b8..0c4e704811f 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -2,6 +2,30 @@ Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `AGENTS.md` +## What good looks like + +Only what a real provider proves. If it holds against our scripted upstream: `tests/integration` + +```python +def test_pre_call_masks_pii_on_chat_completions(self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str) -> None: + name = f"e2e-presidio-pre-chat-{unique_marker()}" + _register_presidio(client, resources, name=name) + email = _fake_email() + _assert_eventually_masked( + lambda: client.chat(scoped_key, MODEL, _pii_prompt(unique_marker(), email), guardrails=[name], max_tokens=128), + _first_content, + email=email, + ) +``` + +Marker per run, so a leftover guardrail cannot pass it. `resources.defer(...)` at creation, so a failed +assert still tears down. Assert what the caller receives + +## Where it goes + +By the surface a customer would name: `guardrails`, `llm_translation`, `management`. Mutation check +deferred; it needs credentials + ## Suite folders Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family or behavior area. If you add a new folder, you must add a line here describing what kind of tests belong in it, so the layout stays self-describing. `gateway/` is the exception: it holds proxy configuration only and never tests diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md new file mode 100644 index 00000000000..57b69f3830d --- /dev/null +++ b/tests/integration/AGENTS.md @@ -0,0 +1,25 @@ +# tests/integration + +Real proxy, Postgres, Redis, scripted upstream. `README.md` has shards and CI wiring + +## What good looks like + +The root example is from here. The spend row lands async: poll, never sleep + +```python +rows = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],)), + lambda values: len(values) == 1, + seconds=70, +) +assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0.002 + 4 * 0.002) +``` + +`sleep(3)` fails on a slow runner and taxes every fast one. Assert the outbound body in the upstream +handler; a leaked field is invisible from the response. `monkeypatch.setenv` is fine; patching our own +function in a full stack is not + +## Where it goes + +By the domain a user would name: `pricing`, `spend`, `routing`. Add the node and its `covers` ids to +`contracts.json` or collection fails. Needs no proxy, DB or Redis: `tests/unit` diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md new file mode 100644 index 00000000000..191777f3e83 --- /dev/null +++ b/tests/unit/AGENTS.md @@ -0,0 +1,37 @@ +# tests/unit + +In-process. No network, clock or subprocess + +## What good looks like + +```python +def test_send_result_same_version_is_identity_passthrough(): + rpc = _rpc(V03_MESSAGE) + out = normalize_jsonrpc_response(rpc, "0.3", method="message/send") + assert out is rpc +``` + +`is`, because `==` passes on a copy. Many inputs: parametrize +(`test_an_incomplete_reservation_accrues_nothing`, fifteen cases, fifteen results) + +No doubles on our own code + +```python +with patch.object(streamer, "_group_by_date") as mock_group, patch.object(streamer, "_send_daily_batch") as mock_send: + mock_group.return_value = {"2025-01-19": pl.DataFrame({"test": ["data1"]}), "2025-01-20": pl.DataFrame({"test": ["data2"]})} + streamer.send_batched(pl.DataFrame({"test": ["data"]}), "replace_hourly") + assert mock_send.call_count == 2 +``` + +Green if `send_batched` drops every row. pydantic doubles in 12 of 203 files, fastapi 11 of 594; +`tests/test_litellm` 59 percent. Exception: the count is the behaviour +(`test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads`, fifty readers, `call_count == 1`) + +## Where it goes + +`tests/unit/` mirrors `litellm/`, so a changed file selects its tests by path, not a mapping +file. Empty today; new unit tests go here. The examples above live in `tests/test_litellm` + +## Writing it so a human can read it + +A class only when tests share an arrange From 446bd1b250fee07538803a2fe813a601eb7064d4 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:05:42 +0000 Subject: [PATCH 038/127] ci(tests): wire tests/unit into CircleCI and drain legacy unit shards green Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/config.yml | 29 ++++ .github/workflows/_test-unit-base.yml | 60 ++++---- pyproject.toml | 1 + scripts/check_test_quality.py | 66 --------- test-quality-budget.json | 3 - tests/AGENTS.md | 2 +- tests/test_litellm/test_check_test_quality.py | 137 +----------------- tests/test_litellm/test_test_quality_gate.py | 2 +- tests/unit/conftest.py | 16 ++ uv.lock | 14 ++ 10 files changed, 101 insertions(+), 229 deletions(-) create mode 100644 tests/unit/conftest.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 4b24ab58930..46aacd56caa 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3026,6 +3026,33 @@ jobs: - store_artifacts: path: test-results + unit: + machine: + image: ubuntu-2204:2024.04.1 + resource_class: large + working_directory: ~/project + steps: + - setup_litellm_test_deps + - run: + name: Generate Prisma client + command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + - run: + name: Run unit tests + command: | + mkdir -p test-results/unit + mapfile -t files < <(find tests/unit -name 'test_*.py' | sort) + if [ "${#files[@]}" -eq 0 ]; then echo "tests/unit holds no test_*.py files; nothing to run"; exit 0; fi + set +e + LITELLM_LOCAL_MODEL_COST_MAP=True uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --junitxml=test-results/unit/junit.xml + status=$? + set -e + if [ "$status" -eq 5 ]; then echo "pytest collected no tests from tests/unit; passing"; exit 0; fi + exit "$status" + - store_test_results: + path: test-results + - store_artifacts: + path: test-results + workflows: integration: jobs: @@ -3047,6 +3074,8 @@ workflows: only: - main - /litellm_.*/ + - unit: + filters: *main_branches - provider_replay_harness - base_sdk_install: filters: *main_branches diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index d4e9a65e7c0..f4d4fb54ba6 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -165,33 +165,41 @@ jobs: DIST: ${{ inputs.dist }} COVERAGE_CORE: sysmon run: | - if [ "${WORKERS}" = "0" ]; then - uv run --no-sync pytest ${TEST_PATH:?} \ - --tb=short -vv \ - --maxfail="${MAX_FAILURES}" \ - --reruns "${RERUNS}" \ - --reruns-delay 1 \ - --timeout="${TEST_TIMEOUT_SECONDS}" \ - --rerun-except "from pytest-timeout" \ - --durations=20 \ - --cov=./litellm --cov=./enterprise/litellm_enterprise \ - --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml - else - uv run --no-sync pytest ${TEST_PATH:?} \ - --tb=short -vv \ - --maxfail="${MAX_FAILURES}" \ - -n "${WORKERS}" \ - --reruns "${RERUNS}" \ - --reruns-delay 1 \ - --timeout="${TEST_TIMEOUT_SECONDS}" \ - --rerun-except "from pytest-timeout" \ - --dist="${DIST}" \ - --durations=20 \ - --cov=./litellm --cov=./enterprise/litellm_enterprise \ - --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml + found_path=false + for path in ${TEST_PATH}; do + if [ -e "${path%%::*}" ]; then + found_path=true + break + fi + done + if [ "$found_path" = false ]; then + echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run" + exit 0 fi + xdist_args=() + if [ "${WORKERS}" != "0" ]; then + xdist_args=(-n "${WORKERS}" --dist="${DIST}") + fi + set +e + uv run --no-sync pytest ${TEST_PATH:?} \ + --tb=short -vv \ + --maxfail="${MAX_FAILURES}" \ + "${xdist_args[@]}" \ + --reruns "${RERUNS}" \ + --reruns-delay 1 \ + --timeout="${TEST_TIMEOUT_SECONDS}" \ + --rerun-except "from pytest-timeout" \ + --durations=20 \ + --cov=./litellm --cov=./enterprise/litellm_enterprise \ + --cov-report=xml:coverage.xml \ + --cov-config=pyproject.toml + status=$? + set -e + if [ "$status" -eq 5 ]; then + echo "pytest collected no tests from ${TEST_PATH}; passing" + exit 0 + fi + exit "$status" - name: Save coverage report if: always() && steps.changes.outputs.decision != 'skip' diff --git a/pyproject.toml b/pyproject.toml index 821f885dbfc..3b17a397f3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -196,6 +196,7 @@ dev = [ "mypy==1.20.1", "keyring==25.7.0", "pytest==9.0.3", + "pytest-socket==0.8.1", "tomli==2.4.1; python_version < '3.11'", "pytest-mock==3.15.1", "pytest-asyncio==1.3.0", diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 1ef4aed8675..dddc9d61982 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -48,10 +48,6 @@ TQ006 A `pytest.skip` reached only when a credential-shaped environment variab deliberate branch. The gate follows one local or module-level binding, which is the `key = os.getenv(...)` then `if not key: pytest.skip(...)` shape most of these use. -TQ008 A `patch(...)` whose target is a `litellm.` internal. Patching the SDK's own - functions pins the test to the current wiring instead of the behaviour, and it - is the idiom the suite reaches for instead of faking the HTTP boundary. Mocking - a third-party client, a transport, or anything outside `litellm.` is untouched. TQ007 A module global that a conftest saves before every test and restores after it. The save/restore list is a hand-maintained inventory of the leaks the suite already knows about, so it is allowed to shrink and never to grow: a new entry @@ -483,67 +479,6 @@ def iter_global_mutation_violations(path: Path, tree: ast.Module) -> Iterator[Vi ) -def _is_sdk_internal(dotted: str) -> bool: - return dotted == SDK_MODULE or dotted.startswith(f"{SDK_MODULE}.") - - -def _sdk_import_bindings(tree: ast.Module) -> Iterator[tuple[str, str]]: - """(local name, dotted path) for every import that binds something under `litellm`.""" - for node in ast.walk(tree): - if isinstance(node, ast.Import): - yield from ( - (alias.asname, alias.name) if alias.asname else (root, root) - for alias in node.names - if _is_sdk_internal(alias.name) - for root in (alias.name.partition(".")[0],) - ) - elif isinstance(node, ast.ImportFrom) and node.module and _is_sdk_internal(node.module): - yield from ((alias.asname or alias.name, f"{node.module}.{alias.name}") for alias in node.names) - - -def _sdk_aliases(tree: ast.Module) -> Mapping[str, str]: - """Local names bound to something under `litellm`, mapped to the path they stand for. - - `from litellm.llms.openai.chat import handler` then `patch.object(handler.X, ...)` - reaches the same internal as the dotted string form and has to read the same way. - """ - return MappingProxyType({name: dotted for name, dotted in _sdk_import_bindings(tree)}) - - -def _resolved(dotted: str, aliases: Mapping[str, str]) -> str: - root, _, rest = dotted.partition(".") - base: Final = aliases.get(root, root) - return f"{base}.{rest}" if rest else base - - -def _patch_targets(call: ast.Call, aliases: Mapping[str, str]) -> Iterator[str]: - """What a patch installer is replacing: the dotted string it names, or the - attribute chain handed to `patch.object` / `patch.dict`, resolved through the - module's imports so a locally bound SDK object reads as its full path.""" - for first in call.args[:1]: - if isinstance(first, ast.Constant) and isinstance(first.value, str): - yield first.value - elif dotted := _dotted_name(first): - yield _resolved(dotted, aliases) - - -def iter_internal_patch_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: - aliases: Final = _sdk_aliases(tree) - for node in ast.walk(tree): - if not (isinstance(node, ast.Call) and _is_patch_installer(_dotted_name(node.func))): - continue - for target in _patch_targets(node, aliases): - if _is_sdk_internal(target): - yield Violation( - path, - node.lineno, - "TQ008", - f"patches `{target}`, an SDK internal, so the test is pinned to how the code is " - "wired rather than what it does; fake the HTTP boundary (respx / MockTransport) " - f"or inject the collaborator (suppress: `# {SUPPRESSION_TOKEN}: `)", - ) - - def _environ_keys(node: ast.AST) -> Iterator[str]: for inner in ast.walk(node): if isinstance(inner, ast.Call) and _dotted_name(inner.func) in ENVIRON_READERS: @@ -784,7 +719,6 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_global_mutation_violations(path, tree), *iter_credential_skip_violations(path, tree), *iter_conftest_inventory_violations(path, tree), - *iter_internal_patch_violations(path, tree), *iter_child_interpreter_violations(path, tree), ) if violation.line not in skip diff --git a/test-quality-budget.json b/test-quality-budget.json index ae4ea4d31be..6f3dde8461b 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -20,9 +20,6 @@ "TQ007": { "limit": 117 }, - "TQ008": { - "limit": 10993 - }, "TQ009": { "limit": 59 } diff --git a/tests/AGENTS.md b/tests/AGENTS.md index ad2b8d95eaf..13b4789003f 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -22,7 +22,7 @@ Rates in the test, expected computed by hand, one call, `response.text` in the a Assert the whole value. Iterating `expected_body.items()` (`test_responses_api_request_body.py`) cannot see an extra key; that is the shape of `stream_options.include_usage` (#19777, #28553) -The linter catches no-assert, mock-echo, credential skips and patched internals. It cannot see an assert +The linter catches no-assert, mock-echo and credential skips. It cannot see an assert behind an `if` (a poll that ends in `pytest.fail` is fine), `except Exception` around the call (`test_router.py`: `except Exception as e: print(f"FAILED TEST")`), or blanket `--reruns` diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index bfe503e74d1..05c25fb19fb 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -187,7 +187,7 @@ def test_mock_echo_is_flagged(tmp_path): " run()\n" " mock_completion.assert_called_once()\n" ) - assert _codes(tmp_path, source) == ["TQ002", "TQ008"] + assert _codes(tmp_path, source) == ["TQ002"] def test_call_args_inspection_is_mock_echo(tmp_path): @@ -200,7 +200,7 @@ def test_call_args_inspection_is_mock_echo(tmp_path): " run()\n" " assert mock_completion.call_args[1]['model'] == 'gpt-4o'\n" ) - assert _codes(tmp_path, source) == ["TQ002", "TQ008"] + assert _codes(tmp_path, source) == ["TQ002"] def test_patch_decorator_counts_as_installing_a_patch(tmp_path): @@ -213,7 +213,7 @@ def test_patch_decorator_counts_as_installing_a_patch(tmp_path): " run()\n" " mock_completion.assert_called_once()\n" ) - assert _codes(tmp_path, source) == ["TQ002", "TQ008"] + assert _codes(tmp_path, source) == ["TQ002"] def test_patching_but_asserting_the_output_is_not_mock_echo(tmp_path): @@ -227,7 +227,7 @@ def test_patching_but_asserting_the_output_is_not_mock_echo(tmp_path): " mock_completion.assert_called_once()\n" " assert result.choices[0].message.content == 'pong'\n" ) - assert _codes(tmp_path, source) == ["TQ008"] + assert _codes(tmp_path, source) == [] def test_asserting_without_patching_is_not_mock_echo(tmp_path): @@ -244,7 +244,7 @@ def test_a_test_with_no_assertions_is_tq001_not_tq002(tmp_path): " with patch('litellm.completion'):\n" " run()\n" ) - assert _codes(tmp_path, source) == ["TQ001", "TQ008"] + assert _codes(tmp_path, source) == ["TQ001"] def test_sys_path_insert_is_flagged(tmp_path): @@ -554,133 +554,6 @@ def test_a_loop_storing_under_a_key_that_is_not_the_loop_variable_is_not_an_inve assert [v.code for v in checker.check_file(_written(tmp_path, source))] == [] -def test_patching_an_sdk_function_by_string_is_flagged(tmp_path): - source = 'from unittest.mock import patch\n\n\n@patch("litellm.completion")\ndef test_x(m):\n assert m\n' - assert "TQ008" in _codes(tmp_path, source) - - -def test_patching_a_deep_sdk_path_is_flagged(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch("litellm.llms.openai.chat.handler.OpenAIChatCompletion.completion"):\n' - " assert True\n" - ) - assert "TQ008" in _codes(tmp_path, source) - - -def test_patch_object_rooted_at_the_sdk_is_flagged(tmp_path): - source = ( - "import litellm\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(litellm, "api_key", "x"):\n' - " assert True\n" - ) - assert "TQ008" in _codes(tmp_path, source) - - -def test_patch_object_on_a_from_imported_sdk_module_is_flagged(tmp_path): - source = ( - "from litellm.llms.openai.chat import handler\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(handler.OpenAIChatCompletion, "completion"):\n' - " assert True\n" - ) - assert "TQ008" in _codes(tmp_path, source) - - -def test_patch_object_on_an_aliased_sdk_module_is_flagged(tmp_path): - source = ( - "import litellm.llms.openai.chat.handler as oai\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(oai.OpenAIChatCompletion, "completion"):\n' - " assert True\n" - ) - assert "TQ008" in _codes(tmp_path, source) - - -def test_patch_object_on_a_renamed_sdk_symbol_is_flagged(tmp_path): - source = ( - "from litellm.utils import get_llm_provider as glp\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(glp, "__wrapped__"):\n' - " assert True\n" - ) - assert "TQ008" in _codes(tmp_path, source) - - -def test_the_reported_target_is_the_resolved_sdk_path(tmp_path): - source = ( - "from litellm.llms.openai.chat import handler\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(handler.OpenAIChatCompletion, "completion"):\n' - " assert True\n" - ) - reported = [v.message for v in checker.check_file(_written(tmp_path, source)) if v.code == "TQ008"] - assert reported - assert "litellm.llms.openai.chat.handler.OpenAIChatCompletion" in reported[0] - - -def test_patch_object_on_a_from_imported_third_party_is_not_flagged(tmp_path): - source = ( - "from openai import OpenAI\nfrom unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch.object(OpenAI, "chat"):\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - -def test_a_local_name_with_no_sdk_import_behind_it_is_not_flagged(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x(handler):\n" - ' with patch.object(handler, "completion"):\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - -def test_mocking_a_third_party_client_is_not_flagged(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch("openai.OpenAI.chat"):\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - -def test_mocking_the_http_transport_is_not_flagged(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch("httpx.AsyncClient.send"):\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - -def test_a_name_merely_starting_with_litellm_is_not_the_sdk(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch("litellm_enterprise.thing.go"):\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - -def test_an_sdk_patch_can_be_suppressed(tmp_path): - source = ( - "from unittest.mock import patch\n\n\n" - "def test_x():\n" - ' with patch("litellm.completion"): # test-quality-ok: pinning the router seam\n' - " assert True\n" - ) - assert "TQ008" not in _codes(tmp_path, source) - - _FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1 _SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare" diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index cde33787c6c..873e7b9b336 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -137,7 +137,7 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) assert set(budget) == { - "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008", "TQ009" + "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009" } assert all(spec["limit"] >= 0 for spec in budget.values()) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000000..253ae7a6119 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,16 @@ +from collections.abc import Iterator + +import pytest +from pytest_socket import enable_socket, socket_allow_hosts + + +@pytest.fixture(autouse=True, scope="session") +def block_external_sockets() -> Iterator[None]: + socket_allow_hosts(["127.0.0.1", "::1"], allow_unix_socket=True) + yield + enable_socket() + + +@pytest.hookimpl(trylast=True) +def pytest_runtest_setup() -> None: + socket_allow_hosts(["127.0.0.1", "::1"], allow_unix_socket=True) diff --git a/uv.lock b/uv.lock index f1a58500a61..dfba6bfc258 100644 --- a/uv.lock +++ b/uv.lock @@ -4702,6 +4702,7 @@ dev = [ { name = "pytest-postgresql" }, { name = "pytest-recording" }, { name = "pytest-rerunfailures" }, + { name = "pytest-socket" }, { name = "pytest-timeout" }, { name = "pytest-xdist" }, { name = "reportlab" }, @@ -4902,6 +4903,7 @@ dev = [ { name = "pytest-postgresql", specifier = "==7.0.2" }, { name = "pytest-recording", specifier = "==0.13.4" }, { name = "pytest-rerunfailures", specifier = "==15.1" }, + { name = "pytest-socket", specifier = "==0.8.1" }, { name = "pytest-timeout", specifier = "==2.4.0" }, { name = "pytest-xdist", specifier = "==3.8.0" }, { name = "reportlab", specifier = "==5.0.1" }, @@ -8087,6 +8089,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7c/ff/3266c8a73b9b93c4b14160a7e2b31d1e1088e28ed29f4c2d93ae34093bfd/pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4", size = 13775, upload-time = "2025-01-19T01:56:11.199Z" }, ] +[[package]] +name = "pytest-socket" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ba/ce/4ef7b049852c95a8727b4a7e6496f762df1ac0b47bc0320d10293f5e95ec/pytest_socket-0.8.1.tar.gz", hash = "sha256:2f57787914ad2e1308d09ce141b95c3e55741fbb4fb7b7556593a6b063e0c9c7", size = 17313, upload-time = "2026-08-19T15:16:25.653Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/ef/ab507f117b3d19b54e3c9c632a99c28c3b284562ec6e02e274581d530d92/pytest_socket-0.8.1-py3-none-any.whl", hash = "sha256:f9846bed1dcd96eed459e5e14795bbaf96715cf4e827891fe70773817ecb8ed4", size = 8751, upload-time = "2026-08-19T15:16:24.426Z" }, +] + [[package]] name = "pytest-timeout" version = "2.4.0" From df7a3d2d1ea1612691f4afefad16e9373c338fbe Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:06:39 +0000 Subject: [PATCH 039/127] ci(tests): share the loopback allow list in the unit conftest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 253ae7a6119..3bdab1d231a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,16 +1,23 @@ from collections.abc import Iterator +from typing import Final import pytest from pytest_socket import enable_socket, socket_allow_hosts +LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] + + +def _allow_loopback_only() -> None: + socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) + @pytest.fixture(autouse=True, scope="session") def block_external_sockets() -> Iterator[None]: - socket_allow_hosts(["127.0.0.1", "::1"], allow_unix_socket=True) + _allow_loopback_only() yield enable_socket() @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: - socket_allow_hosts(["127.0.0.1", "::1"], allow_unix_socket=True) + _allow_loopback_only() From 3c094dbaaa85ea39d5156285fa10ddc340a89c5d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:08:09 +0000 Subject: [PATCH 040/127] ci(tests): temporarily point one shard at an empty directory Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index aa82a0bf3ee..3897a0df3e8 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -51,7 +51,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/mcp_tests" + test-path: "tests/unit" workers: 2 reruns: 0 timeout-minutes: 20 From adadeac24507d689d4f6f6992ac393d11a6633cc Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:10:43 +0000 Subject: [PATCH 041/127] Revert "ci(tests): temporarily point one shard at an empty directory" This reverts commit 3c094dbaaa85ea39d5156285fa10ddc340a89c5d. --- .github/workflows/test-unit.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 3897a0df3e8..aa82a0bf3ee 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -51,7 +51,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/unit" + test-path: "tests/mcp_tests" workers: 2 reruns: 0 timeout-minutes: 20 From e1d2789d29650fc95533eae7b7b5a824ec5e87c4 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:35:23 +0000 Subject: [PATCH 042/127] test(llms): migrate phase 6 provider unit tests to tests/unit Migrate 18 provider test files from tests/test_litellm/llms to tests/unit/llms. 194 kept tests move as-is after mutation testing; 1 test deleted (test_completion_datarobot_with_environment_variables, env-gated no-assert); the fixture-only fal_ai cost calculator file is removed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/fal_ai/test_cost_calculator.py | 19 --------------- .../llms/chatgpt/chat/test_streaming_utils.py | 0 .../test_chatgpt_responses_transformation.py | 9 +++++++ .../test_cloudflare_transformation.py | 0 .../cohere/chat/test_cohere_transformation.py | 0 .../cohere/embed/test_v1_transformation.py | 0 .../ocr/test_cohere_parse_transformation.py | 0 .../rerank/test_rerank_guardrail_handler.py | 0 .../llms/crusoe/test_crusoe.py | 0 .../test_databricks_chat_transformation.py | 0 ...est_databricks_responses_transformation.py | 0 .../test_datarobot_chat_transformation.py | 0 .../llms/datarobot/test_datarobot.py | 24 ------------------- .../chat/test_deepseek_chat_transformation.py | 0 ...pseek_anthropic_messages_transformation.py | 0 .../deepseek/test_deepseek_cost_calculator.py | 10 ++++++++ ...docker_model_runner_chat_transformation.py | 0 ...levenlabs_text_to_speech_transformation.py | 0 .../fastcrw/search/test_transformation.py | 0 19 files changed, 19 insertions(+), 43 deletions(-) delete mode 100644 tests/test_litellm/llms/fal_ai/test_cost_calculator.py rename tests/{test_litellm => unit}/llms/chatgpt/chat/test_streaming_utils.py (100%) rename tests/{test_litellm => unit}/llms/chatgpt/responses/test_chatgpt_responses_transformation.py (97%) rename tests/{test_litellm => unit}/llms/cloudflare/test_cloudflare_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/chat/test_cohere_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/embed/test_v1_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/ocr/test_cohere_parse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cohere/rerank/test_rerank_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/crusoe/test_crusoe.py (100%) rename tests/{test_litellm => unit}/llms/databricks/chat/test_databricks_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/databricks/responses/test_databricks_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/datarobot/chat/test_datarobot_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/datarobot/test_datarobot.py (75%) rename tests/{test_litellm => unit}/llms/deepseek/chat/test_deepseek_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepseek/test_deepseek_cost_calculator.py (91%) rename tests/{test_litellm => unit}/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fastcrw/search/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py deleted file mode 100644 index 419aff42059..00000000000 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ /dev/null @@ -1,19 +0,0 @@ -import pytest - -import litellm -from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils -from litellm.llms.fal_ai.cost_calculator import cost_calculator -from litellm.types.utils import ImageObject, ImageResponse - - -@pytest.fixture(autouse=True) -def _use_local_model_cost_map(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - yield - litellm.get_model_info.cache_clear() - - -def _image_response(num_images: int = 1) -> ImageResponse: - return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)]) diff --git a/tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py b/tests/unit/llms/chatgpt/chat/test_streaming_utils.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py rename to tests/unit/llms/chatgpt/chat/test_streaming_utils.py diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py similarity index 97% rename from tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py rename to tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 9bf3eec61f9..c01ec312796 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -19,6 +19,15 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class TestChatGPTResponsesAPITransformation: @pytest.mark.parametrize( "model_name", diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py similarity index 100% rename from tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py rename to tests/unit/llms/cloudflare/test_cloudflare_transformation.py diff --git a/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py rename to tests/unit/llms/cohere/chat/test_cohere_transformation.py diff --git a/tests/test_litellm/llms/cohere/embed/test_v1_transformation.py b/tests/unit/llms/cohere/embed/test_v1_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/embed/test_v1_transformation.py rename to tests/unit/llms/cohere/embed/test_v1_transformation.py diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py b/tests/unit/llms/cohere/ocr/test_cohere_parse_transformation.py similarity index 100% rename from tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py rename to tests/unit/llms/cohere/ocr/test_cohere_parse_transformation.py diff --git a/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py b/tests/unit/llms/cohere/rerank/test_rerank_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py rename to tests/unit/llms/cohere/rerank/test_rerank_guardrail_handler.py diff --git a/tests/test_litellm/llms/crusoe/test_crusoe.py b/tests/unit/llms/crusoe/test_crusoe.py similarity index 100% rename from tests/test_litellm/llms/crusoe/test_crusoe.py rename to tests/unit/llms/crusoe/test_crusoe.py diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py rename to tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py diff --git a/tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py b/tests/unit/llms/databricks/responses/test_databricks_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py rename to tests/unit/llms/databricks/responses/test_databricks_responses_transformation.py diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/unit/llms/datarobot/chat/test_datarobot_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py rename to tests/unit/llms/datarobot/chat/test_datarobot_chat_transformation.py diff --git a/tests/test_litellm/llms/datarobot/test_datarobot.py b/tests/unit/llms/datarobot/test_datarobot.py similarity index 75% rename from tests/test_litellm/llms/datarobot/test_datarobot.py rename to tests/unit/llms/datarobot/test_datarobot.py index d9f42960601..c98faf0151e 100644 --- a/tests/test_litellm/llms/datarobot/test_datarobot.py +++ b/tests/unit/llms/datarobot/test_datarobot.py @@ -78,27 +78,3 @@ def test_completion_datarobot_with_deployment(): except Exception as e: pytest.fail(f"Error occurred: {e}") - -def test_completion_datarobot_with_environment_variables(): - """Allow the test to run with environment variables if they are set for integrations.""" - # If keys are not set, the test will be skipped - if os.environ.get("DATAROBOT_API_TOKEN") is None: - return - - messages = [ - {"role": "user", "content": "What's the weather like in San Francisco?"} - ] - try: - response = completion( - model="datarobot/vertex_ai/gemini-1.5-flash-002", - messages=messages, - max_tokens=5, - clientId="custom-model", - ) - print(response) - assert response["object"] == "chat.completion" - assert response["model"] == "gemini-1.5-flash-002" - assert len(response["choices"]) == 1 - assert len(response["choices"][0]["message"]["content"]) > 0 - except Exception as e: - pytest.fail(f"Error occurred: {e}") diff --git a/tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py rename to tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py diff --git a/tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py b/tests/unit/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py rename to tests/unit/llms/deepseek/messages/test_deepseek_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py similarity index 91% rename from tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py rename to tests/unit/llms/deepseek/test_deepseek_cost_calculator.py index c3a4cdad0ac..cde4a8a4244 100644 --- a/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py +++ b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py @@ -7,6 +7,16 @@ import litellm from litellm._internal_context import pinned_billing_time from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + PEAK_MOMENTS: Final = ( pytest.param(datetime(2026, 9, 22, 8, 0, tzinfo=timezone.utc), id="tuesday-08:00"), pytest.param(datetime(2026, 9, 25, 9, 59, tzinfo=timezone.utc), id="friday-09:59"), diff --git a/tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py b/tests/unit/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py rename to tests/unit/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py diff --git a/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py rename to tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/fastcrw/search/test_transformation.py b/tests/unit/llms/fastcrw/search/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/fastcrw/search/test_transformation.py rename to tests/unit/llms/fastcrw/search/test_transformation.py From 4048062e537329784c8197fc9b16ba77a08b7611 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:48:07 +0000 Subject: [PATCH 043/127] test(llms): annotate local_model_cost_map fixtures in phase 6 tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chatgpt/responses/test_chatgpt_responses_transformation.py | 3 ++- tests/unit/llms/deepseek/test_deepseek_cost_calculator.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index c01ec312796..0b04dd0ed78 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -5,6 +5,7 @@ Source: litellm/llms/chatgpt/responses/transformation.py """ import json +from collections.abc import Generator from unittest.mock import MagicMock, patch import httpx @@ -20,7 +21,7 @@ from litellm.utils import ProviderConfigManager @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) litellm.get_model_info.cache_clear() diff --git a/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py index cde4a8a4244..e61c15c3746 100644 --- a/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py +++ b/tests/unit/llms/deepseek/test_deepseek_cost_calculator.py @@ -1,3 +1,4 @@ +from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -9,7 +10,7 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) litellm.get_model_info.cache_clear() From b450baa402527c5be47a532508d7708fd1ca9a41 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:47:33 +0000 Subject: [PATCH 044/127] test(llms): migrate phase 5 provider unit tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_bearer_token.py | 158 ------------------ .../test_amazon_nova_canvas_transformation.py | 0 .../test_amazon_stability3_transformation.py | 0 .../image/test_bedrock_image_bearer_token.py | 21 +++ .../test_bedrock_image_prepare_request.py | 0 .../test_amazon_nova_canvas_image_edit.py | 0 .../test_bedrock_agent_transformation.py | 0 .../guardrail_translation/test_handler.py | 0 ...test_bedrock_passthrough_transformation.py | 2 - .../realtime/test_bedrock_realtime_handler.py | 0 .../test_bedrock_realtime_transformation.py | 0 .../test_bedrock_rerank_header_forwarding.py | 0 ...est_bedrock_vector_store_transformation.py | 0 ...drock_mantle_passthrough_transformation.py | 0 .../test_bfl_image_edit_transformation.py | 0 ...est_bfl_image_generation_transformation.py | 0 .../test_bfl_common_utils.py | 0 .../chat/test_bytez_chat_transformation.py | 0 .../test_cerebras_chat_transformation.py | 0 .../llms/chat/test_converse_handler.py | 0 .../chatgpt/test_chatgpt_authenticator.py | 0 21 files changed, 21 insertions(+), 160 deletions(-) delete mode 100644 tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_nova_canvas_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_stability3_transformation.py (100%) create mode 100644 tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_bedrock_image_prepare_request.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/guardrail_translation/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/test_bfl_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/bytez/chat/test_bytez_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cerebras/test_cerebras_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/chat/test_converse_handler.py (100%) rename tests/{test_litellm => unit}/llms/chatgpt/test_chatgpt_authenticator.py (100%) diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py deleted file mode 100644 index 0b11a66c100..00000000000 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ /dev/null @@ -1,158 +0,0 @@ -import json -import os -from unittest.mock import Mock, patch -import pytest - - -import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler - -# Mock response for Bedrock image generation -mock_image_response = {"images": ["base64_encoded_image_data"], "error": None} - - -class TestBedrockImageGeneration: - def test_image_generation_with_api_key_bearer_token(self): - """Test image generation with bearer token authentication""" - test_api_key = "test-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - # Setup mock response - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_env_variable_bearer_token(self, monkeypatch): - """Test image generation with bearer token from environment variable""" - test_api_key = "env-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - # Mock the environment variable - with ( - patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}), - patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen, - ): - - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - @pytest.mark.asyncio - async def test_async_image_generation_with_bearer_token(self): - """Test async image generation with bearer token authentication""" - test_api_key = "async-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.async_image_generation" - ) as mock_async_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_async_bedrock_image_gen.return_value = mock_image_response_obj - - # Call async image generation with api_key parameter - response = await litellm.aimage_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_async_bedrock_image_gen.assert_called_once() - for call in mock_async_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_sigv4(self): - """Test image generation falls back to SigV4 auth when no bearer token is provided""" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - mock_bedrock_image_gen.assert_called_once() - - -def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): - """The deployment's AWS profile does not exist, so resolving SigV4 credentials - raises; a bearer-token deployment must still sign the request with the - bearer token alone.""" - from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration - - monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") - - request = BedrockImageGeneration()._prepare_request( - model="amazon.nova-canvas-v1:0", - prompt="A cute baby sea otter", - optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, - api_base=None, - extra_headers=None, - api_key=None, - logging_obj=Mock(), - ) - - assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py new file mode 100644 index 00000000000..599507da03d --- /dev/null +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -0,0 +1,21 @@ +from unittest.mock import Mock + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py rename to tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py rename to tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py rename to tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py rename to tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py rename to tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index f2a9af11af7..854ef92fa4b 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -367,8 +367,6 @@ def test_bedrock_passthrough_region_extraction_from_inference_profile_arn(): assert ( "us-west-2" in api_base ), f"Expected region 'us-west-2' from ARN in base URL, but got: {api_base}" - - def test_bedrock_passthrough_model_id_arn_encoding(): """ Test that model_id ARNs are properly URL-encoded when used in endpoints. diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py rename to tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py rename to tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py rename to tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py rename to tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py rename to tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py b/tests/unit/llms/black_forest_labs/test_bfl_common_utils.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py rename to tests/unit/llms/black_forest_labs/test_bfl_common_utils.py diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py rename to tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/unit/llms/cerebras/test_cerebras_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py rename to tests/unit/llms/cerebras/test_cerebras_chat_transformation.py diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py similarity index 100% rename from tests/test_litellm/llms/chat/test_converse_handler.py rename to tests/unit/llms/chat/test_converse_handler.py diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py rename to tests/unit/llms/chatgpt/test_chatgpt_authenticator.py From 522c3e3ed6a5ce39b742eed3cbea68303683aeda Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:45:29 +0000 Subject: [PATCH 045/127] test: migrate wave 1 phase 2 legacy unit tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_amazon_nova_chat_completion.py | 195 ------------------ .../integrations/levo/test_levo.py | 42 ---- .../test_litellm_agent_model_resolver.py | 0 .../test_mavvrik_focus_logger.py | 0 .../integrations/opik/test_opik_extractors.py | 0 .../integrations/pointfive/test_logger.py | 0 .../integrations/pointfive/test_payload.py | 0 .../pointfive/test_upload_client.py | 0 .../test_vector_store_pre_call_hook.py | 0 .../audio_utils/test_subtitle_utils.py | 0 .../test_convert_dict_to_response.py | 0 .../test_convert_to_streaming_response.py | 0 .../test_get_formatted_prompt.py | 0 .../test_response_metadata.py | 0 .../test_a2a_guardrail_handler.py | 0 .../chat/test_a2a_chat_streaming_iterator.py | 0 .../a2a/chat/test_a2a_chat_transformation.py | 0 .../llms/a2a/test_common_utils.py | 0 .../llms/anthropic/batches/test_handler.py | 0 .../anthropic/batches/test_transformation.py | 0 20 files changed, 237 deletions(-) delete mode 100644 tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py rename tests/{test_litellm => unit}/integrations/levo/test_levo.py (88%) rename tests/{test_litellm => unit}/integrations/litellm_agent/test_litellm_agent_model_resolver.py (100%) rename tests/{test_litellm => unit}/integrations/mavvrik_focus/test_mavvrik_focus_logger.py (100%) rename tests/{test_litellm => unit}/integrations/opik/test_opik_extractors.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_logger.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_payload.py (100%) rename tests/{test_litellm => unit}/integrations/pointfive/test_upload_client.py (100%) rename tests/{test_litellm => unit}/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/audio_utils/test_subtitle_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_response_metadata.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/test_a2a_chat_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/llms/a2a/chat/test_a2a_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/a2a/test_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/batches/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/batches/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py b/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py deleted file mode 100644 index ecdd1b36333..00000000000 --- a/tests/test_litellm/llms/amazon_nova/chat/test_amazon_nova_chat_completion.py +++ /dev/null @@ -1,195 +0,0 @@ -import os -import pytest - -# Ensure the project root is on the import path - -from litellm import completion -from litellm.types.utils import ModelResponse, Usage, Choices, Message - - -def _has_api_key() -> bool: - """Check if Amazon Nova API key is available""" - return ( - "AMAZON_NOVA_API_KEY" in os.environ - and os.environ["AMAZON_NOVA_API_KEY"] is not None - ) - - -def _create_mock_nova_response(): - """Helper function to create mock Amazon Nova response for testing""" - return ModelResponse( - id="chatcmpl-test-nova-micro", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="I am Amazon Nova Micro. 777 times 9 equals 6993.", - role="assistant", - ), - ) - ], - created=1234567890, - model="amazon-nova/nova-micro-v1", - object="chat.completion", - usage=Usage(prompt_tokens=25, completion_tokens=15, total_tokens=40), - ) - - -def test_amazon_nova_chat_completion_nova_micro(): - if _has_api_key(): - response: ModelResponse = completion( - model="amazon-nova/nova-micro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Can you calculate 777 times 9?", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - else: - # Use mock response when API key is not available - response = _create_mock_nova_response() - # Additional mock-specific assertions for code review reference - assert ( - response.choices[0].message.content - == "I am Amazon Nova Micro. 777 times 9 equals 6993." - ) - assert response.model == "amazon-nova/nova-micro-v1" - assert response.usage.prompt_tokens == 25 - assert response.usage.completion_tokens == 15 - assert response.object == "chat.completion" - assert response.choices[0].finish_reason == "stop" - assert response.choices[0].message.role == "assistant" - - # Common assertions for both real and mock responses - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_lite(): - response: ModelResponse = completion( - model="amazon-nova/nova-lite-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Please tell me a poem on rain", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_pro(): - response: ModelResponse = completion( - model="amazon-nova/nova-pro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? What is MCP server and how does that help in building GenAI applications?", - }, - ], - timeout=30, - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_nova_premier(): - response: ModelResponse = completion( - model="amazon-nova/nova-premier-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What model are you? Can you help me understand what Trigonometry is?", - }, - ], - timeout=60, - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - print(response.choices[0].message.content) - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message.content is not None - assert response.usage.total_tokens > 0 - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_with_tool_usage(): - response: ModelResponse = completion( - model="amazon-nova/nova-micro-v1", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "What is the temperature in SFO?"}, - ], - tools=[ - { - "type": "function", - "function": { - "name": "getCurrentWeather", - "description": "Get the current weather in a given city", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "City and country e.g. Bogotá, Colombia", - } - }, - "required": ["location"], - }, - }, - } - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert response.choices[0].message is not None - - -@pytest.mark.skipif(not _has_api_key(), reason="Amazon Nova API key not available") -def test_amazon_nova_chat_completion_with_stream_response(): - response = completion( - model="amazon-nova/nova-micro-v1", - stream=True, - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - { - "role": "user", - "content": "What are MMO games? Can you give me some sample references?", - }, - ], - api_key=os.environ["AMAZON_NOVA_API_KEY"], - ) - - assert response is not None - chunks = list(response) - assert chunks is not None - assert len(chunks) > 0 diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/unit/integrations/levo/test_levo.py similarity index 88% rename from tests/test_litellm/integrations/levo/test_levo.py rename to tests/unit/integrations/levo/test_levo.py index 903be644671..647bcb3154e 100644 --- a/tests/test_litellm/integrations/levo/test_levo.py +++ b/tests/unit/integrations/levo/test_levo.py @@ -151,48 +151,6 @@ class TestLevoConfig(unittest.TestCase): class TestLevoIntegration(unittest.TestCase): """Integration tests for LevoLogger.""" - @patch.dict( - "os.environ", - { - "LEVOAI_API_KEY": "test-api-key", - "LEVOAI_ORG_ID": "test-org-id", - "LEVOAI_WORKSPACE_ID": "test-workspace-id", - "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", - }, - ) - @pytest.mark.skipif( - not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed" - ) - @patch( - "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy" - ) - @pytest.mark.asyncio - async def test_levo_logger_health_check_healthy(self, mock_init_proxy): - """Test health check returns healthy status when config is valid.""" - # Mock the proxy initialization to avoid importing proxy code - mock_init_proxy.return_value = None - - config = LevoLogger.get_levo_config() - otel_config = OpenTelemetryConfig( - exporter=config.protocol, - endpoint=config.endpoint, - headers=config.otlp_auth_headers, - ) - - # Create tracer provider with in-memory exporter - tracer_provider = TracerProvider() - tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) - - levo_logger = LevoLogger( - config=otel_config, callback_name="levo", tracer_provider=tracer_provider - ) - - # Run health check - result = await levo_logger.async_health_check() - - self.assertEqual(result["status"], "healthy") - self.assertIn("message", result) - @patch.dict("os.environ", {}, clear=True) def test_levo_logger_health_check_unhealthy(self): """Test health check returns unhealthy status when required vars are missing.""" diff --git a/tests/test_litellm/integrations/litellm_agent/test_litellm_agent_model_resolver.py b/tests/unit/integrations/litellm_agent/test_litellm_agent_model_resolver.py similarity index 100% rename from tests/test_litellm/integrations/litellm_agent/test_litellm_agent_model_resolver.py rename to tests/unit/integrations/litellm_agent/test_litellm_agent_model_resolver.py diff --git a/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py b/tests/unit/integrations/mavvrik_focus/test_mavvrik_focus_logger.py similarity index 100% rename from tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py rename to tests/unit/integrations/mavvrik_focus/test_mavvrik_focus_logger.py diff --git a/tests/test_litellm/integrations/opik/test_opik_extractors.py b/tests/unit/integrations/opik/test_opik_extractors.py similarity index 100% rename from tests/test_litellm/integrations/opik/test_opik_extractors.py rename to tests/unit/integrations/opik/test_opik_extractors.py diff --git a/tests/test_litellm/integrations/pointfive/test_logger.py b/tests/unit/integrations/pointfive/test_logger.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_logger.py rename to tests/unit/integrations/pointfive/test_logger.py diff --git a/tests/test_litellm/integrations/pointfive/test_payload.py b/tests/unit/integrations/pointfive/test_payload.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_payload.py rename to tests/unit/integrations/pointfive/test_payload.py diff --git a/tests/test_litellm/integrations/pointfive/test_upload_client.py b/tests/unit/integrations/pointfive/test_upload_client.py similarity index 100% rename from tests/test_litellm/integrations/pointfive/test_upload_client.py rename to tests/unit/integrations/pointfive/test_upload_client.py diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py similarity index 100% rename from tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py rename to tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py b/tests/unit/litellm_core_utils/audio_utils/test_subtitle_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py rename to tests/unit/litellm_core_utils/audio_utils/test_subtitle_utils.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_get_formatted_prompt.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py rename to tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_streaming_iterator.py b/tests/unit/llms/a2a/chat/test_a2a_chat_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/test_a2a_chat_streaming_iterator.py rename to tests/unit/llms/a2a/chat/test_a2a_chat_streaming_iterator.py diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py rename to tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py diff --git a/tests/test_litellm/llms/a2a/test_common_utils.py b/tests/unit/llms/a2a/test_common_utils.py similarity index 100% rename from tests/test_litellm/llms/a2a/test_common_utils.py rename to tests/unit/llms/a2a/test_common_utils.py diff --git a/tests/test_litellm/llms/anthropic/batches/test_handler.py b/tests/unit/llms/anthropic/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/test_handler.py rename to tests/unit/llms/anthropic/batches/test_handler.py diff --git a/tests/test_litellm/llms/anthropic/batches/test_transformation.py b/tests/unit/llms/anthropic/batches/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/test_transformation.py rename to tests/unit/llms/anthropic/batches/test_transformation.py From 02ccdbae906dfc21edb29421d799692c62d7054d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:50:50 +0000 Subject: [PATCH 046/127] test(llms): migrate bedrock, baseten and base_llm batch tests to tests/unit --- .../files/test_bedrock_files_integration.py | 115 ------------------ .../base_llm/batches/test_transformation.py | 16 --- .../realtime/test_transcription_protocol.py | 0 .../baseten/chat/test_baseten_completions.py | 0 .../test_agentcore_transformation.py | 0 .../test_amazon_moonshot_transformation.py | 0 .../test_amazon_nova_transformation.py | 22 ++++ .../test_amazon_qwen2_transformation.py | 0 .../test_amazon_qwen3_transformation.py | 0 .../test_base_invoke_transformation.py | 0 ...ations_anthropic_claude3_transformation.py | 76 ++++++++++++ .../test_twelvelabs_pegasus_transformation.py | 0 ...test_bedrock_chat_mantle_transformation.py | 43 +++++++ .../test_bedrock_count_tokens_handler.py | 0 ...est_bedrock_count_tokens_transformation.py | 0 .../expected_bedrock_batch_completions.jsonl | 0 .../expected_bedrock_batch_embeddings.jsonl | 0 .../files/input_batch_completions.jsonl | 0 .../files/input_batch_embeddings.jsonl | 0 .../files/test_bedrock_files_handler.py | 0 .../test_bedrock_files_transformation.py | 0 21 files changed, 141 insertions(+), 131 deletions(-) delete mode 100644 tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py rename tests/{test_litellm => unit}/llms/base_llm/batches/test_transformation.py (92%) rename tests/{test_litellm => unit}/llms/base_llm/realtime/test_transcription_protocol.py (100%) rename tests/{test_litellm => unit}/llms/baseten/chat/test_baseten_completions.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/agentcore/test_agentcore_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py (85%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py (92%) rename tests/{test_litellm => unit}/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py (68%) rename tests/{test_litellm => unit}/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/expected_bedrock_batch_completions.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/input_batch_completions.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/input_batch_embeddings.jsonl (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/test_bedrock_files_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/files/test_bedrock_files_transformation.py (100%) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py deleted file mode 100644 index 6d37d43b028..00000000000 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Test Bedrock files integration with main files API -""" - -import base64 -from unittest.mock import MagicMock, patch - -import pytest - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent -from litellm.types.utils import SpecialEnums - - -class TestBedrockFilesIntegration: - """Test integration of Bedrock files with main litellm API""" - - @pytest.mark.asyncio - async def test_litellm_afile_content_bedrock_provider_with_s3_uri(self): - """Test litellm.afile_content with bedrock provider using direct S3 URI""" - file_id = "s3://test-bucket/test-file.jsonl" - expected_content = ( - b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - ) - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="s3://test-bucket/test-file.jsonl"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content since the code - # now routes through ProviderConfigManager -> base_llm_http_handler - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call litellm.afile_content - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called with correct parameters - mock_retrieve.assert_called_once() - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["_is_async"] is True - assert call_kwargs["file_content_request"]["file_id"] == file_id - - @pytest.mark.asyncio - async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self): - """Test litellm.afile_content with bedrock provider using unified file ID""" - # Create a unified file ID - s3_uri = "s3://test-bucket/batch-outputs/output.jsonl" - unified_id = "test-unified-id-123" - model_id = "test-model-id-456" - - unified_file_id_str = f"litellm_proxy:application/json;unified_id,{unified_id};target_model_names,;llm_output_file_id,{s3_uri};llm_output_file_model_id,{model_id}" - encoded_file_id = ( - base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") - ) - - expected_content = ( - b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - ) - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=s3_uri), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call litellm.afile_content with unified file ID - result = await litellm.afile_content( - file_id=encoded_file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["_is_async"] is True - # The handler passes the encoded file_id as-is - assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id diff --git a/tests/test_litellm/llms/base_llm/batches/test_transformation.py b/tests/unit/llms/base_llm/batches/test_transformation.py similarity index 92% rename from tests/test_litellm/llms/base_llm/batches/test_transformation.py rename to tests/unit/llms/base_llm/batches/test_transformation.py index d84c820228f..0c360ce2ed9 100644 --- a/tests/test_litellm/llms/base_llm/batches/test_transformation.py +++ b/tests/unit/llms/base_llm/batches/test_transformation.py @@ -129,22 +129,6 @@ def test_subclass_missing_any_abstract_member_cannot_instantiate(missing_member) Incomplete() -def test_concrete_instance_methods_run(): - """Sanity: the trivial overrides actually execute through the base contract.""" - instance = _ConcreteBatchesConfig() - assert instance.custom_llm_provider == LlmProviders.OPENAI - assert instance.validate_environment( - headers={"x": "1"}, - model="m", - messages=[], - optional_params={}, - litellm_params={}, - ) == {"x": "1"} - assert instance.transform_retrieve_batch_request( - batch_id="b-1", optional_params={}, litellm_params={} - ) == {"batch_id": "b-1"} - - # =========================================================================== # # get_config() # =========================================================================== # diff --git a/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py b/tests/unit/llms/base_llm/realtime/test_transcription_protocol.py similarity index 100% rename from tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py rename to tests/unit/llms/base_llm/realtime/test_transcription_protocol.py diff --git a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py similarity index 100% rename from tests/test_litellm/llms/baseten/chat/test_baseten_completions.py rename to tests/unit/llms/baseten/chat/test_baseten_completions.py diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py rename to tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py similarity index 85% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py index 6c370344ae7..f0f0f9160fb 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py @@ -1,5 +1,8 @@ import json +import pytest + +import litellm from litellm.llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import ( AmazonInvokeNovaConfig, ) @@ -13,6 +16,25 @@ TOOL_CALL = {"id": "call_1", "type": "function", "function": {"name": "f", "argu PNG_DATA_URL = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _transform_request(messages, optional_params, litellm_params=None): return AmazonInvokeNovaConfig().transform_request( model=MODEL, diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py similarity index 92% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 84db0733227..2c74d23a6a2 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,6 +1,8 @@ import asyncio +import base64 import json import uuid +from types import SimpleNamespace from typing import Final from unittest.mock import patch @@ -17,6 +19,80 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.fixture +def local_beta_headers_config(monkeypatch): + """Pin the bundled ``anthropic_beta_headers_config.json`` so beta header assertions + do not depend on the network-fetched copy or on what earlier tests left cached.""" + from litellm.anthropic_beta_headers_manager import reload_beta_headers_config + + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + reload_beta_headers_config() + try: + yield + finally: + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + reload_beta_headers_config() + + def test_get_supported_params_thinking(): config = AmazonAnthropicClaudeConfig() params = config.get_supported_openai_params( diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py similarity index 68% rename from tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py rename to tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index a8448f5fa7a..cb892b1ea11 100644 --- a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -1,12 +1,55 @@ +import base64 import json import uuid +from types import SimpleNamespace import httpx +import pytest import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): image_url = f"http://img.example/{uuid.uuid4()}.png" captured = {} diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py rename to tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py rename to tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_completions.jsonl b/tests/unit/llms/bedrock/files/expected_bedrock_batch_completions.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_completions.jsonl rename to tests/unit/llms/bedrock/files/expected_bedrock_batch_completions.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl b/tests/unit/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl rename to tests/unit/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/input_batch_completions.jsonl b/tests/unit/llms/bedrock/files/input_batch_completions.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/input_batch_completions.jsonl rename to tests/unit/llms/bedrock/files/input_batch_completions.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl b/tests/unit/llms/bedrock/files/input_batch_embeddings.jsonl similarity index 100% rename from tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl rename to tests/unit/llms/bedrock/files/input_batch_embeddings.jsonl diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py rename to tests/unit/llms/bedrock/files/test_bedrock_files_handler.py diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py rename to tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py From 4defed7f2e7eaacdf8e130857eddc6018c7bc3f7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:03:54 +0000 Subject: [PATCH 047/127] test: migrate wave 1 phase 8 legacy llm tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_hosted_vllm_ssl_verify.py | 147 ------------------ .../test_hosted_vllm_embedding_ssl_verify.py | 135 ---------------- ..._github_copilot_messages_transformation.py | 7 - ...github_copilot_responses_transformation.py | 116 +++++--------- .../test_gradient_ai_chat_transformation.py | 0 .../chat/test_groq_chat_transformation.py | 2 - .../llms/groq/test_groq_cost_calculator.py | 0 .../test_hosted_vllm_chat_transformation.py | 71 +-------- ...st_hosted_vllm_embedding_transformation.py | 8 +- ...t_hosted_vllm_image_edit_transformation.py | 0 .../responses/test_hosted_vllm_responses.py | 9 +- .../test_hosted_vllm_rerank_transformation.py | 0 .../test_hosted_vllm_video_transformation.py | 0 .../test_huggingface_rerank_transformation.py | 40 +---- .../test_inception_chat_transformation.py | 18 +-- ...est_inception_completion_transformation.py | 18 +-- .../test_jina_embedding_transformation.py | 0 .../chat/test_langflow_chat_transformation.py | 31 +--- .../litellm_proxy/test_sandbox_executor.py | 25 +-- .../litellm_proxy/test_skills_ownership.py | 73 ++------- 20 files changed, 84 insertions(+), 616 deletions(-) delete mode 100644 tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py delete mode 100644 tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py rename tests/{test_litellm => unit}/llms/github_copilot/messages/test_github_copilot_messages_transformation.py (98%) rename tests/{test_litellm => unit}/llms/github_copilot/responses/test_github_copilot_responses_transformation.py (89%) rename tests/{test_litellm => unit}/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/groq/chat/test_groq_chat_transformation.py (99%) rename tests/{test_litellm => unit}/llms/groq/test_groq_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py (82%) rename tests/{test_litellm => unit}/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py (97%) rename tests/{test_litellm => unit}/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/responses/test_hosted_vllm_responses.py (96%) rename tests/{test_litellm => unit}/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/huggingface/rerank/test_huggingface_rerank_transformation.py (91%) rename tests/{test_litellm => unit}/llms/inception/test_inception_chat_transformation.py (96%) rename tests/{test_litellm => unit}/llms/inception/test_inception_completion_transformation.py (95%) rename tests/{test_litellm => unit}/llms/jina_ai/embedding/test_jina_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/langflow/chat/test_langflow_chat_transformation.py (93%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_sandbox_executor.py (84%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_skills_ownership.py (88%) diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py deleted file mode 100644 index 2364468efe1..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the OpenAI catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMSSLVerify: - """Test suite for SSL verification in hosted_vllm provider.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.completion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.acompletion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py deleted file mode 100644 index de94da49384..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py +++ /dev/null @@ -1,135 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider embeddings. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider for embeddings. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the openai_like catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMEmbeddingSSLVerify: - """Test suite for SSL verification in hosted_vllm provider embeddings.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_embedding_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync embedding calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.embedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_embedding_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async embedding calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.aembedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py similarity index 98% rename from tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py rename to tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 8039e744f46..9e9760650cf 100644 --- a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -10,13 +10,6 @@ from litellm.llms.github_copilot.messages.transformation import ( ) -def test_github_copilot_anthropic_messages_config_init(): - """Test GithubCopilotAnthropicMessagesConfig initialization.""" - config = GithubCopilotAnthropicMessagesConfig() - assert config is not None - assert hasattr(config, "authenticator") - - def test_github_copilot_anthropic_messages_get_complete_url(): """get_complete_url builds the /v1/messages URL from the base it is handed. diff --git a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py similarity index 89% rename from tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py rename to tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py index 0174465b0cc..b8380b7adb4 100644 --- a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py +++ b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py @@ -26,9 +26,7 @@ def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): """Pin litellm.model_cost to the bundled local backup so tests don't depend on remote catalog fetches (and don't change behavior across remote refreshes).""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr( - litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url) - ) + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) litellm.add_known_models(model_cost_map=litellm.model_cost) @@ -44,49 +42,35 @@ class TestGithubCopilotResponsesAPITransformation: provider=LlmProviders.GITHUB_COPILOT, ) - assert ( - config is not None - ), "Config should not be None for GitHub Copilot provider" - assert isinstance( - config, GithubCopilotResponsesAPIConfig - ), f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.GITHUB_COPILOT - ), "custom_llm_provider should be GITHUB_COPILOT" + assert config is not None, "Config should not be None for GitHub Copilot provider" + assert isinstance(config, GithubCopilotResponsesAPIConfig), ( + f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" + ) + assert config.custom_llm_provider == LlmProviders.GITHUB_COPILOT, "custom_llm_provider should be GITHUB_COPILOT" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_github_copilot_responses_endpoint_url(self, mock_authenticator_class): """Test that get_complete_url returns correct GitHub Copilot endpoint""" # Mock authenticator to return default base mock_auth_instance = MagicMock() - mock_auth_instance.get_api_base.return_value = ( - "https://api.individual.githubcopilot.com" - ) + mock_auth_instance.get_api_base.return_value = "https://api.individual.githubcopilot.com" mock_authenticator_class.return_value = mock_auth_instance config = GithubCopilotResponsesAPIConfig() # Test with default GitHub Copilot API base (from authenticator) url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.individual.githubcopilot.com/responses" - ), f"Expected GitHub Copilot responses endpoint, got {url}" + assert url == "https://api.individual.githubcopilot.com/responses", ( + f"Expected GitHub Copilot responses endpoint, got {url}" + ) # Test with custom api_base (overrides authenticator) - custom_url = config.get_complete_url( - api_base="https://custom.githubcopilot.com", litellm_params={} - ) - assert ( - custom_url == "https://custom.githubcopilot.com/responses" - ), f"Expected custom endpoint, got {custom_url}" + custom_url = config.get_complete_url(api_base="https://custom.githubcopilot.com", litellm_params={}) + assert custom_url == "https://custom.githubcopilot.com/responses", f"Expected custom endpoint, got {custom_url}" # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.githubcopilot.com/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.githubcopilot.com/responses" - ), "Should handle trailing slash" + url_with_slash = config.get_complete_url(api_base="https://api.githubcopilot.com/", litellm_params={}) + assert url_with_slash == "https://api.githubcopilot.com/responses", "Should handle trailing slash" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_default_headers(self, mock_authenticator_class): @@ -98,9 +82,7 @@ class TestGithubCopilotResponsesAPITransformation: config = GithubCopilotResponsesAPIConfig() - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params={}) # Check required headers assert headers["Authorization"] == "Bearer test-api-key-123" @@ -127,9 +109,7 @@ class TestGithubCopilotResponsesAPITransformation: "custom-header": "custom-value", } - headers = config.validate_environment( - headers=custom_headers, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers=custom_headers, model="gpt-5.1-codex", litellm_params={}) # User header should override default assert headers["editor-version"] == "custom/2.0.0" @@ -182,9 +162,7 @@ class TestGithubCopilotResponsesAPITransformation: """Test _has_vision_input detects input_image type""" config = GithubCopilotResponsesAPIConfig() - input_with_vision = [ - {"role": "user", "content": [{"type": "input_image", "data": "base64..."}]} - ] + input_with_vision = [{"role": "user", "content": [{"type": "input_image", "data": "base64..."}]}] has_vision = config._has_vision_input(input_with_vision) assert has_vision is True, "Should detect input_image type" @@ -246,13 +224,11 @@ class TestGithubCopilotResponsesAPITransformation: } ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("copilot-vision-request") == "true" - ), "Should add copilot-vision-request header for vision input" + assert headers.get("copilot-vision-request") == "true", ( + "Should add copilot-vision-request header for vision input" + ) @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_with_x_initiator(self, mock_authenticator_class): @@ -270,21 +246,15 @@ class TestGithubCopilotResponsesAPITransformation: {"role": "assistant", "content": "Hi"}, ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("X-Initiator") == "agent" - ), "Should set X-Initiator to 'agent' for assistant role" + assert headers.get("X-Initiator") == "agent", "Should set X-Initiator to 'agent' for assistant role" def test_map_openai_params_no_transformation(self): """Test that map_openai_params passes through parameters unchanged""" config = GithubCopilotResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - temperature=0.7, max_output_tokens=1000, stream=False - ) + params = ResponsesAPIOptionalRequestParams(temperature=0.7, max_output_tokens=1000, stream=False) result = config.map_openai_params( response_api_optional_params=params, @@ -338,9 +308,9 @@ class TestGithubCopilotResponsesAPITransformation: result = config._handle_reasoning_item(reasoning_item) # encrypted_content should be preserved - assert ( - result.get("encrypted_content") == "encrypted-blob-abc123" - ), "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + assert result.get("encrypted_content") == "encrypted-blob-abc123", ( + "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + ) # status=None should be filtered out assert "status" not in result, "status=None should be filtered out" # content=None should be filtered out @@ -393,9 +363,7 @@ class TestGithubCopilotResponsesAPIRouting: in the (already-merged) model info; otherwise returns None so the dispatcher routes through the chat-completions translation bridge.""" - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_config_when_mode_is_responses(self, mock_get_info): """``mode=responses`` returns native config.""" mock_get_info.return_value = {"mode": "responses"} @@ -405,9 +373,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_chat(self, mock_get_info): """``mode=chat`` returns None so dispatcher uses bridge.""" mock_get_info.return_value = {"mode": "chat"} @@ -417,9 +383,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_unset_and_no_endpoints(self, mock_get_info): """Entry without ``mode`` and without ``supported_endpoints`` returns None (conservative default).""" @@ -499,9 +463,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_get_model_info_raises(self, mock_get_info): """Catalog lookup failure (model not registered) returns None (conservative default; bridge handles unknown models safely).""" @@ -512,9 +474,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_user_override_via_register_model(self, mock_get_info): """User-supplied per-deployment ``model_info`` flows through ``litellm.register_model`` (called by the router) into the merged @@ -528,9 +488,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_chat_only_entry_returns_none(self, mock_get_info): """Realistic ``model_prices_and_context_window.json`` shape for a chat-only Copilot model (e.g. github_copilot/gemini-3.1-pro-preview) @@ -554,9 +512,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_responses_only_entry_returns_config(self, mock_get_info): """Realistic catalog entry for a Responses-only Copilot model (e.g. github_copilot/gpt-5.5) returns the native config.""" @@ -592,9 +548,7 @@ class TestGithubCopilotReasoningStreamItemIdNormalization: output_index group to the id from its output_item.added.""" def _config(self): - with patch( - "litellm.llms.github_copilot.responses.transformation.Authenticator" - ): + with patch("litellm.llms.github_copilot.responses.transformation.Authenticator"): return GithubCopilotResponsesAPIConfig() def _transform(self, config, chunk): diff --git a/tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py b/tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py rename to tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py similarity index 99% rename from tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py rename to tests/unit/llms/groq/chat/test_groq_chat_transformation.py index f605958b979..f5a7a920124 100644 --- a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -202,5 +202,3 @@ class TestGroqWebSearchUsageSignal: model_response = litellm.ModelResponse() GroqChatConfig()._add_web_search_usage(model_response=model_response) assert getattr(model_response, "usage", None) is None - - diff --git a/tests/test_litellm/llms/groq/test_groq_cost_calculator.py b/tests/unit/llms/groq/test_groq_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/groq/test_groq_cost_calculator.py rename to tests/unit/llms/groq/test_groq_cost_calculator.py diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py similarity index 82% rename from tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py rename to tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index 82b05601a85..1cc6a1457fc 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -41,74 +41,9 @@ def test_hosted_vllm_chat_transformation_file_url(): ] -def test_hosted_vllm_chat_transformation_with_audio_url(): - from litellm import completion - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "llama-3.1-70b-instruct", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - mock_response.text = json.dumps(mock_response.json.return_value) - mock_client.post.return_value = mock_response - - with patch( - "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", - return_value=mock_client, - ): - try: - completion( - model="hosted_vllm/llama-3.1-70b-instruct", - messages=[ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - }, - ], - }, - ], - api_base="https://test-vllm.example.com/v1", - ) - except Exception: - pass - - mock_client.post.assert_called_once() - call_kwargs = mock_client.post.call_args[1] - request_data = json.loads(call_kwargs["data"]) - assert request_data["messages"] == [ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - } - ], - } - ] - - def test_hosted_vllm_supports_reasoning_effort(): config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/gpt-oss-120b" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/gpt-oss-120b") assert "reasoning_effort" in supported_params optional_params = config.map_openai_params( non_default_params={"reasoning_effort": "high"}, @@ -129,9 +64,7 @@ def test_hosted_vllm_supports_thinking(): Related issue: https://github.com/BerriAI/litellm/issues/19761 """ config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/GLM-4.6-FP8" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/GLM-4.6-FP8") assert "thinking" in supported_params # Test thinking below the low threshold -> "minimal" diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py similarity index 97% rename from tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py rename to tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py index 34be3e12abd..5854b1596b4 100644 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py +++ b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py @@ -87,9 +87,7 @@ class TestHostedVLLMEmbeddingTransformation: headers={}, ) - assert ( - "encoding_format" not in result - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in result, "encoding_format should not be in request when not provided" def test_encoding_format_not_included_when_none(self): """ @@ -278,9 +276,7 @@ class TestHostedVLLMEmbeddingTransformation: sent_data = json.loads(call_kwargs["data"]) # Assert that encoding_format is NOT in the sent data - assert ( - "encoding_format" not in sent_data - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in sent_data, "encoding_format should not be in request when not provided" assert sent_data["model"] == "BAAI/bge-small-en-v1.5" assert sent_data["input"] == ["Hello world"] diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py rename to tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py similarity index 96% rename from tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py rename to tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py index e81bf0c4f1f..55d0ce1e68e 100644 --- a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py +++ b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py @@ -68,9 +68,7 @@ def test_hosted_vllm_responses_create_with_string_input(): Test that hosted_vllm routes directly to the native /v1/responses endpoint when the Responses API config is registered, and correctly parses the response. """ - mock_client = _make_mock_http_client( - _make_mock_responses_api_response("I'm doing well, thanks!") - ) + mock_client = _make_mock_http_client(_make_mock_responses_api_response("I'm doing well, thanks!")) with patch( "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", @@ -109,10 +107,7 @@ def test_hosted_vllm_responses_create_with_explicit_none_extra_body(): ) # extra_body=None should be normalized to an empty dict (or absent) - assert ( - optional_params.get("extra_body") is not None - or "extra_body" not in optional_params - ) + assert optional_params.get("extra_body") is not None or "extra_body" not in optional_params def test_hosted_vllm_provider_config_registration(): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py rename to tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py b/tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py rename to tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py diff --git a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py similarity index 91% rename from tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py rename to tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 9d6b7290eb6..6fd2b006fef 100644 --- a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -219,29 +219,6 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_error_handling(mock_post): - """Test HuggingFace rerank error handling.""" - - def return_val(): - return {"error": "Unauthorized"} - - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.json = return_val - mock_response.text = "Unauthorized" - mock_post.return_value = mock_response - - with pytest.raises(litellm.APIConnectionError): - litellm.rerank( - model="huggingface/BAAI/bge-reranker-base", - query="hello", - documents=["hello", "world"], - top_n=2, - api_key="invalid_key", - ) - - def test_huggingface_rerank_config(): """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig @@ -249,10 +226,7 @@ def test_huggingface_rerank_config(): config = HuggingFaceRerankConfig() # Test complete URL generation - assert ( - config.get_complete_url(None, "test") - == "https://api-inference.huggingface.co/rerank" - ) + assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" # Test custom API base custom_url = config.get_complete_url("https://custom.huggingface.co", "test") @@ -292,13 +266,9 @@ def test_request_transformation(): config = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams( - query="hello", texts=["hello", "world"], top_n=2, return_text=True - ) + optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request( - model="test", optional_rerank_params=optional_params, headers={} - ) + request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) assert request_body["query"] == "hello" assert request_body["texts"] == ["hello", "world"] @@ -368,9 +338,7 @@ def test_validate_environment(): # Test headers override custom_headers = {"custom": "header"} - headers = config.validate_environment( - headers=custom_headers, model="test", api_key="test_key" - ) + headers = config.validate_environment(headers=custom_headers, model="test", api_key="test_key") assert "custom" in headers assert headers["custom"] == "header" diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/unit/llms/inception/test_inception_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/inception/test_inception_chat_transformation.py rename to tests/unit/llms/inception/test_inception_chat_transformation.py index 1d12be2adee..c4c023077fc 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/unit/llms/inception/test_inception_chat_transformation.py @@ -188,21 +188,15 @@ def test_inception_does_not_leak_key_to_caller_api_base(): caller also supplies their own key. """ config = InceptionChatConfig() - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True): with mock.patch.object(litellm, "inception_key", "module-secret"): # caller overrides api_base without a key -> server key withheld - api_base, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", None - ) + api_base, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", None) assert api_base == "https://attacker.example/v1" assert api_key is None # caller overrides api_base AND supplies their own key -> used as-is - _, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", "caller-key" - ) + _, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", "caller-key") assert api_key == "caller-key" # default/server base -> server-managed key resolved @@ -217,9 +211,7 @@ def test_get_llm_provider_inception(): assert model == "mercury-2" assert provider == "inception" - model, provider, _, api_base = get_llm_provider( - "mercury-2", api_base="https://api.inceptionlabs.ai/v1" - ) + model, provider, _, api_base = get_llm_provider("mercury-2", api_base="https://api.inceptionlabs.ai/v1") assert model == "mercury-2" assert provider == "inception" assert api_base == "https://api.inceptionlabs.ai/v1" @@ -293,5 +285,3 @@ def test_inception_completion_targets_inception_endpoint(): assert captured["body"]["model"] == "mercury-2" assert captured["body"]["tool_choice"] == "auto" assert response.choices[0].message.content == "hi" - - diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/unit/llms/inception/test_inception_completion_transformation.py similarity index 95% rename from tests/test_litellm/llms/inception/test_inception_completion_transformation.py rename to tests/unit/llms/inception/test_inception_completion_transformation.py index ed3f34fc744..84923229e20 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/unit/llms/inception/test_inception_completion_transformation.py @@ -22,9 +22,7 @@ def _fim_response_bytes(): "object": "text_completion", "created": 1, "model": "mercury-edit-2", - "choices": [ - {"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None} - ], + "choices": [{"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None}], "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, } ).encode() @@ -47,9 +45,7 @@ def test_inception_fim_supports_suffix_param(): def test_inception_fim_supported_params_match_schema(): """FIM exposes the OpenAI subset of Inception's FIMCompletionRequest only""" - params = InceptionTextCompletionConfig().get_supported_openai_params( - "mercury-edit-2" - ) + params = InceptionTextCompletionConfig().get_supported_openai_params("mercury-edit-2") for p in ("suffix", "top_p", "frequency_penalty", "presence_penalty", "stop"): assert p in params # Chat-only sampling controls are not part of Inception's FIM schema @@ -75,11 +71,7 @@ def test_inception_get_supported_openai_params_dispatch(): @pytest.mark.parametrize("provider", ["inception", "text-completion-inception"]) def test_inception_validate_environment(provider): - model = ( - "inception/mercury-2" - if provider == "inception" - else "text-completion-inception/mercury-edit-2" - ) + model = "inception/mercury-2" if provider == "inception" else "text-completion-inception/mercury-edit-2" with mock.patch.dict(os.environ, {}, clear=True): result = litellm.validate_environment(model) @@ -217,9 +209,7 @@ def test_inception_fim_does_not_leak_global_api_key(): content=_fim_response_bytes(), ) - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True): with mock.patch.object(litellm, "inception_key", None): with mock.patch.object(litellm, "api_key", "sk-global-should-not-leak"): with mock.patch("httpx.Client.send", new=fake_send): diff --git a/tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py b/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py rename to tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py similarity index 93% rename from tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py rename to tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py index 383a7afbe93..179a6cad4aa 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py @@ -46,7 +46,7 @@ def test_langflow_config_get_complete_url(): def test_langflow_config_get_complete_url_requires_api_base(): config = LangFlowConfig() - with pytest.raises(ValueError, match='api_base is required for LangFlow\\. Set it via'): + with pytest.raises(ValueError, match="api_base is required for LangFlow\\. Set it via"): config.get_complete_url( api_base=None, api_key=None, @@ -225,9 +225,7 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): posted_bodies.append(json.loads(body) if isinstance(body, str) else body) resp = MagicMock(spec=httpx.Response) resp.status_code = 200 - resp.json.return_value = { - "outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}] - } + resp.json.return_value = {"outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}]} resp.headers = {} resp.text = "{}" return resp @@ -275,9 +273,7 @@ def test_langflow_config_extract_response_from_outputs_dict(): "outputs": [ { "results": {}, - "outputs": { - "message": {"message": {"text": "via outputs dict"}} - }, + "outputs": {"message": {"message": {"text": "via outputs dict"}}}, } ] } @@ -292,14 +288,9 @@ def test_langflow_extract_response_returns_none_when_no_message(): assert config._extract_content_from_response({"outputs": []}) is None assert config._extract_content_from_response({"detail": "flow failed"}) is None assert config._extract_content_from_response({"outputs": ["not-a-dict"]}) is None + assert config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) is None assert ( - config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) - is None - ) - assert ( - config._extract_content_from_response( - {"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]} - ) + config._extract_content_from_response({"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]}) is None ) @@ -310,9 +301,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): status_code=200, json={ "session_id": "sess-abc", - "outputs": [ - {"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]} - ], + "outputs": [{"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]}], }, ) @@ -332,9 +321,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): assert result.choices[0].finish_reason == "stop" assert result.model == "langflow/my-flow-id" assert result.usage.completion_tokens > 0 - assert result.usage.total_tokens == ( - result.usage.prompt_tokens + result.usage.completion_tokens - ) + assert result.usage.total_tokens == (result.usage.prompt_tokens + result.usage.completion_tokens) def test_langflow_transform_response_raises_on_unparseable_body(): @@ -357,9 +344,7 @@ def test_langflow_transform_response_raises_on_unparseable_body(): def test_langflow_transform_response_raises_on_non_json_body(): config = LangFlowConfig() - raw_response = httpx.Response( - status_code=200, content=b"not json", headers={"content-type": "text/plain"} - ) + raw_response = httpx.Response(status_code=200, content=b"not json", headers={"content-type": "text/plain"}) with pytest.raises(LangFlowError): config.transform_response( diff --git a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py similarity index 84% rename from tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py rename to tests/unit/llms/litellm_proxy/test_sandbox_executor.py index 422e7a3cf4d..e7a03b9231a 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py +++ b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py @@ -55,9 +55,7 @@ def _install_fake_sandbox(monkeypatch, session_cls=_FakeSandboxSession): def test_execute_installs_inline_requirements_file(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) requirements = "git+https://example.com/repo.git#egg=foo\n-r extra.txt\n-e ./pkg\n" result = executor.execute( @@ -69,22 +67,15 @@ def test_execute_installs_inline_requirements_file(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - assert created_session.copied_contents[ - "/sandbox/.litellm_requirements.txt" - ] == requirements.encode("utf-8") - assert ( - "pip', 'install', '-r', '.litellm_requirements.txt'" - in created_session.run_calls[0] - ) + assert created_session.copied_contents["/sandbox/.litellm_requirements.txt"] == requirements.encode("utf-8") + assert "pip', 'install', '-r', '.litellm_requirements.txt'" in created_session.run_calls[0] assert "os.chdir('/sandbox')" in created_session.run_calls[1] def test_execute_uses_skill_requirements_txt(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", @@ -97,9 +88,7 @@ def test_execute_uses_skill_requirements_txt(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - copied_paths = { - sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls - } + copied_paths = {sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls} assert "/sandbox/requirements.txt" in copied_paths assert "/sandbox/.litellm_requirements.txt" not in copied_paths assert "pip', 'install', '-r', 'requirements.txt'" in created_session.run_calls[0] @@ -118,9 +107,7 @@ def test_execute_returns_install_failure(monkeypatch): _install_fake_sandbox(monkeypatch, session_cls=_FailingSandboxSession) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/unit/llms/litellm_proxy/test_skills_ownership.py similarity index 88% rename from tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py rename to tests/unit/llms/litellm_proxy/test_skills_ownership.py index e538c50cde8..6caa2da3169 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/unit/llms/litellm_proxy/test_skills_ownership.py @@ -37,12 +37,7 @@ def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: def test_should_extract_skill_auth_from_supported_metadata_fields(): auth = UserAPIKeyAuth(user_id="user-1") - assert ( - skills_main._get_user_api_key_auth_from_kwargs( - {"metadata": {"user_api_key_auth": auth}} - ) - is auth - ) + assert skills_main._get_user_api_key_auth_from_kwargs({"metadata": {"user_api_key_auth": auth}}) is auth assert ( skills_main._get_user_api_key_auth_from_kwargs( {"metadata": {}, "litellm_metadata": {"user_api_key_auth": auth}} @@ -122,9 +117,7 @@ def test_should_forward_skill_auth_through_sdk_entrypoints(monkeypatch): == "deleted" ) - assert handler.create_skill_handler.call_args.kwargs["metadata"] == { - "source": "request" - } + assert handler.create_skill_handler.call_args.kwargs["metadata"] == {"source": "request"} assert handler.create_skill_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.list_skills_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.get_skill_handler.call_args.kwargs["user_api_key_dict"] is auth @@ -149,9 +142,7 @@ def test_should_build_resource_owner_scopes_for_auth_context(): ] assert resource_ownership.get_primary_resource_owner_scope(auth) == "user-1" assert resource_ownership.user_can_access_resource_owner("team:team-1", auth) - assert resource_ownership.get_resource_owner_scopes( - UserAPIKeyAuth(token="token-hash") - ) == ["key:token-hash"] + assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth(token="token-hash")) == ["key:token-hash"] # Identity-less callers get an empty scope set — sharing a sentinel # would collapse every identity-less caller into the same logical # owner, which is a cross-tenant data-access primitive. @@ -165,9 +156,7 @@ def test_should_allow_admin_and_anonymous_resource_owner_paths(): assert resource_ownership.is_proxy_admin(admin) assert resource_ownership.user_can_access_resource_owner(None, admin) assert resource_ownership.user_can_access_resource_owner(None, None) - assert not resource_ownership.user_can_access_resource_owner( - None, UserAPIKeyAuth(user_id="user-1") - ) + assert not resource_ownership.user_can_access_resource_owner(None, UserAPIKeyAuth(user_id="user-1")) @pytest.mark.asyncio @@ -218,9 +207,7 @@ async def test_should_forward_skill_auth_through_transformation_handler(monkeypa async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -242,9 +229,7 @@ async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -268,9 +253,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc sentinel as ``created_by`` would let any two such callers see each other's skills via the resulting shared owner scope.""" table = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -291,9 +274,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch): table = AsyncMock() table.find_many.return_value = [_skill("litellm_skill_owner", "user-1")] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -318,9 +299,7 @@ async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypat async def test_should_hide_skill_from_different_owner(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_other", "user-2") - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -340,9 +319,7 @@ async def test_should_hide_skill_from_different_owner(monkeypatch): async def test_should_hide_unowned_skill_by_default(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_unowned", None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -364,9 +341,7 @@ async def test_list_skills_excludes_unowned_for_non_admin(monkeypatch): with ``created_by IS NULL`` are excluded — admin-only.""" table = AsyncMock() table.find_many.return_value = [] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -422,9 +397,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -432,10 +405,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): ) for _ in range(3): - assert ( - await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - is fake_skill - ) + assert await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") is fake_skill assert table.find_unique.await_count == 1 @@ -445,9 +415,7 @@ async def test_load_skill_caches_negative_lookups(monkeypatch): the DB and the caller still sees ``None``.""" table = AsyncMock() table.find_unique = AsyncMock(return_value=None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -466,9 +434,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) table.delete = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -480,12 +446,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") is fake_skill auth = UserAPIKeyAuth(user_id="user-1") - await skills_handler.LiteLLMSkillsHandler.delete_skill( - "litellm_skill_a", user_api_key_dict=auth - ) + await skills_handler.LiteLLMSkillsHandler.delete_skill("litellm_skill_a", user_api_key_dict=auth) # Post-delete, the cache holds the negative sentinel — not the stale row. - assert ( - skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") - == skills_handler._NEGATIVE_SKILL_SENTINEL - ) + assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") == skills_handler._NEGATIVE_SKILL_SENTINEL From 99c2ef4d73efbfa657da009ef0a6a70205b03073 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:07:50 +0000 Subject: [PATCH 048/127] test(unit): block external sockets at import time and add a socket policy regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++------ tests/unit/test_socket_policy.py | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_socket_policy.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..017e63ed1b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,9 +1,11 @@ -from collections.abc import Iterator +import os from typing import Final import pytest from pytest_socket import enable_socket, socket_allow_hosts +os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -11,13 +13,13 @@ def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) -@pytest.fixture(autouse=True, scope="session") -def block_external_sockets() -> Iterator[None]: - _allow_loopback_only() - yield - enable_socket() +_allow_loopback_only() @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() + + +def pytest_sessionfinish() -> None: + enable_socket() diff --git a/tests/unit/test_socket_policy.py b/tests/unit/test_socket_policy.py new file mode 100644 index 00000000000..f93794d1ba8 --- /dev/null +++ b/tests/unit/test_socket_policy.py @@ -0,0 +1,17 @@ +import socket + +import pytest +from pytest_socket import SocketConnectBlockedError + + +def test_external_connect_is_refused_before_a_packet_leaves() -> None: + with pytest.raises(SocketConnectBlockedError): + socket.create_connection(("192.0.2.1", 9), timeout=1) + + +def test_loopback_connect_is_allowed() -> None: + with socket.socket() as server: + server.bind(("127.0.0.1", 0)) + server.listen() + with socket.create_connection(server.getsockname(), timeout=1) as client: + assert client.getpeername() == server.getsockname() From d47008129c175a068b23c9e0afe306463e39974f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:08:12 +0000 Subject: [PATCH 049/127] test(llms): migrate phase 7 provider unit tests to tests/unit Move the wave 1 phase 7 batch (fireworks_ai, gemini, gigachat, github_copilot; 20 files) from tests/test_litellm to tests/unit after judging every test function under a behaviour mutation. Seven wiring or mock-echo tests that stayed green are deleted. The fireworks cost calculator tests get a local model_cost save/restore fixture since the tests/unit tree has no shared conftest for it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_fireworks_ai_chat_transformation.py | 0 ...test_fireworks_ai_rerank_transformation.py | 0 ...t_fireworks_ai_responses_transformation.py | 15 -------- .../test_fireworks_ai_cache_pricing.py | 0 .../test_fireworks_ai_common_utils.py | 0 .../test_fireworks_ai_cost_calculator.py | 22 ++++++++--- ...mini_audio_transcription_transformation.py | 0 .../files/test_gemini_files_transformation.py | 0 .../test_google_genai_guardrail_handler.py | 0 .../test_gemini_image_edit_transformation.py | 0 .../test_gemini_realtime_transformation.py | 0 .../test_gemini_video_transformation.py | 0 .../chat/test_gigachat_chat_streaming.py | 0 .../chat/test_gigachat_chat_transformation.py | 28 -------------- .../test_gigachat_embedding_transformation.py | 30 --------------- ...est_gigachat_passthrough_transformation.py | 0 .../llms/gigachat/test_authenticator.py | 0 .../llms/gigachat/test_file_handler.py | 38 ------------------- .../llms/gigachat/test_utils.py | 0 ...github_copilot_embedding_transformation.py | 0 20 files changed, 16 insertions(+), 117 deletions(-) rename tests/{test_litellm => unit}/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py (97%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py (90%) rename tests/{test_litellm => unit}/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/files/test_gemini_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/gemini/image_edit/test_gemini_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/realtime/test_gemini_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/videos/test_gemini_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_streaming.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_transformation.py (95%) rename tests/{test_litellm => unit}/llms/gigachat/embedding/test_gigachat_embedding_transformation.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_authenticator.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_file_handler.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/test_utils.py (100%) rename tests/{test_litellm => unit}/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py (100%) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py rename to tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py rename to tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py similarity index 97% rename from tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py rename to tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index d0697ca9b0e..05e3812152e 100644 --- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -406,21 +406,6 @@ def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: assert headers["x-session-affinity"] == "sess-42" -def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) - pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) - with patch(HTTPX_CLIENT_FACTORY, return_value=client): - litellm.responses( - model="fireworks_ai/kimi-k3", - input="hi", - api_key="fw-test-key", - litellm_session_id="sess-42", - extra_headers=pinned, - ) - _, headers, _ = _sent_request(client) - assert headers["x-session-affinity"] == "explicit-node" - - def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: client: Final = MagicMock() request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py similarity index 90% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 52222f22a51..c6096ba2745 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -1,4 +1,5 @@ import math +from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -24,6 +25,15 @@ CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="firew OUTPUT_COST = 4.4e-06 +@pytest.fixture(autouse=True) +def restore_model_cost() -> Generator[None, None, None]: + original: Final = litellm.model_cost + litellm.get_model_info.cache_clear() + yield + litellm.model_cost = original + litellm.get_model_info.cache_clear() + + def _usage(prompt_tokens: int, cached_tokens: int, completion_tokens: int) -> Usage: return Usage( prompt_tokens=prompt_tokens, @@ -57,7 +67,7 @@ def _register_off_peak_model( cache_read_cost: float | None = STANDARD_CACHE_READ_COST, model: str = OFF_PEAK_MODEL, ) -> None: - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -151,7 +161,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente """Fireworks documents a default 50% cached-token discount for serverless models: https://docs.fireworks.ai/guides/prompt-caching, accessed 2026-09-19.""" model = "accounts/fireworks/models/default-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -171,7 +181,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): model = "accounts/fireworks/models/breakdown-cache-read-test" - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -204,7 +214,7 @@ def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): def test_generic_cost_per_token_applies_fireworks_cache_read_default_with_or_without_model_info(): model = "accounts/fireworks/models/generic-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -257,7 +267,7 @@ COMPONENT_AUDIO_OUT_COST = 6e-06 def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{COMPONENT_MODEL}": { "litellm_provider": "fireworks_ai", @@ -302,7 +312,7 @@ def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_ra def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # the SDK types model_cost as dict[Unknown, Unknown] "fireworks_ai/accounts/fireworks/models/no-input-rate-test": { "litellm_provider": "fireworks_ai", diff --git a/tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py b/tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py rename to tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/unit/llms/gemini/files/test_gemini_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py rename to tests/unit/llms/gemini/files/test_gemini_files_transformation.py diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py rename to tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py rename to tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py rename to tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py diff --git a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py b/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py rename to tests/unit/llms/gemini/videos/test_gemini_video_transformation.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py similarity index 100% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py similarity index 95% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 2f9511e642c..8e84072e549 100644 --- a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,22 +141,6 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" - @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") - @patch(f"{TRANSFORM_MODULE}.get_secret_str") - def test_falls_back_to_env_for_credentials( # test-quality-ok: mock-echo of internal wiring - self, mock_get_secret, mock_get_token - ): - mock_get_secret.return_value = "env-creds" - self.config.validate_environment( - headers={}, - model="GigaChat", - messages=[], - optional_params={}, - litellm_params={}, - api_key=None, - api_base=None, - ) - mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") # test-quality-ok: mock-echo of internal wiring class TestGetSupportedOpenAiParams: @@ -865,18 +849,6 @@ class TestUploadImage: def setup_method(self): self.config = GigaChatConfig() - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", return_value="file-uploaded") - def test_upload_image_success(self, mock_upload): - self.config._current_credentials = "creds" - self.config._current_api_base = "https://api.example.com" - result = self.config._upload_image("https://example.com/img.jpg") - assert result == "file-uploaded" - mock_upload.assert_called_once_with( - image_url="https://example.com/img.jpg", - credentials="creds", - api_base="https://api.example.com", - ) - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", side_effect=Exception("fail")) def test_upload_image_failure_returns_none(self, mock_upload): result = self.config._upload_image("https://example.com/img.jpg") diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py similarity index 91% rename from tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py rename to tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 8537793ea72..01fe66ca4c7 100644 --- a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -37,17 +37,6 @@ def _make_httpx_response(body: dict, status_code: int = 200) -> httpx.Response: # --------------------------------------------------------------------------- -class TestGetConfig: - def setup_method(self): - self.config = GigaChatEmbeddingConfig() - - def test_contains_only_abc_impl(self): - """get_config returns ABC internal data due to inheritance.""" - result = self.config.get_config() - # The only key should be _abc_impl from ABC base class - assert set(result.keys()) == {"_abc_impl"} - - class TestGetSupportedOpenAiParams: def setup_method(self): self.config = GigaChatEmbeddingConfig() @@ -287,25 +276,6 @@ class TestTransformEmbeddingResponse: ) assert result.model == "Embeddings" - def test_calls_logging_post_call(self): - raw = self._make_gigachat_response([ - {"object": "embedding", "embedding": [0.1], "index": 0}, - ]) - model_response = EmbeddingResponse() - self.config.transform_embedding_response( - model="gigachat/Embeddings", - raw_response=raw, - model_response=model_response, - logging_obj=self.logging_obj, - api_key="test-api-key", - request_data={"input": ["hello"]}, - optional_params={}, - litellm_params={}, - ) - self.logging_obj.post_call.assert_called_once() - args = self.logging_obj.post_call.call_args.kwargs - assert args["api_key"] == "test-api-key" - assert args["input"] == ["hello"] class TestValidateEnvironment: diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py rename to tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/unit/llms/gigachat/test_authenticator.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_authenticator.py rename to tests/unit/llms/gigachat/test_authenticator.py diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/unit/llms/gigachat/test_file_handler.py similarity index 91% rename from tests/test_litellm/llms/gigachat/test_file_handler.py rename to tests/unit/llms/gigachat/test_file_handler.py index ce9505f11f2..de83b2ddf5f 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/unit/llms/gigachat/test_file_handler.py @@ -344,25 +344,6 @@ class TestUploadFileSync: assert result is None - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_uploads_without_optional_args( - self, mock_http_handler_cls, mock_get_token, mock_get_api_base - ): - """Verify that credentials, api_base, and litellm_params are optional.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json.return_value = {"id": "file-no-args"} - mock_response.raise_for_status = MagicMock() - mock_client.post.return_value = mock_response - mock_http_handler_cls.return_value = mock_client - - result = upload_file_sync(image_url=_RED_PNG_DATA_URL) - - assert result == "file-no-args" - # Should still have called get_access_token without args - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) # --------------------------------------------------------------------------- @@ -483,22 +464,3 @@ class TestUploadFileAsync: ) assert result is None - - @pytest.mark.asyncio - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") - @patch(f"{FILE_MODULE}.get_async_httpx_client") - async def test_uploads_without_optional_args( - self, mock_get_client, mock_get_token, mock_get_api_base - ): - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json = MagicMock(return_value={"id": "async-no-args"}) - mock_response.raise_for_status = MagicMock() - mock_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_client - - result = await upload_file_async(image_url=_RED_PNG_DATA_URL) - - assert result == "async-no-args" - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) \ No newline at end of file diff --git a/tests/test_litellm/llms/gigachat/test_utils.py b/tests/unit/llms/gigachat/test_utils.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_utils.py rename to tests/unit/llms/gigachat/test_utils.py diff --git a/tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py rename to tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py From 443c9f838533027f9d09d60afc5021c6561e8197 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:09:49 +0000 Subject: [PATCH 050/127] test: migrate nvidia, oci, ocr, oobabooga and openai legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../audio_transcription/__init__.py | 0 tests/test_litellm/llms/oci/embed/__init__.py | 0 .../ocr/guardrail_translation/__init__.py | 0 .../test_litellm/llms/openai/chat/__init__.py | 0 .../chat/guardrail_translation/__init__.py | 0 ...t_nvidia_nim_passthrough_transformation.py | 0 .../test_nvidia_nim_rerank_transformation.py | 0 .../audio_transcription/test_audio_utils.py | 13 -- .../audio_transcription/test_handler.py | 0 .../test_transformation.py | 0 .../oci/chat/test_oci_chat_transformation.py | 150 ------------------ .../test_oci_chat_transformation_for_14158.py | 0 .../oci/chat/test_oci_cohere_tool_calls.py | 20 --- .../llms/oci/chat/test_oci_generic_chat.py | 12 -- .../llms/oci/chat/test_oci_sse_splitter.py | 0 .../oci/chat/test_oci_streaming_tool_calls.py | 0 .../embed/test_oci_embed_transformation.py | 22 --- .../llms/oci/embed/test_oci_embedding.py | 0 .../test_ocr_guardrail_handler.py | 0 .../llms/oobabooga/chat/test_oobabooga.py | 0 .../test_openai_guardrail_handler.py | 19 --- .../chat/test_openai_gpt_transformation.py | 0 .../completion/test_completion_handler.py | 0 .../test_text_completion_guardrail_handler.py | 0 .../test_text_completion_token_ids.py | 0 25 files changed, 236 deletions(-) delete mode 100644 tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py delete mode 100644 tests/test_litellm/llms/oci/embed/__init__.py delete mode 100644 tests/test_litellm/llms/ocr/guardrail_translation/__init__.py delete mode 100644 tests/test_litellm/llms/openai/chat/__init__.py delete mode 100644 tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py rename tests/{test_litellm => unit}/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_audio_utils.py (90%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/nvidia_riva/audio_transcription/test_transformation.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_chat_transformation.py (91%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_chat_transformation_for_14158.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_cohere_tool_calls.py (97%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_generic_chat.py (97%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_sse_splitter.py (100%) rename tests/{test_litellm => unit}/llms/oci/chat/test_oci_streaming_tool_calls.py (100%) rename tests/{test_litellm => unit}/llms/oci/embed/test_oci_embed_transformation.py (95%) rename tests/{test_litellm => unit}/llms/oci/embed/test_oci_embedding.py (100%) rename tests/{test_litellm => unit}/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/oobabooga/chat/test_oobabooga.py (100%) rename tests/{test_litellm => unit}/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py (99%) rename tests/{test_litellm => unit}/llms/openai/chat/test_openai_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_completion_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_text_completion_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/completion/test_text_completion_token_ids.py (100%) diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py b/tests/test_litellm/llms/nvidia_riva/audio_transcription/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/oci/embed/__init__.py b/tests/test_litellm/llms/oci/embed/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py b/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/chat/__init__.py b/tests/test_litellm/llms/openai/chat/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py b/tests/unit/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py rename to tests/unit/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py diff --git a/tests/test_litellm/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py rename to tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py similarity index 90% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py index 63a53c2c97b..54fc30e6f2d 100644 --- a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_audio_utils.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_audio_utils.py @@ -62,19 +62,6 @@ def test_resample_16khz_mono_passes_through_int16_bytes_match_length(): assert resampled.duration_seconds == pytest.approx(1.0, abs=0.001) -def test_resample_preserves_int16_clip_range(): - sample_rate = 16000 - samples = np.array([2.0, -2.0, 0.0, 1.0], dtype=np.float32) - wav_in = _wav_bytes(samples, sample_rate) - - resampled = resample_to_riva_pcm(wav_in) - - decoded = np.frombuffer(resampled.pcm_bytes, dtype="= -32767 - - def test_unknown_format_raises_clear_error(): # 4 random bytes are not valid audio in any container we can decode. with pytest.raises(NvidiaRivaException) as excinfo: diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_handler.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py diff --git a/tests/test_litellm/llms/nvidia_riva/audio_transcription/test_transformation.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/audio_transcription/test_transformation.py rename to tests/unit/llms/nvidia_riva/audio_transcription/test_transformation.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py similarity index 91% rename from tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py rename to tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 4c9bd29b337..708187b8ae1 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -922,91 +922,6 @@ class TestOCISignerSupport: assert wrapper.path_url == "/api/v1/chat" -class TestOCISplitChunks: - """ - Unit tests for the SSE split_chunks helpers used in sync and async streaming. - - These validate the fix for: - - Sync: JSONDecodeError when iter_text() returns chunks spanning multiple events - - Async: whitespace-only chunks being yielded before stripping (Greptile P2) - """ - - def _run_sync_split(self, raw_chunks): - """Invoke the sync split_chunks logic directly (extracted for testability).""" - results = [] - for item in raw_chunks: - for chunk in item.split("\n\n"): - stripped = chunk.strip() - if stripped: - results.append(stripped) - return results - - async def _run_async_split(self, raw_chunks): - """Invoke the async split_chunks logic directly.""" - results = [] - - async def _gen(): - for c in raw_chunks: - yield c - - async for item in _gen(): - for chunk in item.split("\n\n"): - stripped = chunk.strip() - if stripped: - results.append(stripped) - return results - - def test_sync_single_event_per_chunk(self): - """Normal case: one SSE event per iter_text() chunk.""" - chunks = ['data: {"text":"hello"}', 'data: {"text":"world"}'] - assert self._run_sync_split(chunks) == [ - 'data: {"text":"hello"}', - 'data: {"text":"world"}', - ] - - def test_sync_multiple_events_in_one_chunk(self): - """iter_text() returns two SSE events concatenated — must be split.""" - chunks = ['data: {"text":"a"}\n\ndata: {"text":"b"}'] - assert self._run_sync_split(chunks) == [ - 'data: {"text":"a"}', - 'data: {"text":"b"}', - ] - - def test_sync_whitespace_only_chunks_discarded(self): - """Whitespace between events must not be yielded.""" - chunks = ["data: {}\n\n \n\ndata: {}"] - result = self._run_sync_split(chunks) - assert result == ["data: {}", "data: {}"] - - def test_sync_empty_string_discarded(self): - """Empty string produced by splitting trailing \\n\\n must be discarded.""" - chunks = ["data: {}\n\n"] - assert self._run_sync_split(chunks) == ["data: {}"] - - @pytest.mark.asyncio - async def test_async_whitespace_only_chunks_discarded(self): - """ - Regression test for Greptile P2: async version was checking `if not chunk` - BEFORE stripping, so '\\n ' would pass the guard and yield '' downstream, - causing ValueError in chunk_creator ('Chunk does not start with data:'). - """ - chunks = ["data: {}\n\n \n\ndata: {}"] - result = await self._run_async_split(chunks) - assert result == ["data: {}", "data: {}"] - - @pytest.mark.asyncio - async def test_async_empty_string_discarded(self): - """Trailing \\n\\n must not produce an empty yielded chunk in async path.""" - chunks = ["data: {}\n\n"] - result = await self._run_async_split(chunks) - assert result == ["data: {}"] - - @pytest.mark.asyncio - async def test_async_multiple_events_in_one_chunk(self): - """Async path must split concatenated SSE events just like sync.""" - chunks = ['data: {"text":"x"}\n\ndata: {"text":"y"}'] - result = await self._run_async_split(chunks) - assert result == ['data: {"text":"x"}', 'data: {"text":"y"}'] class TestOCIProviderEmbeddingConfig: @@ -1026,21 +941,6 @@ class TestOCIProviderEmbeddingConfig: ) assert isinstance(config, OCIEmbedConfig) - def test_no_duplicate_oci_branch(self): - """ - Ensure utils.py does not contain two separate OCI embedding branches. - The dead code was removed in commit 64dfbe2b; this test guards against - regression (e.g. a future merge re-introducing it). - """ - import inspect - from litellm.utils import ProviderConfigManager - - source = inspect.getsource(ProviderConfigManager.get_provider_embedding_config) - oci_count = source.count("LlmProviders.OCI") - assert oci_count == 1, ( - f"Expected exactly 1 OCI branch in get_provider_embedding_config, found {oci_count}. " - "A duplicate dead-code branch may have been reintroduced." - ) class TestOCICohereParamMapping: @@ -1586,57 +1486,7 @@ def config(): class TestOCIKeyNormalization: """Tests for OCI private key content normalization.""" - def test_oci_key_with_escaped_newlines(self, config): - """Test that escaped newlines (\\n) are converted to actual newlines.""" - # Simulate PEM content with escaped newlines (as would come from JSON/UI input) - escaped_pem = "-----BEGIN RSA PRIVATE KEY-----\\nMIIEowIBAAKCAQEA...\\n-----END RSA PRIVATE KEY-----" - optional_params = { - "oci_user": "ocid1.user.oc1..test", - "oci_fingerprint": "aa:bb:cc:dd", - "oci_tenancy": "ocid1.tenancy.oc1..test", - "oci_region": "us-ashburn-1", - "oci_key": escaped_pem, - } - - # We can't fully test signing without a real key, but we can verify - # the error message indicates the key was processed (not a type error) - with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: - sign_with_manual_credentials( - headers={}, - optional_params=optional_params, - request_data={"test": "data"}, - api_base="https://test.oci.oraclecloud.com/api", - ) - - # The error should be about key format/loading, not about type - # This confirms the string was processed and newlines were normalized - error_message = str(exc_info.value) - assert "must be a string" not in error_message.lower() - - def test_oci_key_with_crlf_newlines(self, config): - """Test that Windows-style CRLF newlines are normalized to LF.""" - # Simulate PEM content with CRLF newlines - crlf_pem = "-----BEGIN RSA PRIVATE KEY-----\r\nMIIEowIBAAKCAQEA...\r\n-----END RSA PRIVATE KEY-----" - - optional_params = { - "oci_user": "ocid1.user.oc1..test", - "oci_fingerprint": "aa:bb:cc:dd", - "oci_tenancy": "ocid1.tenancy.oc1..test", - "oci_region": "us-ashburn-1", - "oci_key": crlf_pem, - } - - with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: - sign_with_manual_credentials( - headers={}, - optional_params=optional_params, - request_data={"test": "data"}, - api_base="https://test.oci.oraclecloud.com/api", - ) - - error_message = str(exc_info.value) - assert "must be a string" not in error_message.lower() def test_oci_key_rejects_non_string_type(self, config): """Test that non-string oci_key values raise OCIError.""" diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation_for_14158.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py rename to tests/unit/llms/oci/chat/test_oci_chat_transformation_for_14158.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py similarity index 97% rename from tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py rename to tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py index 729a2d25f41..9b06b01aa00 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -966,13 +966,6 @@ class TestOCICohereStreaming: completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging ) - def test_cohere_streaming_wrapper_initialization(self): - """Test OCIStreamWrapper initialization""" - stream_wrapper = self._create_stream_wrapper() - - # chunk_creator is the public dispatch entry point - assert hasattr(stream_wrapper, "chunk_creator") - assert callable(stream_wrapper.chunk_creator) def test_cohere_streaming_chunk_parsing(self): """Test parsing of Cohere streaming chunks""" @@ -1003,16 +996,3 @@ class TestOCICohereStreaming: # Test non-JSON chunk with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): stream_wrapper.chunk_creator("data: invalid json") - - def test_cohere_streaming_generic_chunk_fallback(self): - """Test fallback to generic chunk handling for non-Cohere chunks""" - stream_wrapper = self._create_stream_wrapper() - - # Test generic chunk (no apiFormat or different apiFormat) - generic_chunk = {"apiFormat": "GEMINI", "text": "Hello from Gemini"} - chunk_data = f"data: {json.dumps(generic_chunk)}" - - # This should fall back to generic handling - result = stream_wrapper.chunk_creator(chunk_data) - # The exact structure depends on the generic handler implementation - assert hasattr(result, "choices") diff --git a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py b/tests/unit/llms/oci/chat/test_oci_generic_chat.py similarity index 97% rename from tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py rename to tests/unit/llms/oci/chat/test_oci_generic_chat.py index 0a47852d085..9ec5ab9aed4 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py +++ b/tests/unit/llms/oci/chat/test_oci_generic_chat.py @@ -450,15 +450,3 @@ class TestGpt5MaxCompletionTokens: ) assert out.get("maxTokens") == 64 assert "maxCompletionTokens" not in out - - def test_payload_serializes_max_completion_tokens(self): - from litellm.types.llms.oci import OCIChatRequestPayload - - payload = OCIChatRequestPayload( - apiFormat="GENERIC", - messages=[], - maxCompletionTokens=64, - ) - dumped = payload.model_dump(exclude_none=True) - assert dumped["maxCompletionTokens"] == 64 - assert "maxTokens" not in dumped diff --git a/tests/test_litellm/llms/oci/chat/test_oci_sse_splitter.py b/tests/unit/llms/oci/chat/test_oci_sse_splitter.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_sse_splitter.py rename to tests/unit/llms/oci/chat/test_oci_sse_splitter.py diff --git a/tests/test_litellm/llms/oci/chat/test_oci_streaming_tool_calls.py b/tests/unit/llms/oci/chat/test_oci_streaming_tool_calls.py similarity index 100% rename from tests/test_litellm/llms/oci/chat/test_oci_streaming_tool_calls.py rename to tests/unit/llms/oci/chat/test_oci_streaming_tool_calls.py diff --git a/tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py similarity index 95% rename from tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py rename to tests/unit/llms/oci/embed/test_oci_embed_transformation.py index 363c0b46809..4ffd79ff147 100644 --- a/tests/test_litellm/llms/oci/embed/test_oci_embed_transformation.py +++ b/tests/unit/llms/oci/embed/test_oci_embed_transformation.py @@ -269,28 +269,6 @@ class TestOCIEmbedConfig: assert result.model == "cohere.embed-v3.0" assert result.usage.prompt_tokens == 10 - def test_transform_response_no_usage(self): - cfg = self._config() - model_response = EmbeddingResponse() - raw = self._mock_response( - 200, - { - "embeddings": [[0.1]], - "modelId": "cohere.embed-v3.0", - "modelVersion": "3.0.0", - }, - ) - result = cfg.transform_embedding_response( - model="cohere.embed-v3.0", - raw_response=raw, - model_response=model_response, - logging_obj=MagicMock(), - api_key=None, - request_data={}, - optional_params={}, - litellm_params={}, - ) - assert len(result.data) == 1 def test_transform_response_http_error_raises(self): cfg = self._config() diff --git a/tests/test_litellm/llms/oci/embed/test_oci_embedding.py b/tests/unit/llms/oci/embed/test_oci_embedding.py similarity index 100% rename from tests/test_litellm/llms/oci/embed/test_oci_embedding.py rename to tests/unit/llms/oci/embed/test_oci_embedding.py diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py b/tests/unit/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py rename to tests/unit/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py diff --git a/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py similarity index 100% rename from tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py rename to tests/unit/llms/oobabooga/chat/test_oobabooga.py diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py similarity index 99% rename from tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py rename to tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 258226ae22c..5c85faa5e13 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -545,25 +545,6 @@ class TestOpenAIChatCompletionsHandlerToolCallsInput: assert data["messages"][0]["content"] == "HELLO" assert data["messages"][1]["content"] == "HI THERE!" - @pytest.mark.asyncio - async def test_empty_tool_calls_list(self): - """Test that empty tool_calls list is handled correctly""" - handler = OpenAIChatCompletionsHandler() - guardrail = MockGuardrail() - - data = { - "messages": [ - {"role": "assistant", "content": "Hello", "tool_calls": []}, - ] - } - - # Process the input - await handler.process_input_messages(data, guardrail) - - # Verify empty tool_calls doesn't cause issues - assert guardrail.last_inputs is not None - tool_calls = guardrail.last_inputs.get("tool_calls", []) - assert len(tool_calls) == 0 class TestOpenAIChatCompletionsHandlerToolCallsOutput: diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py rename to tests/unit/llms/openai/chat/test_openai_gpt_transformation.py diff --git a/tests/test_litellm/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_completion_handler.py rename to tests/unit/llms/openai/completion/test_completion_handler.py diff --git a/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py rename to tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/completion/test_text_completion_token_ids.py b/tests/unit/llms/openai/completion/test_text_completion_token_ids.py similarity index 100% rename from tests/test_litellm/llms/openai/completion/test_text_completion_token_ids.py rename to tests/unit/llms/openai/completion/test_text_completion_token_ids.py From feca00248cf78fda0da5a1fa5b7f1a164fbe5608 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:11:56 +0000 Subject: [PATCH 051/127] test(bedrock): isolate host AWS config in realtime and rerank unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/realtime/test_bedrock_realtime_handler.py | 9 +++++++++ .../rerank/test_bedrock_rerank_header_forwarding.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index a7f0f64ef68..73a78a94e9f 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -18,6 +18,15 @@ from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + + class FakePayloadPart: def __init__(self, bytes_): self.bytes_ = bytes_ diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 2ea61b5e978..aa93ddb21b8 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -15,6 +15,15 @@ from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + # Mock response for Bedrock rerank # Format based on Bedrock rerank API response structure bedrock_rerank_response = { From 65a4a009585973a6d328cd981866b4d8e5188c1a Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:12:39 +0000 Subject: [PATCH 052/127] fix(ci): excuse retired test-quality rules in the budget ratchet Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/budget_ratchet_check.py | 43 +++++++++++++++---- scripts/check_test_quality.py | 4 ++ .../test_litellm/test_budget_ratchet_check.py | 30 +++++++++++++ tests/test_litellm/test_check_test_quality.py | 15 +++++++ 4 files changed, 84 insertions(+), 8 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 485e118efd2..470adddfcd3 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -7,13 +7,17 @@ driven DOWN over time. This check compares every budget file against its own content at the merge-base with the target branch and fails (exits 1, red) if: * a rule's `limit` went up, - * a rule was dropped from a budget (its ceiling effectively became infinite), or + * a rule was dropped from a budget (its ceiling effectively became infinite) while + its checker still emits it, or * an entire budget file was deleted. New rules and lowered/equal limits are fine. So is a rule that graduated: once a paired config (ruff.toml for the ruff-strict budget) selects the rule outright it hard-fails at the first violation, which is stricter than any ceiling the budget could hold, so dropping its entry tightens the guard rather than removing it. +Likewise a retired rule: once the paired checker (check_test_quality.py for the +test-quality budget) no longer emits a code, its entry has no ceiling left to +loosen. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -29,11 +33,12 @@ Usage: from __future__ import annotations import argparse +import importlib.util import json import subprocess import sys from pathlib import Path -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import Final, NamedTuple if sys.version_info >= (3, 11): @@ -49,6 +54,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "test-quality-budget.json", ) GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) +RETIREMENT_SOURCES = MappingProxyType({"test-quality-budget.json": "check_test_quality"}) class Regression(NamedTuple): @@ -139,20 +145,40 @@ def graduated_selectors(rel: str) -> tuple[str, ...]: ) +def _load_script(name: str) -> ModuleType: + if name in sys.modules: + return sys.modules[name] + spec: Final = importlib.util.spec_from_file_location(name, REPO_ROOT / "scripts" / f"{name}.py") + assert spec is not None and spec.loader is not None + module: Final = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def retired_rules(rel: str, base: dict) -> frozenset[str]: + """Rules in the base budget that the paired checker can no longer emit, so there is no ceiling to loosen.""" + source: Final = RETIREMENT_SOURCES.get(rel) + if source is None: + return frozenset() + return frozenset(_limits(base)) - _load_script(source).RULE_CODES + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], graduated: tuple[str, ...], + retired: frozenset[str] = frozenset(), ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. + """Why `rule` regressed vs base, or None when it held flat, fell, or left the budget legitimately. - A dropped rule is terminal unless it graduated; otherwise the only loosening - left is a raised limit. + A dropped rule is terminal unless it graduated or retired; otherwise the only + loosening left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: - if graduated and rule.startswith(graduated): + if rule in retired or (graduated and rule.startswith(graduated)): return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: @@ -165,6 +191,7 @@ def regressions_for( base: dict | None, head: dict | None, graduated: tuple[str, ...] = (), + retired: frozenset[str] = frozenset(), ) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet @@ -175,7 +202,7 @@ def regressions_for( return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated, retired)) is not None ] @@ -209,7 +236,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {base_ref} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel), retired_rules(rel, base))) if regressions: print( diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index dddc9d61982..9f93023cd53 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -146,6 +146,10 @@ SDK_MODULE: Final = "litellm" SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call")) INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) +RULE_CODES: Final = frozenset(( + "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", +)) + CREDENTIAL_NAME_RE: Final = re.compile( r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$" ) diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 22d05f4d00d..d6809b22161 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -92,6 +92,36 @@ def test_graduation_never_excuses_a_raised_limit(): assert "0 -> 7" in regs[0].detail +def test_dropped_rule_the_checker_retired_is_clean(): + base = {"TQ008": _spec_of(10993)} + assert ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) == [] + + +def test_dropped_rule_the_checker_still_emits_is_a_regression(): + base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + regs = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) + assert [r.rule for r in regs] == ["TQ001"] + assert "dropped" in regs[0].detail + + +def test_retirement_never_excuses_a_raised_limit(): + base = {"TQ008": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) + assert [r.rule for r in regs] == ["TQ008"] + assert "0 -> 7" in regs[0].detail + + +def test_retired_rules_come_from_the_paired_checker(): + base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + assert ratchet.retired_rules("test-quality-budget.json", base) == frozenset({"TQ008"}) + + +def test_budgets_without_a_paired_checker_never_retire(): + base = {"TQ008": _spec_of(1)} + for rel in ("ruff-strict-budget.json", "type-discipline-budget.json", "basedpyright-code-budget.json"): + assert ratchet.retired_rules(rel, base) == frozenset() + + def test_graduated_selectors_come_from_the_paired_ruff_config(): selectors = ratchet.graduated_selectors("ruff-strict-budget.json") assert "UP006" in selectors diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 05c25fb19fb..5a5c53fc31c 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -7,8 +7,10 @@ produced against tests/e2e, where the assertions live in a shared helper rather in the test body. """ +import ast import importlib.util import os +import re import subprocess import sys from pathlib import Path @@ -612,6 +614,19 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): assert all(" TQ001 " in line for line in reported) +def test_rule_codes_match_every_code_the_checker_emits(): + source = _MODULE_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + definition = next( + node + for node in tree.body + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "RULE_CODES" + ) + lines = source.splitlines() + outside = "\n".join(lines[: definition.lineno - 1] + lines[definition.end_lineno :]) + assert frozenset(re.findall(r'"(TQ\d{3})"', outside)) == checker.RULE_CODES + + def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path): source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n' assert _codes(tmp_path, source) == ["TQ009"] From ad523edb254fdb58e3a965339a0e44362d263320 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:13:11 +0000 Subject: [PATCH 053/127] test: migrate phase 9 legacy llm provider tests to tests/unit --- .../test_litellm_proxy_chat_transformation.py | 0 .../skills/test_code_execution.py | 0 .../litellm_proxy/skills/test_skill_search.py | 0 .../test_llamafile_chat_transformation.py | 0 tests/unit/llms/manus/__init__.py | 0 tests/unit/llms/manus/responses/__init__.py | 0 .../test_manus_responses_transformation.py | 0 .../test_meta_realtime_transformation.py | 0 .../test_meta_llama_chat_transformation.py | 0 tests/unit/llms/minimax/__init__.py | 0 tests/unit/llms/minimax/chat/__init__.py | 0 .../llms/minimax/chat/test_transformation.py | 96 ------------------- tests/unit/llms/minimax/messages/__init__.py | 0 .../minimax/messages/test_transformation.py | 74 -------------- tests/unit/llms/mistral/__init__.py | 0 ...est_mistral_audio_speech_transformation.py | 0 tests/unit/llms/mistral/batches/__init__.py | 0 .../test_mistral_batches_transformation.py | 0 tests/unit/llms/mistral/files/__init__.py | 0 .../test_mistral_files_transformation.py | 0 tests/unit/llms/mistral/ocr/__init__.py | 0 .../ocr/test_mistral_ocr_transformation.py | 0 ...est_modelscope_image_gen_transformation.py | 0 .../test_mongodb_transformation.py | 0 .../test_moonshot_chat_transformation.py | 25 ----- .../llms/neosantara/test_neosantara.py | 0 .../test_nimble_search_transformation.py | 0 .../chat/test_novita_chat_transformation.py | 9 -- .../chat/test_nscale_chat_transformation.py | 0 29 files changed, 204 deletions(-) rename tests/{test_litellm => unit}/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/litellm_proxy/skills/test_code_execution.py (100%) rename tests/{test_litellm => unit}/llms/litellm_proxy/skills/test_skill_search.py (100%) rename tests/{test_litellm => unit}/llms/llamafile/chat/test_llamafile_chat_transformation.py (100%) create mode 100644 tests/unit/llms/manus/__init__.py create mode 100644 tests/unit/llms/manus/responses/__init__.py rename tests/{test_litellm => unit}/llms/manus/responses/test_manus_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/meta/realtime/test_meta_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/meta_llama/test_meta_llama_chat_transformation.py (100%) create mode 100644 tests/unit/llms/minimax/__init__.py create mode 100644 tests/unit/llms/minimax/chat/__init__.py rename tests/{test_litellm => unit}/llms/minimax/chat/test_transformation.py (54%) create mode 100644 tests/unit/llms/minimax/messages/__init__.py rename tests/{test_litellm => unit}/llms/minimax/messages/test_transformation.py (57%) create mode 100644 tests/unit/llms/mistral/__init__.py rename tests/{test_litellm => unit}/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/batches/__init__.py rename tests/{test_litellm => unit}/llms/mistral/batches/test_mistral_batches_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/files/__init__.py rename tests/{test_litellm => unit}/llms/mistral/files/test_mistral_files_transformation.py (100%) create mode 100644 tests/unit/llms/mistral/ocr/__init__.py rename tests/{test_litellm => unit}/llms/mistral/ocr/test_mistral_ocr_transformation.py (100%) rename tests/{test_litellm => unit}/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/mongodb/vector_stores/test_mongodb_transformation.py (100%) rename tests/{test_litellm => unit}/llms/moonshot/test_moonshot_chat_transformation.py (96%) rename tests/{test_litellm => unit}/llms/neosantara/test_neosantara.py (100%) rename tests/{test_litellm => unit}/llms/nimble/search/test_nimble_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/novita/chat/test_novita_chat_transformation.py (85%) rename tests/{test_litellm => unit}/llms/nscale/chat/test_nscale_chat_transformation.py (100%) diff --git a/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py rename to tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_code_execution.py b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/skills/test_code_execution.py rename to tests/unit/llms/litellm_proxy/skills/test_code_execution.py diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py similarity index 100% rename from tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py rename to tests/unit/llms/litellm_proxy/skills/test_skill_search.py diff --git a/tests/test_litellm/llms/llamafile/chat/test_llamafile_chat_transformation.py b/tests/unit/llms/llamafile/chat/test_llamafile_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/llamafile/chat/test_llamafile_chat_transformation.py rename to tests/unit/llms/llamafile/chat/test_llamafile_chat_transformation.py diff --git a/tests/unit/llms/manus/__init__.py b/tests/unit/llms/manus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/manus/responses/__init__.py b/tests/unit/llms/manus/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py b/tests/unit/llms/manus/responses/test_manus_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py rename to tests/unit/llms/manus/responses/test_manus_responses_transformation.py diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py rename to tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py diff --git a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py b/tests/unit/llms/meta_llama/test_meta_llama_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py rename to tests/unit/llms/meta_llama/test_meta_llama_chat_transformation.py diff --git a/tests/unit/llms/minimax/__init__.py b/tests/unit/llms/minimax/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/minimax/chat/__init__.py b/tests/unit/llms/minimax/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/minimax/chat/test_transformation.py b/tests/unit/llms/minimax/chat/test_transformation.py similarity index 54% rename from tests/test_litellm/llms/minimax/chat/test_transformation.py rename to tests/unit/llms/minimax/chat/test_transformation.py index 9d51b556500..2645b2832aa 100644 --- a/tests/test_litellm/llms/minimax/chat/test_transformation.py +++ b/tests/unit/llms/minimax/chat/test_transformation.py @@ -2,14 +2,9 @@ Test MiniMax OpenAI-compatible API support """ -import os from unittest.mock import MagicMock, patch -import pytest - - import litellm -from litellm import completion from litellm.llms.minimax.chat.transformation import MinimaxChatConfig @@ -107,97 +102,6 @@ def test_minimax_provider_config_manager(): assert isinstance(config, MinimaxChatConfig) -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_basic(): - """Test basic chat completion with MiniMax OpenAI-compatible API""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello, how are you?"}, - ], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_with_reasoning_split(): - """Test completion with reasoning_split parameter (MiniMax M2.1 feature)""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Solve this problem: 2+2=?"}, - ], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - extra_body={"reasoning_split": True}, - ) - - assert response is not None - # Check if reasoning_details is present in response - if hasattr(response.choices[0].message, "reasoning_details"): - assert response.choices[0].message.reasoning_details is not None - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_with_tools(): - """Test completion with tool calling (function calling)""" - tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather in a location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - }, - } - ] - - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "What's the weather in San Francisco?"}], - tools=tools, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - assert response is not None - assert hasattr(response, "choices") - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_chat_completion_streaming(): - """Test streaming completion""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Count to 5"}], - stream=True, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/v1", - ) - - chunks = [] - for chunk in response: - chunks.append(chunk) - - assert len(chunks) > 0 - - if __name__ == "__main__": # Run basic tests that don't require API key print("Testing MiniMax Chat Config...") diff --git a/tests/unit/llms/minimax/messages/__init__.py b/tests/unit/llms/minimax/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/minimax/messages/test_transformation.py b/tests/unit/llms/minimax/messages/test_transformation.py similarity index 57% rename from tests/test_litellm/llms/minimax/messages/test_transformation.py rename to tests/unit/llms/minimax/messages/test_transformation.py index c7435a52890..a4b075414e3 100644 --- a/tests/test_litellm/llms/minimax/messages/test_transformation.py +++ b/tests/unit/llms/minimax/messages/test_transformation.py @@ -2,14 +2,9 @@ Test MiniMax Anthropic-compatible API support """ -import os from unittest.mock import MagicMock, patch -import pytest - - import litellm -from litellm import completion from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig @@ -58,75 +53,6 @@ def test_minimax_provider_config_manager(): assert config.custom_llm_provider == "minimax" -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_basic(): - """Test basic completion with MiniMax Anthropic-compatible API""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Hello, how are you?"}], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - ) - - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_with_thinking(): - """Test completion with thinking parameter (MiniMax M2.1 feature)""" - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "Solve this problem: 2+2=?"}], - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - thinking={"type": "enabled", "budget_tokens": 1000}, - ) - - assert response is not None - # Check if thinking content is present in response - for choice in response.choices: - if hasattr(choice.message, "content"): - # MiniMax returns thinking blocks similar to Anthropic - assert choice.message.content is not None - - -@pytest.mark.skip(reason="Requires actual MiniMax API key") -def test_minimax_completion_with_tools(): - """Test completion with tool calling (function calling)""" - tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather in a location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - }, - } - ] - - response = completion( - model="minimax/MiniMax-M2.1", - messages=[{"role": "user", "content": "What's the weather in San Francisco?"}], - tools=tools, - api_key=os.getenv("MINIMAX_API_KEY"), - api_base="https://api.minimax.io/anthropic/v1/messages", - ) - - assert response is not None - assert hasattr(response, "choices") - - if __name__ == "__main__": # Run basic tests that don't require API key print("Testing MiniMax Anthropic Config...") diff --git a/tests/unit/llms/mistral/__init__.py b/tests/unit/llms/mistral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py rename to tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py diff --git a/tests/unit/llms/mistral/batches/__init__.py b/tests/unit/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/unit/llms/mistral/batches/test_mistral_batches_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py rename to tests/unit/llms/mistral/batches/test_mistral_batches_transformation.py diff --git a/tests/unit/llms/mistral/files/__init__.py b/tests/unit/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/unit/llms/mistral/files/test_mistral_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py rename to tests/unit/llms/mistral/files/test_mistral_files_transformation.py diff --git a/tests/unit/llms/mistral/ocr/__init__.py b/tests/unit/llms/mistral/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py b/tests/unit/llms/mistral/ocr/test_mistral_ocr_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py rename to tests/unit/llms/mistral/ocr/test_mistral_ocr_transformation.py diff --git a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py rename to tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py diff --git a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py similarity index 100% rename from tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py rename to tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py rename to tests/unit/llms/moonshot/test_moonshot_chat_transformation.py index f94ea5e3db2..c39affc18a8 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py @@ -19,31 +19,6 @@ from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig class TestMoonshotConfig: """Test class for Moonshot AI functionality""" - def test_default_api_base(self): - """Test that default API base is used when none is provided""" - config = MoonshotChatConfig() - headers = {} - api_key = "fake-moonshot-key" - - # Call validate_environment without specifying api_base - result = config.validate_environment( - headers=headers, - model="moonshot-v1-8k", - messages=[{"role": "user", "content": "Hey"}], - optional_params={}, - litellm_params={}, - api_key=api_key, - api_base=None, # Not providing api_base - ) - - # Verify headers are still set correctly - assert result["Authorization"] == f"Bearer {api_key}" - assert result["Content-Type"] == "application/json" - - # We can't directly test the api_base value here since validate_environment - # only returns the headers, but we can verify it doesn't raise an exception - # which would happen if api_base handling was incorrect - def test_get_supported_openai_params(self): """Test that get_supported_openai_params returns correct params""" config = MoonshotChatConfig() diff --git a/tests/test_litellm/llms/neosantara/test_neosantara.py b/tests/unit/llms/neosantara/test_neosantara.py similarity index 100% rename from tests/test_litellm/llms/neosantara/test_neosantara.py rename to tests/unit/llms/neosantara/test_neosantara.py diff --git a/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py b/tests/unit/llms/nimble/search/test_nimble_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py rename to tests/unit/llms/nimble/search/test_nimble_search_transformation.py diff --git a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py b/tests/unit/llms/novita/chat/test_novita_chat_transformation.py similarity index 85% rename from tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py rename to tests/unit/llms/novita/chat/test_novita_chat_transformation.py index 3f2a3f77c41..1381cf95a5d 100644 --- a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py +++ b/tests/unit/llms/novita/chat/test_novita_chat_transformation.py @@ -54,12 +54,3 @@ class TestNovitaConfig: ) assert "Missing Novita AI API Key" in str(excinfo.value) - - def test_inheritance(self): - """Test proper inheritance from OpenAIGPTConfig""" - config = NovitaConfig() - - from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig - - assert isinstance(config, OpenAIGPTConfig) - assert hasattr(config, "get_supported_openai_params") diff --git a/tests/test_litellm/llms/nscale/chat/test_nscale_chat_transformation.py b/tests/unit/llms/nscale/chat/test_nscale_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/nscale/chat/test_nscale_chat_transformation.py rename to tests/unit/llms/nscale/chat/test_nscale_chat_transformation.py From fcabb626acdcf240c43ba78dfb92d7822e239cfe Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:19:05 +0000 Subject: [PATCH 054/127] test(unit): migrate wave 1 phase 3 anthropic, apiserpent, azure and azure_ai legacy tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 11 +++ .../test_reasoning_effort_fields.py | 0 .../test_anthropic_files_transformation.py | 0 .../messages/test_advisor_orchestration.py | 0 .../llms/apiserpent/test_apiserpent_search.py | 0 .../test_azure_image_edit_transformation.py | 0 .../test_azure_image_generation_init.py | 82 ------------------- .../test_azure_passthrough_transformation.py | 0 .../realtime/test_azure_realtime_handler.py | 35 -------- .../response/test_azure_transformation.py | 0 .../foundry_responses_web_search_fixture.json | 0 ...st_bing_grounding_search_transformation.py | 0 .../test_azure_tts_transformation.py | 0 ...test_azure_vector_stores_transformation.py | 0 .../chat/test_azure_ai_transformation.py | 15 ---- .../embed/test_azure_ai_embed_handler.py | 0 ...test_azure_ai_image_edit_transformation.py | 0 .../test_mai_image_edit_transformation.py | 0 ...st_azure_ai_cohere_parse_transformation.py | 0 ...est_azure_ai_passthrough_transformation.py | 0 .../test_azure_ai_rerank_transformation.py | 0 .../test_azure_ai_responses_transformation.py | 0 22 files changed, 11 insertions(+), 132 deletions(-) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/files/test_anthropic_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/messages/test_advisor_orchestration.py (100%) rename tests/{test_litellm => unit}/llms/apiserpent/test_apiserpent_search.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_edit/test_azure_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_generation/test_azure_image_generation_init.py (91%) rename tests/{test_litellm => unit}/llms/azure/passthrough/test_azure_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/realtime/test_azure_realtime_handler.py (94%) rename tests/{test_litellm => unit}/llms/azure/response/test_azure_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/search/foundry_responses_web_search_fixture.json (100%) rename tests/{test_litellm => unit}/llms/azure/search/test_bing_grounding_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/text_to_speech/test_azure_tts_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/vector_stores/test_azure_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/chat/test_azure_ai_transformation.py (97%) rename tests/{test_litellm => unit}/llms/azure_ai/embed/test_azure_ai_embed_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/responses/test_azure_ai_responses_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..d002452d14c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,6 +2,8 @@ from collections.abc import Iterator from typing import Final import pytest + +import litellm from pytest_socket import enable_socket, socket_allow_hosts LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -21,3 +23,12 @@ def block_external_sockets() -> Iterator[None]: @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() + + +@pytest.fixture +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py b/tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py rename to tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py rename to tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py similarity index 100% rename from tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py rename to tests/unit/llms/anthropic/messages/test_advisor_orchestration.py diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/unit/llms/apiserpent/test_apiserpent_search.py similarity index 100% rename from tests/test_litellm/llms/apiserpent/test_apiserpent_search.py rename to tests/unit/llms/apiserpent/test_apiserpent_search.py diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py rename to tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py similarity index 91% rename from tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py rename to tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index cfde1760389..eabd5c8427d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -133,88 +133,6 @@ def test_azure_image_generation_flattens_extra_body(): assert data["size"] == "1024x1024" -def test_azure_image_generation_creates_token_provider_from_credentials(): - """ - Test that azure_ad_token_provider is created from tenant_id, client_id, client_secret. - - This test verifies the fix in images/main.py where we now create the - azure_ad_token_provider from credentials in litellm_params if it's not already provided. - """ - # Simulate the fix in images/main.py - litellm_params_dict = { - "tenant_id": "test-tenant-id", - "client_id": "test-client-id", - "client_secret": "test-client-secret", - "azure_scope": None, - } - - azure_ad_token_provider = None - - # This is the logic we added in images/main.py - if azure_ad_token_provider is None: - tenant_id = litellm_params_dict.get("tenant_id") - client_id = litellm_params_dict.get("client_id") - client_secret = litellm_params_dict.get("client_secret") - azure_scope = ( - litellm_params_dict.get("azure_scope") - or "https://cognitiveservices.azure.com/.default" - ) - - # Verify the credentials are extracted correctly - assert tenant_id == "test-tenant-id" - assert client_id == "test-client-id" - assert client_secret == "test-client-secret" - assert azure_scope == "https://cognitiveservices.azure.com/.default" - - # Verify the condition to create token provider is met - assert ( - tenant_id and client_id and client_secret - ), "Credentials should be present to create token provider" - - -def test_azure_image_generation_headers_without_api_key(): - """ - Test that when api_key is None, the api-key header is not added to headers. - - This prevents the httpx TypeError: "Header value must be str or bytes, not " - that was occurring when api_key was None and being set in headers. - - This is a unit test for the fix in images/main.py where we now check: - if api_key is not None: - default_headers["api-key"] = api_key - """ - from litellm.images.main import image_generation - - # Test the header building logic directly - api_key = None - - default_headers = { - "Content-Type": "application/json", - } - - # This is the fix: only add api-key if it's not None - if api_key is not None: - default_headers["api-key"] = api_key - - # Verify api-key is not in headers when api_key is None - assert "api-key" not in default_headers - - # Verify Content-Type is still there - assert default_headers["Content-Type"] == "application/json" - - # Test with a valid api_key - api_key = "valid-key-123" - default_headers_with_key = { - "Content-Type": "application/json", - } - if api_key is not None: - default_headers_with_key["api-key"] = api_key - - # Verify api-key is added when api_key is valid - assert "api-key" in default_headers_with_key - assert default_headers_with_key["api-key"] == "valid-key-123" - - def test_azure_image_generation_drop_params_response_format(): """ Test that unsupported params like response_format are dropped when drop_params=True. diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py rename to tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py similarity index 94% rename from tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py rename to tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index 7d24e604569..c1ba286f8c0 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -426,41 +426,6 @@ async def test_async_realtime_beta_without_api_version_raises(): ) -@pytest.mark.asyncio -async def test_realtime_protocol_env_var_fallback(): - """ - Test that LITELLM_AZURE_REALTIME_PROTOCOL env var is used as fallback. - Fixes #22127: no way to set realtime_protocol from config. - """ - from litellm.realtime_api.main import _arealtime - from litellm.types.router import GenericLiteLLMParams - - with patch.dict(os.environ, {"LITELLM_AZURE_REALTIME_PROTOCOL": "v1"}): - # Create a GenericLiteLLMParams without realtime_protocol - litellm_params = GenericLiteLLMParams() - # The env var should be picked up as fallback - realtime_protocol = ( - {}.get("realtime_protocol") - or litellm_params.get("realtime_protocol") - or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") - or "beta" - ) - assert realtime_protocol == "v1" - - -@pytest.mark.asyncio -async def test_realtime_protocol_from_litellm_params(): - """ - Test that realtime_protocol is read from litellm_params (config.yaml extra field). - Fixes #22127: realtime_protocol in litellm_params was not used. - """ - from litellm.types.router import GenericLiteLLMParams - - # Simulate config.yaml with realtime_protocol as an extra field - litellm_params = GenericLiteLLMParams(realtime_protocol="GA") - assert litellm_params.get("realtime_protocol") == "GA" - - @pytest.mark.asyncio async def test_arealtime_transcription_intent_defaults_to_ga(monkeypatch): """ diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/response/test_azure_transformation.py rename to tests/unit/llms/azure/response/test_azure_transformation.py diff --git a/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json b/tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json similarity index 100% rename from tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json rename to tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json diff --git a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py rename to tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py diff --git a/tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py b/tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py rename to tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py diff --git a/tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py rename to tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py similarity index 97% rename from tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py rename to tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index f8cc0b5071e..e4a33d5772c 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -352,21 +352,6 @@ def test_azure_model_router_stamps_selected_model_on_hidden_params(): ) -def test_azure_model_router_stamp_does_not_leak_across_responses(): - """ - ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written - as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. - """ - from litellm.llms.azure_ai.common_utils import ( - AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, - ) - from litellm.types.utils import ModelResponse - - untouched = ModelResponse() - - assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) - - def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py rename to tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py rename to tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py rename to tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py rename to tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py rename to tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py From 6d327fff6f812b9c7dbbbb1ecc9b997a7807146e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 08:21:28 +0000 Subject: [PATCH 055/127] test(bedrock): keep beta headers fixture teardown off the network --- ...oke_transformations_anthropic_claude3_transformation.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 2c74d23a6a2..cf2fd78a896 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -86,11 +86,8 @@ def local_beta_headers_config(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") reload_beta_headers_config() - try: - yield - finally: - monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) - reload_beta_headers_config() + yield + reload_beta_headers_config() def test_get_supported_params_thinking(): From cf2a9b372cffe0e00e44ee252168384c8f5c058d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:21:51 +0000 Subject: [PATCH 056/127] test(gigachat): cover env credential fallback by its resulting auth header Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_gigachat_chat_transformation.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 8e84072e549..b1307f56336 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,6 +141,25 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" + @patch( + f"{TRANSFORM_MODULE}.get_access_token", + side_effect=lambda credentials, litellm_params: f"token-for-{credentials}", + ) + def test_falls_back_to_env_credentials_when_api_key_missing( + self, mock_get_token, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setenv("GIGACHAT_CREDENTIALS", "env-creds") + result = self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + assert result["Authorization"] == "Bearer token-for-env-creds" + assert self.config._current_credentials == "env-creds" class TestGetSupportedOpenAiParams: From 9850cd14f7e72bbe99f0b025cbb722d9a18cb523 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:27:02 +0000 Subject: [PATCH 057/127] test(ci): prove RULE_CODES by running every checker rule Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/budget_ratchet_check.py | 2 +- tests/test_litellm/test_check_test_quality.py | 49 ++++++++++++++----- 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 470adddfcd3..3ca5e9f3e9d 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -156,7 +156,7 @@ def _load_script(name: str) -> ModuleType: return module -def retired_rules(rel: str, base: dict) -> frozenset[str]: +def retired_rules(rel: str, base: dict[str, object]) -> frozenset[str]: """Rules in the base budget that the paired checker can no longer emit, so there is no ceiling to loosen.""" source: Final = RETIREMENT_SOURCES.get(rel) if source is None: diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 5a5c53fc31c..2a0e32d4de7 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -7,13 +7,13 @@ produced against tests/e2e, where the assertions live in a shared helper rather in the test body. """ -import ast import importlib.util import os -import re import subprocess import sys from pathlib import Path +from types import MappingProxyType +from typing import Final import pytest @@ -614,17 +614,42 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): assert all(" TQ001 " in line for line in reported) -def test_rule_codes_match_every_code_the_checker_emits(): - source = _MODULE_PATH.read_text(encoding="utf-8") - tree = ast.parse(source) - definition = next( - node - for node in tree.body - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "RULE_CODES" +_VIOLATING_SNIPPETS: Final = MappingProxyType( + { + "TQ000": ("test_snippet.py", "def test_broken(:\n pass\n"), + "TQ001": ("test_snippet.py", "def test_nothing():\n compute()\n"), + "TQ002": ( + "test_snippet.py", + "from unittest.mock import patch\n" + "\n" + "\n" + "def test_echo():\n" + " with patch('litellm.completion') as mock_completion:\n" + " run()\n" + " mock_completion.assert_called_once()\n", + ), + "TQ003": ("test_snippet.py", "import sys\n\nsys.path.insert(0, '..')\n"), + "TQ004": ("test_snippet.py", "import os\n\nos.environ['KEY'] = 'v'\n"), + "TQ005": ("test_snippet.py", "import litellm\n\nlitellm.drop_params = True\n"), + "TQ006": ("test_snippet.py", _DIRECT_GATE), + "TQ007": ("conftest.py", _SNAPSHOT_CONFTEST), + "TQ009": ( + "test_snippet.py", + 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n', + ), + } +) + + +def test_rule_codes_match_every_code_the_checker_emits(tmp_path): + emitted = frozenset( + v.code + for name, source in _VIOLATING_SNIPPETS.values() + for v in checker.check_file(_written(tmp_path, source, name)) ) - lines = source.splitlines() - outside = "\n".join(lines[: definition.lineno - 1] + lines[definition.end_lineno :]) - assert frozenset(re.findall(r'"(TQ\d{3})"', outside)) == checker.RULE_CODES + for code, (name, source) in _VIOLATING_SNIPPETS.items(): + assert code in [v.code for v in checker.check_file(_written(tmp_path, source, name))], code + assert emitted == checker.RULE_CODES def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path): From decb28f5b575378cff5e8ce941c999da24df3875 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:34:55 +0000 Subject: [PATCH 058/127] test(unit): restore live router and runtime model cost state between unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index d002452d14c..d47f71b21a8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,9 +2,11 @@ from collections.abc import Iterator from typing import Final import pytest +from pytest_socket import enable_socket, socket_allow_hosts import litellm -from pytest_socket import enable_socket, socket_allow_hosts +import litellm.router as litellm_router_module +import litellm.utils as litellm_utils_module LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -25,6 +27,24 @@ def pytest_runtest_setup() -> None: _allow_loopback_only() +@pytest.fixture(autouse=True) +def isolate_router_model_cost_state() -> Iterator[None]: + original_live_routers: Final = frozenset(litellm_router_module._live_routers) + original_runtime_registered_model_cost: Final = { + model_key: dict(model_value) + for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() + } + yield + for router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(router) + for router in original_live_routers: + litellm_router_module._live_routers.add(router) + litellm_utils_module._runtime_registered_model_cost.clear() + litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + litellm_utils_module._invalidate_model_cost_lowercase_map() + litellm.get_model_info.cache_clear() + + @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") From bbfa853d844b8c54c66746a416cafeef93f15b8f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:07:51 +0000 Subject: [PATCH 059/127] test(ci): annotate new test locals as Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_budget_ratchet_check.py | 15 ++++++++------- tests/test_litellm/test_check_test_quality.py | 2 +- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index d6809b22161..b359b8b42e5 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -9,6 +9,7 @@ import importlib.util import subprocess import sys from pathlib import Path +from typing import Final _MODULE_PATH = ( Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" @@ -93,31 +94,31 @@ def test_graduation_never_excuses_a_raised_limit(): def test_dropped_rule_the_checker_retired_is_clean(): - base = {"TQ008": _spec_of(10993)} + base: Final = {"TQ008": _spec_of(10993)} assert ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) == [] def test_dropped_rule_the_checker_still_emits_is_a_regression(): - base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} - regs = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) + base: Final = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + regs: Final = ratchet.regressions_for("b.json", base, {}, retired=frozenset({"TQ008"})) assert [r.rule for r in regs] == ["TQ001"] assert "dropped" in regs[0].detail def test_retirement_never_excuses_a_raised_limit(): - base = {"TQ008": _spec_of(0)} - regs = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) + base: Final = {"TQ008": _spec_of(0)} + regs: Final = ratchet.regressions_for("b.json", base, {"TQ008": _spec_of(7)}, retired=frozenset({"TQ008"})) assert [r.rule for r in regs] == ["TQ008"] assert "0 -> 7" in regs[0].detail def test_retired_rules_come_from_the_paired_checker(): - base = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} + base: Final = {"TQ001": _spec_of(5), "TQ008": _spec_of(10993)} assert ratchet.retired_rules("test-quality-budget.json", base) == frozenset({"TQ008"}) def test_budgets_without_a_paired_checker_never_retire(): - base = {"TQ008": _spec_of(1)} + base: Final = {"TQ008": _spec_of(1)} for rel in ("ruff-strict-budget.json", "type-discipline-budget.json", "basedpyright-code-budget.json"): assert ratchet.retired_rules(rel, base) == frozenset() diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index 2a0e32d4de7..bf05775d09d 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -642,7 +642,7 @@ _VIOLATING_SNIPPETS: Final = MappingProxyType( def test_rule_codes_match_every_code_the_checker_emits(tmp_path): - emitted = frozenset( + emitted: Final = frozenset( v.code for name, source in _VIOLATING_SNIPPETS.values() for v in checker.check_file(_written(tmp_path, source, name)) From 03a650c8a95802633c792d290a95bb8cbbb1159c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:17:57 +0000 Subject: [PATCH 060/127] test: migrate wave 1 phase 1 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 153 ------------------ .../test_pydantic_ai_agent_headers.py | 0 .../test_pydantic_ai_agent_transformation.py | 0 ...test_watsonx_orchestrate_transformation.py | 18 --- .../test_exception_mapping_utils.py | 0 .../batches/test_batch_utils.py | 0 .../batches/test_main.py | 0 .../batches/test_responses_batch_cost.py | 25 +-- .../chat_completions/test_dispatch.py | 13 +- ...responses_transformation_transformation.py | 0 .../compression/test_compress.py | 0 .../test_transformation.py | 0 .../test_callback_controls.py | 0 .../enterprise_callbacks/test_llm_guard.py | 0 .../test_secret_detection.py | 0 .../test_compression_interception_handler.py | 0 .../gcs_bucket/test_gcs_bucket_base.py | 0 .../integrations/gcs_pubsub/test_pub_sub.py | 0 .../helicone/test_helicone_gemini.py | 34 ---- 19 files changed, 14 insertions(+), 229 deletions(-) delete mode 100644 tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py (95%) rename tests/{test_litellm => unit}/anthropic_interface/exceptions/test_exception_mapping_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_batch_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_main.py (100%) rename tests/{test_litellm => unit}/batches/test_responses_batch_cost.py (87%) rename tests/{test_litellm => unit}/chat_completions/test_dispatch.py (93%) rename tests/{test_litellm => unit}/completion_extras/test_litellm_responses_transformation_transformation.py (100%) rename tests/{test_litellm => unit}/compression/test_compress.py (100%) rename tests/{test_litellm => unit}/endpoints/speech/speech_to_completion_bridge/test_transformation.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_callback_controls.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_llm_guard.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_secret_detection.py (100%) rename tests/{test_litellm => unit}/integrations/compression_interception/test_compression_interception_handler.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_bucket/test_gcs_bucket_base.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_pubsub/test_pub_sub.py (100%) rename tests/{test_litellm => unit}/integrations/helicone/test_helicone_gemini.py (73%) diff --git a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py b/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py deleted file mode 100644 index 8036c72679e..00000000000 --- a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py +++ /dev/null @@ -1,153 +0,0 @@ -""" -Regression test for https://github.com/BerriAI/litellm/issues/28505 - -the Responses API bridge double-strips the provider prefix from the -model name when a Chat Completions request has both `tools` and -`reasoning_effort`. - -Root cause: the bridge handler called `litellm.responses()` / -`litellm.aresponses()` without passing the already-resolved -`custom_llm_provider`. The downstream call then re-invoked -`get_llm_provider()` with `custom_llm_provider=None`, which stripped -a second provider prefix from a `provider/provider/model` deployment -string. - -This test pins both the sync and async bridge handler call sites: -the resolved `custom_llm_provider` must be forwarded to the underlying -`responses` / `aresponses` call so the provider isn't re-detected. -""" - -from unittest.mock import MagicMock, patch - -import pytest - -from litellm.completion_extras.litellm_responses_transformation.handler import ( - ResponsesToCompletionBridgeHandler, -) - - -def _validated_kwargs(): - return { - "model": "openai/openai/openai/gpt-5.5", - "messages": [{"role": "user", "content": "hi"}], - "optional_params": {}, - "litellm_params": {}, - "headers": {}, - "model_response": MagicMock(), - "logging_obj": MagicMock(), - "custom_llm_provider": "openai", - } - - -def test_sync_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - handler.transformation_handler.transform_response.return_value = ( - _validated_kwargs()["model_response"] - ) - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch( - "litellm.responses", - return_value=MagicMock(spec=[]), - ) as mock_responses, - ): - # The handler routes ResponsesAPIResponse through transform_response. - # We just want to verify the kwargs going INTO responses(). - try: - handler.completion(acompletion=False) - except Exception: - # Downstream handling (transform_response, type checks) is not - # the subject of this test. - pass - assert mock_responses.called - kwargs = mock_responses.call_args.kwargs - assert kwargs.get("custom_llm_provider") == "openai", ( - "sync bridge must forward custom_llm_provider to litellm.responses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("custom_llm_provider") == "openai", ( - "async bridge must forward custom_llm_provider to litellm.aresponses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_aws_region_name(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai.gpt-5.5", - "input": [], - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - validated = _validated_kwargs() - validated["custom_llm_provider"] = "bedrock_mantle" - validated["litellm_params"] = { - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - with ( - patch.object(handler, "validate_input_kwargs", return_value=validated), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("aws_region_name") == "us-east-2" diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py similarity index 95% rename from tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py rename to tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 43dfdaba02d..a300560ae9d 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,7 +1,6 @@ import asyncio import json import time -from pathlib import Path import httpx import pytest @@ -571,20 +570,3 @@ def test_config_manager_returns_wxo_provider(): ) assert config is not None assert config.__class__.__name__ == "WatsonxOrchestrateA2AConfig" - - -def test_wxo_dashboard_auth_fields(): - fields_path = ( - Path(__file__).resolve().parents[5] - / "litellm/proxy/public_endpoints/agent_create_fields.json" - ) - agent_fields = json.loads(fields_path.read_text()) - wxo_agent = next( - agent for agent in agent_fields if agent["agent_type"] == "watsonx_orchestrate" - ) - fields_by_key = {field["key"]: field for field in wxo_agent["credential_fields"]} - - assert fields_by_key["auth_mode"]["default_value"] == "cp4d" - # Username is CP4D-only; UI does not require it so ibm_cloud users are not blocked. - assert fields_by_key["username"]["required"] is False - assert "cp4d" in fields_by_key["username"]["tooltip"].lower() diff --git a/tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py similarity index 100% rename from tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py rename to tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py similarity index 100% rename from tests/test_litellm/batches/test_batch_utils.py rename to tests/unit/batches/test_batch_utils.py diff --git a/tests/test_litellm/batches/test_main.py b/tests/unit/batches/test_main.py similarity index 100% rename from tests/test_litellm/batches/test_main.py rename to tests/unit/batches/test_main.py diff --git a/tests/test_litellm/batches/test_responses_batch_cost.py b/tests/unit/batches/test_responses_batch_cost.py similarity index 87% rename from tests/test_litellm/batches/test_responses_batch_cost.py rename to tests/unit/batches/test_responses_batch_cost.py index b634f5f73db..63b28fb3b42 100644 --- a/tests/test_litellm/batches/test_responses_batch_cost.py +++ b/tests/unit/batches/test_responses_batch_cost.py @@ -12,17 +12,28 @@ Line shape decides the parse, not the batch's declared endpoint, so an output file mixing Responses-shaped and chat-shaped lines sums across both. """ -from typing import Literal, get_args, get_type_hints import pytest import litellm import litellm.batches.batch_utils as bu -from litellm.types.llms.openai import CreateBatchRequest MODEL = "gpt-5.6" +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _responses_line(input_tokens: int, output_tokens: int) -> dict: return { "response": { @@ -107,13 +118,3 @@ async def test_mixed_shape_batch_output_sums_across_both_line_shapes(local_model assert result.cost == pytest.approx( 133 * model_info["input_cost_per_token_batches"] + 107 * model_info["output_cost_per_token_batches"] ) - - -def test_create_batch_endpoint_accepts_v1_responses(): - """A type-checked caller can pass endpoint="/v1/responses", which the runtime - already forwarded correctly.""" - endpoint_annotation = get_type_hints(CreateBatchRequest)["endpoint"] - assert "/v1/responses" in get_args(endpoint_annotation) - - for create_fn in (litellm.create_batch, litellm.acreate_batch): - assert "/v1/responses" in get_args(get_type_hints(create_fn)["endpoint"]) diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py similarity index 93% rename from tests/test_litellm/chat_completions/test_dispatch.py rename to tests/unit/chat_completions/test_dispatch.py index d4bfeaf8d70..63821c74208 100644 --- a/tests/test_litellm/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -1,11 +1,9 @@ -import inspect from collections.abc import Awaitable, Callable, Mapping -from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures import pytest import litellm -from litellm import main as python_chat from litellm.chat_completions.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch @@ -40,15 +38,6 @@ def acompletion_binding(native: NativeAcompletion | None) -> NativeBinding[Nativ return binding -def test_public_signature_is_the_legacy_signature() -> None: - public_completion: Final = cast(Callable[..., object], litellm.completion) - legacy_completion: Final = cast(Callable[..., object], python_chat.completion) - public_acompletion: Final = cast(Callable[..., object], litellm.acompletion) - legacy_acompletion: Final = cast(Callable[..., object], python_chat.acompletion) - assert inspect.signature(public_completion) == inspect.signature(legacy_completion) - assert inspect.signature(public_acompletion) == inspect.signature(legacy_acompletion) - - def test_python_route_forwards_original_call_shape() -> None: metadata: Final = {"user_id": "u"} args: Final[tuple[object, ...]] = ("gpt-4o", MESSAGES) diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py similarity index 100% rename from tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py rename to tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py diff --git a/tests/test_litellm/compression/test_compress.py b/tests/unit/compression/test_compress.py similarity index 100% rename from tests/test_litellm/compression/test_compress.py rename to tests/unit/compression/test_compress.py diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py similarity index 100% rename from tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py rename to tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py b/tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py rename to tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py rename to tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py rename to tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py rename to tests/unit/integrations/compression_interception/test_compression_interception_handler.py diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py similarity index 100% rename from tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py rename to tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py diff --git a/tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py similarity index 100% rename from tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py rename to tests/unit/integrations/gcs_pubsub/test_pub_sub.py diff --git a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py b/tests/unit/integrations/helicone/test_helicone_gemini.py similarity index 73% rename from tests/test_litellm/integrations/helicone/test_helicone_gemini.py rename to tests/unit/integrations/helicone/test_helicone_gemini.py index 8ce02784345..667b16a48a1 100644 --- a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/unit/integrations/helicone/test_helicone_gemini.py @@ -3,7 +3,6 @@ Test HeliconeLogger Gemini/Vertex AI support. Fixes: https://github.com/BerriAI/litellm/issues/19093 """ -import pytest def test_helicone_gemini_model_in_list(): @@ -36,39 +35,6 @@ def test_helicone_gemini_models_recognized(): assert is_recognized, f"{model} should be recognized by helicone_model_list" -def test_helicone_vertex_ai_models_recognized(): - """ - Test that Vertex AI models (GLM, DeepSeek, etc.) are recognized via custom_llm_provider. - """ - # Test models that don't contain "gemini" but are vertex_ai - test_models = [ - "vertex_ai/zai-org/glm-4.7-maas", - "vertex_ai/deepseek-ai/deepseek-v3", - "vertex_ai/meta/llama-3.1-405b", - ] - for model in test_models: - is_vertex_ai = model.startswith("vertex_ai/") - assert is_vertex_ai, f"{model} should be recognized as vertex_ai model" - - -def test_helicone_vertex_ai_via_custom_llm_provider(): - """ - Test that vertex_ai models are recognized when custom_llm_provider is set. - """ - # Models without vertex_ai/ prefix but with custom_llm_provider="vertex_ai" - test_cases = [ - ("zai-org/glm-4.7-maas", "vertex_ai"), - ("deepseek-ai/deepseek-v3", "vertex_ai"), - ] - for model, custom_llm_provider in test_cases: - is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith( - "vertex_ai/" - ) - assert ( - is_vertex_ai - ), f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai" - - def test_helicone_vertex_gemini_gets_vertex_provider_url(): """ Test that vertex_ai/gemini-* models route to aiplatform.googleapis.com, From ba629f2537a73f74f5d466e8b7d2703902b74014 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:37:17 +0000 Subject: [PATCH 061/127] test(bedrock): wrap long lines flagged by review in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_prepare_request.py | 9 ++++++--- .../test_bedrock_passthrough_transformation.py | 13 ++++++++++--- .../realtime/test_bedrock_realtime_handler.py | 8 +++++++- .../rerank/test_bedrock_rerank_header_forwarding.py | 11 +++++++++-- .../test_bedrock_vector_store_transformation.py | 3 ++- 5 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py index 1575ccb5739..b010db3a840 100644 --- a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py @@ -11,7 +11,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" @@ -31,7 +32,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: assert ( request.endpoint_url - == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012%3Aapplication-inference-profile%2Fabcdefghi123/invoke" + == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012" + "%3Aapplication-inference-profile%2Fabcdefghi123/invoke" ) @@ -41,7 +43,8 @@ def test_bedrock_image_prepare_request_without_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 854ef92fa4b..d1d636a15f7 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -419,7 +419,9 @@ def test_bedrock_passthrough_model_id_arn_encoding(): ), f"ARN slash should be encoded, but found unencoded version in: {url_str}" # Verify the complete expected URL structure - expected_encoded_model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + expected_encoded_model_id = ( + "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + ) expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/converse" assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}" @@ -515,7 +517,10 @@ def test_bedrock_passthrough_model_id_without_arn(): def _event_frame(event_type: str, payload: dict) -> bytes: def header(name: str, value: str) -> bytes: name_b, value_b = name.encode(), value.encode() - return struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + return ( + struct.pack("!B", len(name_b)) + name_b + + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + ) payload_b = json.dumps(payload, separators=(",", ":")).encode() headers_b = ( @@ -589,7 +594,9 @@ def _feed(collector: PassthroughStreamCollector, stream: bytes, chunk_size: int def test_converse_stream_collector_keeps_usage_without_retaining_the_stream(): texts = [f"tok{i} " for i in range(4000)] - stream = _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + stream = ( + _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + ) _feed(_converse_stream_collector(), stream) tracemalloc.start() diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 73a78a94e9f..3aa827beb80 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -23,7 +23,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index aa93ddb21b8..c40830b238f 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -21,7 +21,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) # Mock response for Bedrock rerank @@ -39,7 +45,8 @@ bedrock_rerank_response = { test_query = "What is the capital of the United States?" test_documents = [ "Carson City is the capital city of the American state of Nevada.", - "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.", + "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. " + "Its capital is Saipan.", "Washington, D.C. is the capital of the United States.", ] diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index ab5a2531461..b45e70e31d3 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -46,7 +46,8 @@ def test_transform_search_request_encodes_vector_store_id(): assert ( url - == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve" + == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother" + "%3Fx%3D1%23frag/retrieve" ) assert body["retrievalQuery"].get("text") == "hello" From f61abe00a6bd3eff086e0fc5fc503e4418850d3e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:58:23 +0000 Subject: [PATCH 062/127] test(llms): wrap remaining lines over 120 chars in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/test_handler.py | 12 ++++-- ...drock_mantle_passthrough_transformation.py | 8 +++- .../chat/test_bytez_chat_transformation.py | 42 +++++++++++++++---- tests/unit/llms/chat/test_converse_handler.py | 5 ++- 4 files changed, 53 insertions(+), 14 deletions(-) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index dee8366ce2d..da7ed635dcb 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1072,7 +1072,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_text_delta_de_anonymized(self): - """Reasoning deltas carry model output; their text must be guardrailed while the reasoning signature is left untouched.""" + """Reasoning deltas carry model output; their text must be guardrailed while the + reasoning signature is left untouched.""" stream_bytes = ( _build_event_stream_frame("messageStart", {"role": "assistant"}) + _build_event_stream_frame( @@ -1105,7 +1106,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_tool_use_input_delta_de_anonymized(self): - """toolUse.input deltas carry model-generated tool arguments and must be guardrailed instead of being forwarded raw.""" + """toolUse.input deltas carry model-generated tool arguments and must be + guardrailed instead of being forwarded raw.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"q":""}'}}}, @@ -1154,7 +1156,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_text_and_reasoning_deltas_de_anonymized_independently(self): - """Distinct delta kinds must each be guardrailed and written back into their own field without bleeding the de-anonymized text across kinds.""" + """Distinct delta kinds must each be guardrailed and written back into their own + field without bleeding the de-anonymized text across kinds.""" captured = {} async def mock_hook(data, user_api_key_dict, response): @@ -1192,7 +1195,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_signature_only_frame_left_unmodified(self): - """A reasoning delta carrying only a signature has no guardrailable text; it must be forwarded untouched and the guardrail must not run.""" + """A reasoning delta carrying only a signature has no guardrailable text; it must + be forwarded untouched and the guardrail must not run.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig"}}}, diff --git a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py index 090de0a9d3e..27f6c9a9140 100644 --- a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py +++ b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -109,7 +109,13 @@ def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), ], ) -def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): +def test_sign_request_uses_the_deployment_bearer_token( + no_ambient_aws, + monkeypatch, + litellm_params, + env, + expected_bearer, +): for name, value in env.items(): monkeypatch.setenv(name, value) headers, body = BedrockMantlePassthroughConfig().sign_request( diff --git a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py index 440304aeac1..157e2e51175 100644 --- a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py @@ -8,6 +8,32 @@ from litellm.llms.bytez.chat.transformation import BytezChatConfig, API_BASE, ve TEST_API_KEY = "MOCK_BYTEZ_API_KEY" TEST_MODEL_NAME = "google/gemma-3-4b-it" TEST_MODEL = f"bytez/{TEST_MODEL_NAME}" +CAT_IMAGE_URL = ( + "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUX" + "VRLHI/male-orange-tabby-cat.jpg" +) +KAGGLE_AUDIO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_" + "SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-1616" + "07.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&" + "X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf" + "81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc39" + "0679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250" + "f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817" + "000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468" + "adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3" +) +KAGGLE_VIDEO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG" + "4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F202507" + "11%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-Signed" + "Headers=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5f" + "c6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72" + "084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb" + "90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d9" + "99f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189" + "c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947" +) TEST_MESSAGES = [{"role": "user", "content": "Hello"}] @@ -148,7 +174,7 @@ class TestBytezChatConfig: "What color is this cat?", { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -160,7 +186,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -174,7 +200,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -186,7 +212,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -200,7 +226,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "input_audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -212,7 +238,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -226,7 +252,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video_url", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } @@ -238,7 +264,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py index 12b5f03aedc..05debee0602 100644 --- a/tests/unit/llms/chat/test_converse_handler.py +++ b/tests/unit/llms/chat/test_converse_handler.py @@ -106,7 +106,10 @@ class TestBedrockRegionInModelPath: ), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}" assert ( optional_params.get("aws_region_name") == expected_region - ), f"region mismatch for {model!r}: got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ), ( + f"region mismatch for {model!r}: " + f"got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ) def test_explicit_aws_region_name_not_overridden(self): """ From ccd8997b0846469ff0e624b2cd837c4d0f8a3da6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:16:48 +0000 Subject: [PATCH 063/127] refactor(types): replace Any with proven types in 34 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_agentcore/transformation.py | 2 +- .../providers/watsonx_orchestrate/config.py | 4 ++-- litellm/a2a_protocol/utils.py | 6 ++--- .../gitlab/gitlab_prompt_manager.py | 4 ++-- litellm/integrations/otel/presets/agentops.py | 7 ++++-- .../integrations/vantage/vantage_logger.py | 4 ++-- litellm/interactions/agents/http_handler.py | 22 +++++++++---------- litellm/litellm_core_utils/logging_utils.py | 4 +++- litellm/litellm_core_utils/url_utils.py | 4 ++-- litellm/llms/anthropic/chat/handler.py | 2 +- litellm/llms/azure/completion/handler.py | 9 ++++---- .../document_intelligence/transformation.py | 8 +++---- .../guardrail_translation/base_translation.py | 2 +- .../llms/bedrock/batches/transformation.py | 10 ++++----- ...mazon_twelvelabs_pegasus_transformation.py | 2 +- litellm/llms/bytez/chat/transformation.py | 6 ++--- litellm/llms/custom_httpx/aiohttp_handler.py | 4 ++-- .../llms/deprecated_providers/aleph_alpha.py | 2 +- litellm/llms/lemonade/chat/transformation.py | 2 +- .../llms/openai/image_edit/transformation.py | 2 +- .../runwayml/text_to_speech/transformation.py | 2 +- .../llms/vertex_ai/files/transformation.py | 6 ++--- .../batch_embed_content_handler.py | 4 ++-- .../llms/vertex_ai/vertex_ai_non_gemini.py | 2 +- .../audio_transcription/transformation.py | 4 ++-- litellm/proxy/a2a/agent_card.py | 16 +++++++------- .../proxy/agent_endpoints/a2a_endpoints.py | 2 +- litellm/proxy/client/credentials.py | 5 +++-- .../guardrails_ai/guardrails_ai.py | 4 ++-- .../guardrail_hooks/singulr/singulr.py | 6 ++--- .../object_permission_utils.py | 13 ++++++----- .../proxy/policy_engine/pipeline_executor.py | 12 +++++----- .../router_strategy/adaptive_router/hooks.py | 2 +- .../auto_router/auto_router.py | 4 ++-- 34 files changed, 98 insertions(+), 90 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 9fa9db48af8..c486f1f6d95 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -85,7 +85,7 @@ def _filter_reserved_headers( def _request_scoped_runtime_session_id( - params: Mapping[str, Any], + params: Mapping[str, object], litellm_params: Mapping[str, Any], ) -> str | None: context_id: Final = get_session_id_from_a2a_params(params) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index ca84d3e07b4..44873edf271 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -20,7 +20,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Handle a non-streaming A2A request via WXO runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: @@ -40,7 +40,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """Handle a streaming A2A request via WXO streaming runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 47f561068cd..7400844bf28 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -17,7 +17,7 @@ class A2ARequestUtils: """Utility class for A2A request/response processing.""" @staticmethod - def extract_text_from_message(message: Any) -> str: + def extract_text_from_message(message: object) -> str: """ Extract text content from A2A message parts. @@ -142,7 +142,7 @@ class A2ARequestUtils: return prompt_tokens, completion_tokens, total_tokens -def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None: +def get_session_id_from_a2a_params(params: Mapping[str, object]) -> str | None: message: Final = params.get("message", {}) if isinstance(message, dict): return message.get("contextId") @@ -166,7 +166,7 @@ def scope_session_to_principal(session_id: str, principal: str | None) -> str: # Backwards compatibility aliases -def extract_text_from_a2a_message(message: Any) -> str: +def extract_text_from_a2a_message(message: object) -> str: return A2ARequestUtils.extract_text_from_message(message) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index d4602176650..817d280074f 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -200,8 +200,8 @@ class GitLabTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: - result: Final[dict[str, Any]] = {} + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, bool | int | float | str]: + result: Final[dict[str, bool | int | float | str]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 965213f2ee4..58123656caa 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,9 +9,12 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ +from collections.abc import Sequence from typing import Any, Final import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -71,7 +74,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> Any: +def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -106,7 +109,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Any) -> Any: + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: self._ensure_authenticated() return super().export(spans) diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index c219ba392ab..48a492fdd72 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -59,7 +59,7 @@ class VantageLogger(FocusLogger): raw_interval, ) - destination_config: Final[dict[str, Any]] = {} + destination_config: Final[dict[str, str]] = {} if resolved_api_key: destination_config["api_key"] = resolved_api_key if resolved_token: @@ -93,7 +93,7 @@ class VantageLogger(FocusLogger): pod_lock_manager = None if proxy_logging_obj is not None: - writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) + writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: pod_lock_manager = getattr(writer, "pod_lock_manager", None) diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index ec9df0fb488..2afedc34d36 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -7,7 +7,7 @@ duplicated. BaseAgentsAPIConfig stays as pure transform code. """ from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final import httpx @@ -38,7 +38,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, @@ -93,7 +93,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, @@ -141,7 +141,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -181,7 +181,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentListResponse: @@ -216,7 +216,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -259,7 +259,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentCreateResponse: @@ -295,7 +295,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -338,7 +338,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentDeleteResult: @@ -374,7 +374,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -417,7 +417,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentVersionsResponse: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 5be9dd7be2f..38a501ecaae 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -67,7 +67,9 @@ def _truncate_base64_in_string(value: str) -> str: return _DATA_URI_RE.sub(_base64_data_uri_replacer, value) -def _truncate_base64_in_value(value: Any) -> Any: +def _truncate_base64_in_value( + value: str | dict[str, object] | list[object] | None, +) -> str | dict[str, object] | list[object] | None: """Iteratively truncate base64 data URIs in a JSON-like value (str/list/dict). Uses an explicit stack instead of recursion to satisfy the project's diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..b94d5a6886d 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -418,7 +418,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -461,7 +461,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 359b8bb08c9..ef0f45d8f8b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -596,7 +596,7 @@ class ModelResponseIterator: self.reasoning_content_chunks: list[str] = [] # Track server tool use inputs and results for code_interpreter_results - self._server_tool_inputs: dict[str, Any] = {} + self._server_tool_inputs: dict[str, object] = {} self.tool_results: list[dict[str, Any]] = [] self._current_server_tool_id: str | None = None self._container_id: str | None = None diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 80934e994f6..23eef51e7ee 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -1,6 +1,7 @@ from collections.abc import Callable -from typing import Any, Final +from typing import Final +import httpx from openai import AsyncAzureOpenAI, AzureOpenAI from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -191,7 +192,7 @@ class AzureTextCompletion(BaseAzureLLM): model: str, api_base: str, data: dict, - timeout: Any, + timeout: float | httpx.Timeout | None, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, max_retries: int, @@ -253,7 +254,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, @@ -306,7 +307,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 3a2af8a5aba..23f532be757 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -12,7 +12,7 @@ import asyncio import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from urllib.parse import quote import httpx @@ -127,7 +127,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def map_ocr_params( self, - non_default_params: dict, + non_default_params: Mapping[str, object], optional_params: dict, model: str, ) -> dict: @@ -164,7 +164,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise UnsupportedParamsError(message=f"{e}", model=model, llm_provider="azure_ai") from e @staticmethod - def _normalize_pages_param(pages: Any) -> str: + def _normalize_pages_param(pages: object) -> str: """ Convert a caller-provided `pages` value to Azure DI's query-string form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`. @@ -412,7 +412,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise ValueError("Document URL is required") # Build Azure DI request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} # Check if it's a data URI (base64) if document_url.startswith("data:"): diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 89ad67f0485..ace5af8124f 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Any | None, + user_api_key_dict: Optional["UserAPIKeyAuth"], ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index ae0f8c5935b..973388ca5bd 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -2,7 +2,7 @@ import os import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError @@ -170,7 +170,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform the batch creation request to Bedrock format. @@ -354,7 +354,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) @staticmethod - def _get_openai_compatible_batch_metadata(metadata: Any) -> dict[str, str]: + def _get_openai_compatible_batch_metadata(metadata: object) -> dict[str, str]: """ OpenAI Batch metadata only accepts string values. """ @@ -379,7 +379,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform batch retrieval request for Bedrock. @@ -523,7 +523,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Enrich metadata with useful Bedrock fields - enriched_metadata_raw: Final[dict[str, Any]] = { + enriched_metadata_raw: Final[dict[str, object]] = { "jobName": response_data.get("jobName"), "clientRequestToken": response_data.get("clientRequestToken"), "modelId": response_data.get("modelId"), diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index d12c8aee48c..39cded4ed64 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -110,7 +110,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): headers: dict, ) -> dict: input_prompt: Final = self._convert_messages_to_prompt(messages=messages) - request_data: Final[dict[str, Any]] = {"inputPrompt": input_prompt} + request_data: Final[dict[str, object]] = {"inputPrompt": input_prompt} media_source: Final = self._build_media_source(optional_params) if media_source is not None: diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index d9a0c98b6db..7977db0f056 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -335,10 +335,10 @@ class BytezChatConfig(BaseConfig): class BytezCustomStreamWrapper(CustomStreamWrapper): - def chunk_creator(self, chunk: Any): + def chunk_creator(self, chunk: object): try: model_response: Final = self.model_response_creator() - response_obj: dict[str, Any] = {} + response_obj: dict[str, object] = {} response_obj = { "text": chunk, @@ -346,7 +346,7 @@ class BytezCustomStreamWrapper(CustomStreamWrapper): "finish_reason": "", } - completion_obj: Final[dict[str, Any]] = {"content": chunk} + completion_obj: Final[dict[str, object]] = {"content": chunk} return self.return_processed_chunk_logic( completion_obj=completion_obj, diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 7035ce58ae1..0809ef5274f 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,5 +1,5 @@ import ssl -from collections.abc import Callable +from collections.abc import AsyncIterable, Callable, Iterable from typing import TYPE_CHECKING, Any, Final, cast import aiohttp @@ -212,7 +212,7 @@ class BaseLLMAIOHTTPHandler: litellm_params: dict, stream: bool = False, files: dict | None = None, - content: Any = None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, params: dict | None = None, ) -> httpx.Response: max_retry_on_unprocessable_entity_error: Final = provider_config.max_retry_on_unprocessable_entity_error diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 4a29549b6aa..2ad9ce4edc8 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -146,7 +146,7 @@ class AlephAlphaConfig: setattr(self.__class__, key, value) @classmethod - def get_config(cls): + def get_config(cls) -> dict[str, object]: return { k: v for k, v in cls.__dict__.items() diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 553478aec16..c01ad2a0edd 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -170,7 +170,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> Any: + ) -> dict[str, object]: if model.startswith("lemonade/"): model = model.split("/", 1)[1] diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index f55084adbde..d54522597a0 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -66,7 +66,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): def _add_image_to_files( self, files_list: list[tuple[str, Any]], - image: Any, + image: object, field_name: str, ) -> None: """Add an image to the files list with appropriate content type""" diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 19e6d8ff494..6769accc1d6 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -78,7 +78,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", Coroutine[object, object, "HttpxBinaryResponseContent"], diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 85ec2911464..80d32289c94 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -651,7 +651,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. @@ -774,7 +774,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -948,7 +948,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index f81d4ca777e..c6ac87d646b 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal import httpx @@ -210,7 +210,7 @@ class GoogleBatchEmbeddings(VertexLLM): ) ### TRANSFORMATION (sync path) ### - request_data: Any + request_data: VertexAIBatchEmbeddingsRequestBody | dict[str, object] if use_embed_content: resolved_files = {} if api_key: diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 1c582c7c376..a7a1ea8d88d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -64,7 +64,7 @@ def _get_client_from_cache(client_cache_key: str): return litellm.in_memory_llm_clients_cache.get_cache(client_cache_key) -def _set_client_in_cache(client_cache_key: str, vertex_llm_model: Any): +def _set_client_in_cache(client_cache_key: str, vertex_llm_model: object): litellm.in_memory_llm_clients_cache.set_cache( key=client_cache_key, value=vertex_llm_model, diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 7d1aba63428..2169b9bf49a 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/audio/transcriptions` to IBM WatsonX's `/ml/v1/aud WatsonX follows the OpenAI spec for audio transcription. """ -from typing import Any, Final +from typing import Final from httpx import Response @@ -124,7 +124,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran } # Convert TypedDict to regular dict for AudioTranscriptionRequestData - form_data_dict: Final[dict[str, Any]] = dict(form_data) + form_data_dict: Final[dict[str, object]] = dict(form_data) return AudioTranscriptionRequestData(data=form_data_dict, files=files) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 5bec5158bcc..3718359fbb6 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -10,7 +10,7 @@ and uses LiteLLM auth. import re from collections.abc import Mapping from copy import deepcopy -from typing import Any, Final, Literal +from typing import Final, Literal SupportedA2AVersion = Literal["0.3", "1.0"] @@ -44,7 +44,7 @@ def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) -def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: +def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str: """Return the validated protocol version an agent card pins, else the default.""" normalized: Final = normalize_protocol_version(card.get("protocolVersion") if card else None) return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION @@ -53,7 +53,7 @@ def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces # whatever upstream advertised — the client must authenticate to the proxy, # not the upstream agent. -LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, Any]]] = { +LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, str]]] = { "LiteLLMKey": { "type": "http", "scheme": "bearer", @@ -112,7 +112,7 @@ _ALLOWED_TOP_LEVEL_KEYS: Final = { "url", } -_DEFAULT_SKILLS: Final[list[dict[str, Any]]] = [ +_DEFAULT_SKILLS: Final[list[dict[str, str | list[str]]]] = [ { "id": "chat", "name": "Chat", @@ -129,7 +129,7 @@ _DEFAULT_MODES: Final[list[str]] = ["text"] _DEFAULT_AGENT_VERSION: Final = "1.0.0" -def _filter_capabilities(upstream_capabilities: Any) -> dict[str, Any]: +def _filter_capabilities(upstream_capabilities: object) -> dict[str, object]: """Return a capabilities dict containing only allowlisted, truthy keys.""" if not isinstance(upstream_capabilities, dict): return {} @@ -143,13 +143,13 @@ def _default_litellm_provider(proxy_base_url: str) -> dict[str, str]: def merge_agent_card( - upstream_card: Mapping[str, Any] | None, + upstream_card: Mapping[str, object] | None, *, proxy_url: str, proxy_base_url: str, name: str | None = None, description: str | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build the LiteLLM-fronted agent card. @@ -169,7 +169,7 @@ def merge_agent_card( A dict suitable for serving as the proxy's agent card. Only keys in the v1.0 AgentCard schema (plus ``supportedInterfaces``) are emitted. """ - base: Final[dict[str, Any]] = deepcopy(dict(upstream_card)) if upstream_card else {} + base: Final[dict[str, object]] = deepcopy(dict(upstream_card)) if upstream_card else {} # Keep the upstream ``url`` on the stored card: the runtime A2A # invocation path reads it from ``agent_card_params`` to know where to diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..faf3e98a3a7 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, Any]] = ( + response_dict: Final[dict[str, object]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index a9bff67b1c5..d9edecd2eb7 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import requests @@ -69,8 +70,8 @@ class CredentialsManagementClient: def create( self, credential_name: str, - credential_info: dict[str, Any], - credential_values: dict[str, Any], + credential_info: Mapping[str, object], + credential_values: Mapping[str, object], return_request: bool = False, ) -> dict[str, Any] | requests.Request: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 47324471650..18451df574f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,7 +7,7 @@ import json import os -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict +from typing import TYPE_CHECKING, Final, Literal, TypedDict from fastapi import HTTPException @@ -181,7 +181,7 @@ class GuardrailsAI(CustomGuardrail): ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm return await self.process_input(data=data, call_type=call_type) - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: if call_type == "acompletion" or call_type == "completion": kwargs = await self.process_input(data=kwargs, call_type=call_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 06d4b39f5f6..bd5b18e368d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) @@ -36,7 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( ToolCall, ToolCallFunction, ) -from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, GenericGuardrailAPIInputs _DEFAULT_API_BASE: Final = "http://localhost:8003" _GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2" @@ -339,7 +339,7 @@ class SingulrGuardrail(CustomGuardrail): return inputs @staticmethod - def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None": + def _build_tool_call(tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall) -> "ToolCall | None": tool_call_id: Final = tool_call.get("id") fun: Final = tool_call.get("function") if not tool_call_id or not fun: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index daab38d3662..61e432daa16 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,12 +156,13 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - new_object_permission: dict | str | None = data_json.pop("object_permission", None) - if new_object_permission is None: + raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) + if raw_object_permission is None: return None - if isinstance(new_object_permission, str): - new_object_permission = json.loads(new_object_permission) + new_object_permission: Final[object] = ( + json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission + ) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -230,7 +231,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 0b81e7af84d..e9d23436b59 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -8,7 +8,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. import copy import time from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar +from typing import TYPE_CHECKING, Final, Literal, TypeVar from pydantic import BaseModel @@ -314,11 +314,11 @@ class PipelineExecutor: steps: list[PipelineStep], mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ @@ -490,10 +490,10 @@ class PipelineExecutor: step: PipelineStep, mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], @@ -722,7 +722,7 @@ def _extract_error_message(e: Exception) -> str: if isinstance(e, ModifyResponseException): return str(e) if HTTPException is not None and isinstance(e, HTTPException): - detail: Final = getattr(e, "detail", None) + detail: Final[object] = getattr(e, "detail", None) if detail: return str(detail) return str(e) diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 709910753f2..c4a2eae1ef9 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -86,7 +86,7 @@ def _resolve_session_key(kwargs: dict[str, Any]) -> str | None: return hashlib.sha256(payload.encode("utf-8")).hexdigest() -def _last_user_content(messages: list[dict[str, Any]] | None) -> str | None: +def _last_user_content(messages: Sequence[Mapping[str, object]] | None) -> str | None: if not messages: return None for msg in reversed(messages): diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index d08afa8c1f6..250201a46d1 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -3,7 +3,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ import asyncio -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel, ConfigDict @@ -158,7 +158,7 @@ class AutoRouter(CustomLogger): return await asyncio.shield(build_task) @staticmethod - def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str: + def _extract_text_from_messages(messages: Sequence[Mapping[str, object]]) -> str: """ Extract text content from the last user message for routing. From dbf566873ea32a82a9915876615dad67febd2e57 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:28:06 +0000 Subject: [PATCH 064/127] refactor(types): keep agentops preset imports optional Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/presets/agentops.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 58123656caa..965213f2ee4 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,12 +9,9 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ -from collections.abc import Sequence from typing import Any, Final import httpx -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -74,7 +71,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: +def _build_agentops_exporter(spec: ExporterSpec) -> Any: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -109,7 +106,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + def export(self, spans: Any) -> Any: self._ensure_authenticated() return super().export(spans) From 729a96e4eaff1d6fe4e68d5a1137ba83b07c0fe0 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:40:27 +0000 Subject: [PATCH 065/127] test: migrate openai, openai_like and openrouter legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/openai/evals/__init__.py | 1 - tests/test_litellm/llms/openai_like/embedding/__init__.py | 1 - tests/test_litellm/llms/openai_like/messages/__init__.py | 0 tests/test_litellm/llms/openrouter/image_edit/__init__.py | 0 .../guardrail_translation/test_embeddings_guardrail_handler.py | 0 .../llms/openai/evals/test_openai_evals_transformation.py | 0 .../llms/openai/image_generation/test_gpt_transformation.py | 0 .../image_generation/test_image_generation_guardrail_handler.py | 0 .../test_openai_image_generation_extra_headers.py | 0 .../llms/openai/speech/test_text_to_speech_guardrail_handler.py | 0 .../transcriptions/test_audio_transcription_guardrail_handler.py | 0 .../openai/transcriptions/test_transcription_duration_hidden.py | 0 .../llms/openai/transcriptions/test_whisper_transformation.py | 0 .../test_openai_vector_store_files_transformation.py | 0 .../vector_stores/test_openai_vector_stores_transformation.py | 0 .../llms/openai/videos/test_openai_video_transformation.py | 0 .../openai_like/chat/test_openai_like_chat_transformation.py | 0 .../llms/openai_like/embedding/test_openai_like_embedding.py | 0 .../test_openai_like_anthropic_messages_transformation.py | 0 .../llms/openrouter/chat/test_openrouter_chat_transformation.py | 0 .../image_edit/test_openrouter_image_edit_transformation.py | 0 .../image_generation/test_openrouter_image_gen_transformation.py | 0 .../llms/openrouter/test_openrouter_embedding_transformation.py | 0 .../llms/openrouter/test_openrouter_provider_routing.py | 0 24 files changed, 2 deletions(-) delete mode 100644 tests/test_litellm/llms/openai/evals/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/embedding/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/messages/__init__.py delete mode 100644 tests/test_litellm/llms/openrouter/image_edit/__init__.py rename tests/{test_litellm => unit}/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/evals/test_openai_evals_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_image_generation_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_openai_image_generation_extra_headers.py (100%) rename tests/{test_litellm => unit}/llms/openai/speech/test_text_to_speech_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_transcription_duration_hidden.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_whisper_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_stores/test_openai_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/videos/test_openai_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/chat/test_openai_like_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/embedding/test_openai_like_embedding.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/chat/test_openrouter_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_provider_routing.py (100%) diff --git a/tests/test_litellm/llms/openai/evals/__init__.py b/tests/test_litellm/llms/openai/evals/__init__.py deleted file mode 100644 index 47a8a2f0aed..00000000000 --- a/tests/test_litellm/llms/openai/evals/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""OpenAI Evals API tests""" diff --git a/tests/test_litellm/llms/openai_like/embedding/__init__.py b/tests/test_litellm/llms/openai_like/embedding/__init__.py deleted file mode 100644 index 2cb77227ed0..00000000000 --- a/tests/test_litellm/llms/openai_like/embedding/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Test module for OpenAI-like embedding handler diff --git a/tests/test_litellm/llms/openai_like/messages/__init__.py b/tests/test_litellm/llms/openai_like/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openrouter/image_edit/__init__.py b/tests/test_litellm/llms/openrouter/image_edit/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py b/tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py rename to tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py rename to tests/unit/llms/openai/evals/test_openai_evals_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py b/tests/unit/llms/openai/image_generation/test_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py rename to tests/unit/llms/openai/image_generation/test_gpt_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py b/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py rename to tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py rename to tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py diff --git a/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py b/tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py rename to tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py b/tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py rename to tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py rename to tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py b/tests/unit/llms/openai/transcriptions/test_whisper_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py rename to tests/unit/llms/openai/transcriptions/test_whisper_transformation.py diff --git a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py b/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py rename to tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py diff --git a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py b/tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py rename to tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/unit/llms/openai/videos/test_openai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py rename to tests/unit/llms/openai/videos/test_openai_video_transformation.py diff --git a/tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py rename to tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py diff --git a/tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py b/tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py similarity index 100% rename from tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py rename to tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py rename to tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py rename to tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py b/tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py rename to tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py rename to tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py rename to tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py b/tests/unit/llms/openrouter/test_openrouter_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py rename to tests/unit/llms/openrouter/test_openrouter_provider_routing.py From 7cf9a3035ce130fb4a07a02c7e139df6915fd378 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:02 +0000 Subject: [PATCH 066/127] test: migrate phase 16 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rust_bridge/__init__.py | 0 tests/unit/rust_bridge/ocr/__init__.py | 0 tests/unit/rust_bridge/ocr/test_route_host.py | 85 ++ tests/unit/rust_bridge/responses/__init__.py | 0 .../rust_bridge/responses/test_route_host.py | 57 + tests/unit/sandbox/test_e2b_sandbox.py | 318 +++++ .../unit/sandbox/test_opensandbox_sandbox.py | 647 ++++++++++ tests/unit/sandbox/test_sandbox_tools.py | 181 +++ tests/unit/skills/test_skills_main.py | 57 + .../test_enforce_model_rate_limits.py | 468 ++++++++ .../test_router/test_io_token_rate_limits.py | 1041 +++++++++++++++++ .../types/llms/test_types_llms_bedrock.py | 46 + .../unit/types/llms/test_types_llms_openai.py | 591 ++++++++++ .../types/proxy/policy_engine/__init__.py | 0 .../policy_engine/test_pipeline_types.py | 168 +++ .../proxy/policy_engine/test_policy_types.py | 15 + .../policy_engine/test_resolver_types.py | 115 ++ tests/unit/videos/__init__.py | 0 tests/unit/videos/test_main.py | 455 +++++++ tests/unit/videos/test_utils.py | 181 +++ 20 files changed, 4425 insertions(+) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/test_route_host.py create mode 100644 tests/unit/rust_bridge/responses/__init__.py create mode 100644 tests/unit/rust_bridge/responses/test_route_host.py create mode 100644 tests/unit/sandbox/test_e2b_sandbox.py create mode 100644 tests/unit/sandbox/test_opensandbox_sandbox.py create mode 100644 tests/unit/sandbox/test_sandbox_tools.py create mode 100644 tests/unit/skills/test_skills_main.py create mode 100644 tests/unit/test_router/test_enforce_model_rate_limits.py create mode 100644 tests/unit/test_router/test_io_token_rate_limits.py create mode 100644 tests/unit/types/llms/test_types_llms_bedrock.py create mode 100644 tests/unit/types/llms/test_types_llms_openai.py create mode 100644 tests/unit/types/proxy/policy_engine/__init__.py create mode 100644 tests/unit/types/proxy/policy_engine/test_pipeline_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_policy_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_resolver_types.py create mode 100644 tests/unit/videos/__init__.py create mode 100644 tests/unit/videos/test_main.py create mode 100644 tests/unit/videos/test_utils.py diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/__init__.py b/tests/unit/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/test_route_host.py b/tests/unit/rust_bridge/ocr/test_route_host.py new file mode 100644 index 00000000000..699492e4424 --- /dev/null +++ b/tests/unit/rust_bridge/ocr/test_route_host.py @@ -0,0 +1,85 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure +from litellm.rust_bridge.ocr.route_host import response as build_ocr_response +from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest + +REQUEST: Final = LiteLLMOcrRequest( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"req_format": "markdown"}, +) + + +class RustUpstreamError(Exception): + def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: + super().__init__(status, body) + self.headers: Final = list(headers) + + +class RustFormatError(Exception): + ocr_request_format_error: Final = True + + +def test_rust_ocr_response_retains_provider_native_response(): + provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} + response = build_ocr_response( + { + "pages": [], + "model": "prebuilt-layout", + "document_annotation": None, + "usage_info": {"pages_processed": 0}, + "object": "ocr", + "provider_native_response": provider_response, + } + ) + + assert response.get_provider_native_response() == provider_response + assert response.model_dump().get("provider_native_response") is None + + +def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: + error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.RateLimitError) + assert public_error.status_code == 429 + assert public_error.response.headers["retry-after"] == "7" + assert public_error.response.text == '{"message": "slow down"}' + assert public_error.__context__ is error + assert public_error.llm_provider == "mistral" + + +def test_map_failure_maps_upstream_401_to_authentication_error() -> None: + error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.AuthenticationError) + assert public_error.status_code == 401 + assert public_error.response.text == '{"message": "Unauthorized"}' + assert public_error.__context__ is error + + +def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: + error: Final = RuntimeError("bridge exploded") + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert not isinstance(public_error, UpstreamFailure) + assert isinstance(public_error, litellm.APIConnectionError) + assert "bridge exploded" in str(public_error) + + +def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): + raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/unit/rust_bridge/responses/__init__.py b/tests/unit/rust_bridge/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/responses/test_route_host.py b/tests/unit/rust_bridge/responses/test_route_host.py new file mode 100644 index 00000000000..49bf19e7d8a --- /dev/null +++ b/tests/unit/rust_bridge/responses/test_route_host.py @@ -0,0 +1,57 @@ +from types import MappingProxyType +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.rust_bridge.responses.route_host import arguments, response +from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest +from litellm.types.llms.openai import ResponsesAPIResponse + + +def test_response_validates_into_the_public_responses_model() -> None: + built: Final = response( + MappingProxyType( + { + "id": "resp_native", + "object": "response", + "created_at": 1, + "model": "gpt-4o", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_native", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "native", "annotations": []}], + } + ], + } + ) + ) + + assert isinstance(built, ResponsesAPIResponse) + assert built.id == "resp_native" + assert built.output[0].content[0].text == "native" + + +def test_response_rejects_a_payload_missing_required_fields() -> None: + with pytest.raises(ValidationError): + response(MappingProxyType({"object": "response"})) + + +def test_arguments_are_the_public_kwargs_view() -> None: + kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) + request: Final = LiteLLMResponsesRequest( + model="gpt-4o", + input="hi", + stream=None, + api_key=None, + api_base=None, + custom_llm_provider="openai", + extra_headers=None, + kwargs=kwargs, + ) + + assert arguments(request) is kwargs diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py new file mode 100644 index 00000000000..e01b9120416 --- /dev/null +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -0,0 +1,318 @@ +""" +Tests for the e2b code execution sandbox primitive. + +Unit tests inject a fake async HTTP client (dependency injection, no +monkeypatching) and assert request shapes and result mapping. Real-network +integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. +""" + +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.e2b.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + E2BSandboxConfig, +) + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + """Records outbound requests and returns canned responses keyed by URL.""" + + def __init__( + self, + *, + create_json=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "sandboxID": "sbx_123", + "domain": "e2b.app", + "envdAccessToken": "tok_abc", + } + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json)) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/execute"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None)) + if not (200 <= self.delete_status < 300): + raise httpx.HTTPStatusError( + f"status {self.delete_status}", + request=httpx.Request("DELETE", url), + response=httpx.Response(self.delete_status), + ) + return FakeResponse(status_code=self.delete_status) + + +# ---------- pure parser ---------- + + +def test_parse_lines_stdout_and_count(): + lines = [ + json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), + json.dumps({"type": "number_of_executions", "execution_count": 1}), + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "6\n" + assert result.execution_count == 1 + assert result.error is None + + +def test_parse_lines_error_surfaces_name_and_traceback(): + lines = [ + json.dumps( + { + "type": "error", + "name": "ZeroDivisionError", + "value": "division by zero", + "traceback": "Traceback (most recent call last): ...", + } + ) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.error["name"] == "ZeroDivisionError" + assert "Traceback" in result.error["traceback"] + + +def test_parse_lines_result_carries_png(): + lines = [ + json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.results and result.results[0]["png"] == "BASE64DATA" + assert "type" not in result.results[0] + + +# ---------- request shapes ---------- + + +@pytest.mark.asyncio +async def test_template_flows_into_create_request_as_templateID(): + client = FakeHTTPClient() + cfg = E2BSandboxConfig() + handle = await cfg.acreate_sandbox( + template="my-custom-template", api_key="e2b_key", client=client + ) + + method, url, headers, body = client.calls[0] + assert method == "POST" + assert url.endswith("/sandboxes") + assert body["templateID"] == "my-custom-template" # not "template" + assert body["secure"] is True + assert headers["X-API-Key"] == "e2b_key" + assert handle.id == "sbx_123" + assert handle._hidden_params["envd_access_token"] == "tok_abc" + + +@pytest.mark.asyncio +async def test_create_defaults_template_when_omitted(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) + _, _, _, body = client.calls[0] + assert body["templateID"] == "code-interpreter-v1" + + +@pytest.mark.asyncio +async def test_run_code_targets_jupyter_host_with_access_token(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] + ) + handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok_run"} + + result = await E2BSandboxConfig().arun_code( + container=handle, code="print(6*7)", client=client + ) + + method, url, headers, body = client.calls[0] + assert url == "https://49999-sbx_xyz.e2b.app/execute" + assert headers["X-Access-Token"] == "tok_run" + assert body["code"] == "print(6*7)" + assert result.stdout.strip() == "42" + + +@pytest.mark.asyncio +async def test_delete_issues_delete_to_sandbox_id(): + client = FakeHTTPClient(delete_status=204) + handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + method, url, headers, _ = client.calls[0] + assert method == "DELETE" + assert url.endswith("/sandboxes/sbx_del") + assert ok is True + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + assert ok is False + + +# ---------- ephemeral teardown ---------- + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="e2b", code="1/0", api_key="e2b_key", client=client + ) + + methods = [c[0] for c in client.calls] + urls = [c[1] for c in client.calls] + assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete + assert urls[0].endswith("/sandboxes") + assert urls[1].endswith("/execute") + assert urls[2].endswith("/sandboxes/sbx_123") + + +# ---------- correctness guards ---------- + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + with pytest.raises(httpx.HTTPStatusError): + await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + +@pytest.mark.asyncio +async def test_create_preserves_explicit_zero_timeout(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + timeout=0, api_key="e2b_key", client=client + ) + _, _, _, body = client.calls[0] + assert body["timeout"] == 0 + + +@pytest.mark.asyncio +async def test_run_code_rejects_bare_id_without_access_token(): + client = FakeHTTPClient() + with pytest.raises(ValueError, match="access token"): + await E2BSandboxConfig().arun_code( + container="sbx_no_token", code="print(1)", client=client + ) + assert client.calls == [] # never reached the network + + +def test_parse_lines_skips_non_json_lines(): + lines = [ + "not-json-heartbeat", + json.dumps({"type": "stdout", "text": "ok\n"}), + "", + "{partial", + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "ok\n" + assert result.error is None + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + big_line = "x" * (MAX_OUTPUT_BYTES + 1) + client = FakeHTTPClient(execute_lines=[big_line]) + handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + with pytest.raises(ValueError, match="exceeded"): + await E2BSandboxConfig().arun_code( + container=handle, code="print('x'*999)", client=client + ) + + +# ---------- public entrypoints ---------- + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] + ) + container = await litellm.acreate_sandbox( + provider="e2b", api_key="e2b_key", client=client + ) + assert container.id == "sbx_123" + + result = await litellm.arun_code( + provider="e2b", + container=container, + api_key="e2b_key", + code="print(6*7)", + client=client, + ) + assert result.stdout.strip() == "42" + + assert ( + await litellm.adelete_sandbox( + provider="e2b", container=container, api_key="e2b_key", client=client + ) + is True + ) + + +@pytest.mark.asyncio +async def test_unsupported_provider_raises(): + with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): + await litellm.acreate_sandbox(provider="not-a-provider") + + +# ---------- api_base override ---------- + + +@pytest.mark.asyncio +async def test_create_uses_api_base_override(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + api_base="http://my-sandbox:8080", api_key="k", client=client + ) + _, url, _, _ = client.calls[0] + assert url == "http://my-sandbox:8080/sandboxes" + + +@pytest.mark.asyncio +async def test_create_defaults_to_e2b_api_base(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) + _, url, _, _ = client.calls[0] + assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py new file mode 100644 index 00000000000..2928dea100e --- /dev/null +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -0,0 +1,647 @@ +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.opensandbox.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OpenSandboxSandboxConfig, +) +from litellm.utils import ProviderConfigManager + +TEST_API_BASE = "https://sandbox.test/v1" + + +def http_status_error(status_code, url="http://test"): + return httpx.HTTPStatusError( + f"status {status_code}", + request=httpx.Request("GET", url), + response=httpx.Response(status_code), + ) + + +def sse(data): + return f"data: {json.dumps(data)}" + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + raise http_status_error(self.status_code) + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + def __init__( + self, + *, + create_json=None, + sandbox_states=None, + endpoint_json=None, + endpoint_responses=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + self.sandbox_states = list( + sandbox_states + or [ + { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + ] + ) + self.endpoint_json = endpoint_json or { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + } + self.endpoint_responses = ( + list(endpoint_responses) if endpoint_responses is not None else None + ) + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json, {"stream": stream})) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/code"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def get(self, url, headers=None, params=None, **kwargs): + self.calls.append(("GET", url, headers, None, params)) + if "/endpoints/44772" in url: + if self.endpoint_responses is not None and self.endpoint_responses: + response = self.endpoint_responses.pop(0) + if isinstance(response, Exception): + raise response + if isinstance(response, FakeResponse): + return response + return FakeResponse(json_data=response) + return FakeResponse(json_data=self.endpoint_json) + if "/sandboxes/" in url: + state = self.sandbox_states.pop(0) + return FakeResponse(json_data=state) + raise AssertionError(f"unexpected GET {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None, None)) + if not (200 <= self.delete_status < 300): + raise http_status_error(self.delete_status, url) + return FakeResponse(status_code=self.delete_status) + + +def test_parse_sse_lines_maps_output_result_count_and_error(): + lines = [ + sse({"type": "stdout", "text": "hello\n"}), + sse({"type": "stderr", "text": "warn\n"}), + sse({"type": "result", "results": {"text/plain": "4"}}), + sse({"type": "execution_count", "execution_count": 7}), + sse( + { + "type": "error", + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback"], + }, + } + ), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "hello\n" + assert result.stderr == "warn\n" + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 7 + assert result.error == { + "name": "ValueError", + "value": "bad", + "traceback": ["Traceback"], + } + + +def test_parse_sse_lines_skips_non_json_and_control_lines(): + lines = [ + "event: message", + "not-json", + "", + sse({"type": "stdout", "text": "ok\n"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "ok\n" + assert result.error is None + + +def test_parse_sse_lines_maps_fallback_shapes(): + lines = [ + "data:", + sse(["not-a-dict"]), + sse({"code": "BadRequest", "message": "nope"}), + sse({"type": "result", "text/plain": "4"}), + sse({"type": "error", "name": "RuntimeError", "text": "boom"}), + sse({"type": "execution_count", "execution_count": "8"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 8 + assert result.error == { + "name": "BadRequest", + "value": "nope", + "traceback": [], + } + fallback_error = OpenSandboxSandboxConfig._parse_lines( + [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] + ) + assert fallback_error.error == { + "name": "RuntimeError", + "value": "boom", + "traceback": [], + } + empty_string_error = OpenSandboxSandboxConfig._parse_lines( + [ + sse( + { + "type": "error", + "error": { + "ename": "", + "name": "FallbackName", + "evalue": "", + "value": "fallback value", + "traceback": [], + }, + } + ) + ] + ) + assert empty_string_error.error == { + "name": "", + "value": "", + "traceback": [], + } + + +def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): + def fake_secret(key): + if key == "OPEN_SANDBOX_API_KEY": + return "env-key" + if key == "OPEN_SANDBOX_API_BASE": + return TEST_API_BASE + return None + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + fake_secret, + ) + config = OpenSandboxSandboxConfig() + handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") + + assert config.validate_environment() == "env-key" + assert config.validate_environment(api_key="") == "" + assert config._api_key(api_key=None, handle=handle) == "env-key" + + handle._hidden_params = {"api_key": "stored-key"} + assert config._api_key(api_key=None, handle=handle) == "stored-key" + assert config._http(None) is not None + + body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy={"egress": [{"domain": "example.com"}]}, + secure_access=True, + ) + assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} + assert body["secureAccess"] is True + + other_body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy=None, + secure_access=False, + ) + assert body["resourceLimits"] is not other_body["resourceLimits"] + + assert config._sandbox_state(None) is None + assert config._sandbox_state({"status": "Running"}) is None + assert config._as_str_dict(None) == {} + assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( + "http://execd.local" + ) + assert config._api_base(None) == TEST_API_BASE + assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" + assert config._as_int("9") == 9 + assert config._as_int("nope") is None + assert config._as_int(None) is None + assert isinstance( + ProviderConfigManager.get_provider_sandbox_config("opensandbox"), + OpenSandboxSandboxConfig, + ) + + +def test_api_base_requires_kwarg_or_env(monkeypatch): + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + lambda key: None, + ) + + with pytest.raises(ValueError, match="api_base is required"): + OpenSandboxSandboxConfig._api_base(None) + + +@pytest.mark.asyncio +async def test_create_posts_default_body_and_omits_empty_api_key(): + client = FakeHTTPClient() + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + method, url, headers, body, _ = client.calls[0] + assert method == "POST" + assert url == f"{TEST_API_BASE}/sandboxes" + assert "OPEN-SANDBOX-API-KEY" not in headers + assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} + assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] + assert body["timeout"] == 300 + assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert handle.id == "osb_123" + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + + +@pytest.mark.asyncio +async def test_create_can_opt_into_internet_access(): + client = FakeHTTPClient() + + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + allow_internet_access=True, + client=client, + ) + + _, _, _, body, _ = client.calls[0] + assert "networkPolicy" not in body + + +@pytest.mark.asyncio +async def test_create_custom_options_poll_and_endpoint_resolution(): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + }, + sandbox_states=[ + { + "id": "osb_pending", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + } + ], + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + template="custom/image:latest", + timeout=600, + allow_internet_access=False, + api_key="osb-key", + api_base="https://sandbox.example/v1", + metadata={"suite": "unit"}, + env_vars={"PYTHONUNBUFFERED": "1"}, + resource_limits={"cpu": "500m", "memory": "512Mi"}, + resource_requests={"cpu": "250m", "memory": "256Mi"}, + entrypoint=["/bin/sh", "-lc", "sleep 3600"], + use_server_proxy=True, + client=client, + ) + + _, create_url, create_headers, body, _ = client.calls[0] + _, poll_url, poll_headers, _, _ = client.calls[1] + _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] + + assert create_url == "https://sandbox.example/v1/sandboxes" + assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert body["image"] == {"uri": "custom/image:latest"} + assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] + assert body["metadata"] == {"suite": "unit"} + assert body["env"] == {"PYTHONUNBUFFERED": "1"} + assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} + assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" + assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") + assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_params == {"use_server_proxy": True} + assert handle.id == "osb_pending" + + +@pytest.mark.asyncio +async def test_create_waits_across_pending_state(monkeypatch): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + }, + sandbox_states=[ + {"id": "osb_pending", "status": {"state": "Pending"}}, + {"id": "osb_pending", "status": {"state": "Running"}}, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + assert handle.id == "osb_pending" + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_for_terminal_state(): + client = FakeHTTPClient( + create_json={"id": "osb_failed", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_failed", "status": {"state": "Failed"}}, + ], + ) + + with pytest.raises(ValueError, match="entered Failed"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_create_times_out_waiting_for_running(): + client = FakeHTTPClient( + create_json={"id": "osb_slow", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_slow", "status": {"state": "Pending"}}, + ], + ) + + with pytest.raises(TimeoutError, match="was not Running"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=0, + poll_interval=0, + client=client, + ) + + +@pytest.mark.asyncio +async def test_create_waits_for_endpoint_resolution(monkeypatch): + client = FakeHTTPClient( + endpoint_responses=[ + http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), + { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + }, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + assert len(endpoint_calls) == 2 + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_when_endpoint_is_missing(): + client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) + + with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client + ) + + +@pytest.mark.asyncio +async def test_create_reraises_non_404_endpoint_error(): + client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_run_code_resolves_bare_id_and_posts_sse_request(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + result = await OpenSandboxSandboxConfig().arun_code( + container="osb_bare", + code="print(6*7)", + language="python", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + endpoint_call = client.calls[0] + run_call = client.calls[1] + assert endpoint_call[0] == "GET" + assert ( + endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" + ) + assert run_call[0] == "POST" + assert run_call[1] == "http://execd.local:44772/code" + assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" + assert run_call[3] == { + "code": "print(6*7)", + "context": {"language": "python"}, + } + assert run_call[4] == {"stream": True} + assert result.stdout == "42\n" + + +@pytest.mark.asyncio +async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): + client = FakeHTTPClient() + handle = ContainerHandle( + id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" + ) + handle._hidden_params = { + "execd_endpoint": "execd.example/route/44772", + "execd_headers": {}, + } + + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print(1)", client=client + ) + + assert client.calls[0][1] == "https://execd.example/route/44772/code" + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) + handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") + handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} + + with pytest.raises(ValueError, match="exceeded"): + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print('x')", client=client + ) + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + + ok = await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_gone", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + assert ok is False + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_err", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + container = await litellm.acreate_sandbox( + provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client + ) + result = await litellm.arun_code( + provider="opensandbox", + container=container, + code="print(6*7)", + api_key="", + client=client, + ) + ok = await litellm.adelete_sandbox( + provider="opensandbox", + container=container, + api_key="", + client=client, + ) + + assert container.id == "osb_123" + assert result.stdout == "42\n" + assert ok is True + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="opensandbox", + code="1/0", + api_key="", + api_base=TEST_API_BASE, + client=client, + ) + + assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] + assert client.calls[0][1].endswith("/sandboxes") + assert client.calls[1][1].endswith("/endpoints/44772") + assert client.calls[2][1].endswith("/code") + assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/unit/sandbox/test_sandbox_tools.py b/tests/unit/sandbox/test_sandbox_tools.py new file mode 100644 index 00000000000..06136534b13 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_tools.py @@ -0,0 +1,181 @@ +"""Unit tests for the sandbox-tool registry.""" + +from litellm.sandbox import sandbox_tools + + +def _reset(): + sandbox_tools.clear_sandbox_tools() + + +def test_register_resolves_provider_key_and_base(): + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved == { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + } + _reset() + + +def test_register_clears_stale_entries_on_reload(): + """A tool removed from the config must not survive a re-registration.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "old", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("old") is not None + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "new", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("new") is not None + assert ( + sandbox_tools.resolve_sandbox_tool("old") is None + ), "stale tool must be gone after the config is reloaded" + _reset() + + +def test_register_empty_list_clears_removed_tools(): + """Reloading a config with sandbox_tools removed (the proxy passes an empty + list) must drop previously registered credentials from the process.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None + + sandbox_tools.register_sandbox_tools([]) + + assert ( + sandbox_tools.resolve_sandbox_tool("e2b_default") is None + ), "removing sandbox_tools from config must clear stale credentials" + _reset() + + +def test_register_resolves_secret_from_env(monkeypatch): + _reset() + monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "os.environ/MY_SANDBOX_KEY", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved is not None + assert resolved["api_key"] == "sk-from-env" + assert resolved["api_base"] is None + _reset() + + +def test_resolve_unknown_returns_none(): + _reset() + assert sandbox_tools.resolve_sandbox_tool("nope") is None + + +def test_register_skips_malformed_entries_without_crashing(): + """A single malformed entry (missing sandbox_tool_name, or not a dict) must + not crash registration during proxy startup/hot-reload; valid entries in the + same list must still register.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name + "not-a-dict", # wrong type + {"sandbox_tool_name": "", "litellm_params": {}}, # empty name + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("good") is not None + assert sandbox_tools.resolve_sandbox_tool("") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_skips_entry_missing_sandbox_provider(): + """An entry with a name but no sandbox_provider must be skipped at + registration so it cannot later resolve and call acreate_sandbox(provider=None), + which fails with a cryptic runtime error instead of a clear startup warning.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, + { + "sandbox_tool_name": "null_provider", + "litellm_params": {"sandbox_provider": None}, + }, + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("no_provider") is None + assert sandbox_tools.resolve_sandbox_tool("null_provider") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_swaps_registry_atomically(): + """register_sandbox_tools must replace the registry in one rebind so a + concurrent resolve never observes a half-populated or transiently empty + registry between clearing and repopulating.""" + _reset() + sandbox_tools.register_sandbox_tools( + [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] + ) + before = sandbox_tools._SANDBOX_TOOL_REGISTRY + + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, + {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, + ] + ) + after = sandbox_tools._SANDBOX_TOOL_REGISTRY + + assert after is not before, "the registry must be replaced, not mutated in place" + assert set(after) == {"b", "c"} + assert "a" not in after + _reset() diff --git a/tests/unit/skills/test_skills_main.py b/tests/unit/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/unit/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py new file mode 100644 index 00000000000..7577064b7f9 --- /dev/null +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -0,0 +1,468 @@ +""" +Tests for enforce_model_rate_limits feature. + +This feature allows users to enforce TPM/RPM limits set on model deployments +regardless of the routing strategy being used. +""" + +import asyncio +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) + +TPM_DEPLOYMENT = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "replica-test-id"}, + "model_name": "test-model", +} + + +def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache) + now = litellm.utils.get_utc_datetime() + for minute in (now, now + timedelta(minutes=1)): + tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) + dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) + return dual_cache + + +class TestModelRateLimitingCheck: + """Test the ModelRateLimitingCheck class directly.""" + + def test_get_deployment_limits_from_top_level(self): + """Test extracting limits from top-level deployment config.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "tpm": 1000, + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 1000 + assert rpm == 10 + + def test_get_deployment_limits_from_litellm_params(self): + """Test extracting limits from litellm_params.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 2000 + assert rpm == 20 + + def test_get_deployment_limits_from_model_info(self): + """Test extracting limits from model_info.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 3000 + assert rpm == 30 + + def test_get_deployment_limits_none_when_not_set(self): + """Test that None is returned when limits are not set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm is None + assert rpm is None + + def test_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 11 # Over limit after increment + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + assert "current usage=11" in str(exc_info.value) + + def test_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 6 + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 1000 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 1000 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.parametrize( + "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] + ) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() + redis_cache.increment_cache.return_value = 2 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert check.pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + def test_log_success_event_increments_cache(self): + """Test that log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + check.log_success_event(kwargs, None, None, None) + + # Verify increment_cache was called + mock_cache.increment_cache.assert_called_once() + _, kwarg_params = mock_cache.increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestModelRateLimitingCheckAsync: + """Test async methods of ModelRateLimitingCheck.""" + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=6) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=1000) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] + ) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.async_get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( + self, + ): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) + redis_cache.async_increment = AsyncMock(return_value=2) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert await check.async_pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_log_success_event_increments_cache(self): + """Test that async_log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + # Verify async_increment_cache was called + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestRouterWithEnforceModelRateLimits: + """Test Router integration with enforce_model_rate_limits.""" + + def test_router_initializes_with_enforce_model_rate_limits(self): + """Test that Router properly initializes the ModelRateLimitingCheck.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + router = Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Check that the callback was added + assert router.optional_callbacks is not None + assert len(router.optional_callbacks) == 1 + assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + + def test_router_optional_callbacks_contains_model_rate_limiting(self): + """Test that ModelRateLimitingCheck is in the callbacks list.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Find the ModelRateLimitingCheck in litellm.callbacks + found = False + for callback in litellm.callbacks: + if isinstance(callback, ModelRateLimitingCheck): + found = True + break + + assert found, "ModelRateLimitingCheck should be in litellm.callbacks" + + +class TestModelRateLimitConcurrency: + """Test that RPM rate limiting is atomic under concurrent requests.""" + + @pytest.mark.asyncio + async def test_concurrent_requests_respect_rpm_limit(self): + """ + Fire 4 concurrent async requests with RPM limit of 2. + Exactly 2 should succeed and 2 should raise RateLimitError. + + This test validates the atomic increment-first pattern: + the old check-then-increment pattern would let 3+ through + due to a race condition on the local cache read. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + + deployment = { + "rpm": 2, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "concurrent-test-id"}, + "model_name": "test-model", + } + + async def attempt_request(): + return await check.async_pre_call_check(deployment) + + results = await asyncio.gather( + *[attempt_request() for _ in range(4)], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + failures = [r for r in results if isinstance(r, litellm.RateLimitError)] + + assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" + assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py new file mode 100644 index 00000000000..a5a68271111 --- /dev/null +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -0,0 +1,1041 @@ +""" +Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). +""" + +import asyncio + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, + _reservation_value, + _resolve_max_tokens, + async_io_token_pre_call_check, + async_io_token_reconcile_success, + build_io_token_rate_limit_headers, + deployment_has_io_token_limits, + get_io_token_rate_limit_request_kwargs, + io_token_reconcile_success, + io_token_refund_failure, + refund_stale_reservation_before_retry, + set_io_token_rate_limit_request_kwargs, +) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) +from litellm.types.utils import ModelResponse, Usage + + +class TestIOTokenRateLimitHelpers: + def test_deployment_has_io_token_limits(self): + assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) + assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) + + def test_reservation_value_minimal_when_estimate_fails(self): + # A failed/empty estimate (0) must reserve a minimal slot, not the + # entire limit - otherwise one request whose estimate failed fills + # the whole bucket and blocks every concurrent request until it + # completes and reconciles. + assert _reservation_value(0, 100) == 1 + assert _reservation_value(0, 1) == 1 + # A real non-zero estimate is reserved as-is. + assert _reservation_value(42, 100) == 42 + + def test_resolve_max_tokens_respects_explicit_zero(self): + deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} + # An explicit max_tokens=0 is honored, not replaced by the model default. + assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 + # max_completion_tokens is the fallback only when max_tokens is absent. + assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 + assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + + def test_build_io_token_rate_limit_headers(self): + headers = build_io_token_rate_limit_headers( + itpm_limit=200, + otpm_limit=40, + current_itpm=15, + current_otpm=4, + ) + assert headers["x-ratelimit-limit-input-tokens"] == 200 + assert headers["x-ratelimit-remaining-input-tokens"] == 185 + assert headers["x-ratelimit-limit-output-tokens"] == 40 + assert headers["x-ratelimit-remaining-output-tokens"] == 36 + + +class TestModelRateLimitingCheckIOTokens: + @pytest.mark.asyncio + async def test_itpm_reservation_and_reconcile(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-test-id"}, + "model_name": "opus", + } + + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + + kwargs = { + "standard_logging_object": { + "model_id": "io-test-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "index": 0, + "finish_reason": "stop", + } + ], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + await check.async_log_success_event(kwargs, response, None, None) + + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + # ITPM tracks input tokens only (billable prompt tokens), not output. + assert current_itpm == 5 + assert current_otpm == 3 + + @pytest.mark.asyncio + async def test_itpm_limit_raises_429(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 5, + }, + "model_info": {"id": "io-limit-id"}, + "model_name": "opus", + } + + # ITPM enforces input tokens only; the prompt alone must exceed the limit, + # a large max_tokens must not contribute to the ITPM reservation. + set_io_token_rate_limit_request_kwargs( + { + "messages": [ + { + "role": "user", + "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", + } + ], + "max_tokens": 10, + "metadata": {}, + } + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "ITPM limit=5" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + otpm_limit = 10 + max_tokens = 4 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": otpm_limit, + }, + "model_info": {"id": "io-otpm-race-id"}, + "model_name": "opus", + } + + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": max_tokens, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + + # Atomic reservation must never let concurrent requests overshoot the limit. + assert current_otpm is not None + assert current_otpm <= otpm_limit + assert successes == otpm_limit // max_tokens + assert current_otpm == successes * max_tokens + + @pytest.mark.asyncio + async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): + """ + When input-token estimation yields 0 (no messages/prompt/input field, + unsupported model, tokenizer error), the reservation must be a + minimal 1 token, not the entire itpm limit. Otherwise the first + request whose estimate fails fills the whole bucket and every + concurrent request is rejected until it completes - effectively + serializing traffic to the deployment. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_limit = 5 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": itpm_limit, + }, + "model_info": {"id": "io-itpm-estimate-fail-id"}, + "model_name": "opus", + } + + # No messages/prompt/input field -> _estimate_input_tokens returns 0. + set_io_token_rate_limit_request_kwargs( + { + "max_tokens": 5, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + + # A minimal 1-token reservation per request lets itpm_limit concurrent + # requests through, instead of a single request starving the rest. + assert current_itpm is not None + assert current_itpm <= itpm_limit + assert successes == itpm_limit + + @pytest.mark.asyncio + async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + # Production kwargs commonly carry litellm_params.metadata; the stashed + # reservation lives in the top-level metadata and must still be found. + kwargs = { + "standard_logging_object": { + "model_id": "io-lp-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + @pytest.mark.asyncio + async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + }, + "model_info": {"id": "io-zero-est-id"}, + "model_name": "opus", + } + + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A failed/zero estimate reserves a minimal 1 token, not the full + # itpm limit, so it doesn't starve concurrent requests. + assert await dual_cache.async_get_cache(key=itpm_key) == 1 + + kwargs = { + "standard_logging_object": { + "model_id": "io-zero-est-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + + @pytest.mark.asyncio + async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): + """ + A zero/failed estimate reserves a minimal 1 token rather than the + full itpm limit, so up to itpm_limit such calls are allowed + concurrently instead of the first one claiming the entire bucket. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 2, + }, + "model_info": {"id": "io-zero-cap-id"}, + "model_name": "opus", + } + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + # Second zero-estimate call still fits within the itpm=2 limit. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + await check.async_pre_call_check(deployment) + + # A third exceeds the limit and is rejected. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": 5, + }, + "model_info": {"id": "io-zero-output-id"}, + "model_name": "opus", + } + zero_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 0, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + await check.async_pre_call_check(deployment) + + zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 + assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + + normal_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + await check.async_pre_call_check(deployment) + + normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + + def test_sync_io_pre_call_reserves_and_reconciles(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-sync-id"}, + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + check.pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + kwargs = { + "standard_logging_object": { + "model_id": "io-sync-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 5 + assert dual_cache.get_cache(key=otpm_key) == 3 + + @pytest.mark.asyncio + async def test_reconcile_runs_via_success_event_without_model_id(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + # standard_logging_object has no model_id (only the TPM path needs it); + # IO reconciliation must still run off the stashed cache key. + kwargs = { + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + + @pytest.mark.asyncio + async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + + # Shared request metadata carrying the first (IO) deployment's reservation. + metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + fail_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "io-first", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + } + await check.async_log_failure_event(fail_kwargs, None, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + # Retry succeeds on a non-IO fallback deployment reusing the same metadata. + retry_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "non-io-second", + "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, + "metadata": {}, + "total_tokens": 12, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), + ) + await check.async_log_success_event(retry_kwargs, response, None, None) + + # The first deployment's ITPM counter is not driven negative... + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + # ...and the non-IO deployment's TPM usage is tracked normally. + tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" + assert await dual_cache.async_get_cache(key=tpm_key) == 12 + + @pytest.mark.asyncio + async def test_stale_reservation_refunded_before_retry_overwrites_it(self): + """ + A retry reuses the same mutable kwargs dict for the next deployment. + If deployment A's failure event hasn't run yet (e.g. it was scheduled + as a background task) when the retry calls + set_io_token_rate_limit_request_kwargs for deployment B, the router + must first synchronously refund + clear A's reservation via + refund_stale_reservation_before_retry - otherwise A's counter stays + elevated by the reservation until its TTL expires, and the + now-orphaned sentinels must not leak into B's accounting either. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + minute = get_utc_datetime().strftime("%H-%M") + itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + + # Deployment A's still-unreconciled reservation, stashed on the shared + # kwargs dict the retry loop reuses. + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + + # Router calls this before overwriting kwargs for deployment B's attempt - + # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. + refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + + # A's reservation is refunded immediately, not left stranded for a + # background failure task that may run arbitrarily later (or never, + # if the sentinels get cleared out from under it first). + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] + assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + + # A's own (now-late) failure event finds nothing left to refund and + # is a safe no-op, since the sentinels were already cleared above. + io_token_refund_failure(dual_cache, shared_kwargs) + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + + # The retry proceeds to stash deployment B's own reservation on the + # same dict; it starts clean, unaffected by A's cleared sentinels. + set_io_token_rate_limit_request_kwargs(shared_kwargs) + itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 + shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b + await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) + assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + + @pytest.mark.asyncio + async def test_client_supplied_reservation_keys_are_stripped(self): + # metadata is caller-controlled; the server-only reservation sentinels + # must be removed before the router captures the request kwargs. + forged = { + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "litellm_metadata": {OTPM_RESERVED_KEY: 7}, + "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + } + set_io_token_rate_limit_request_kwargs(forged) + stored = get_io_token_rate_limit_request_kwargs() + + assert ITPM_RESERVED_KEY not in stored["metadata"] + assert ITPM_CACHE_KEY not in stored["metadata"] + assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] + assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + + @pytest.mark.asyncio + async def test_forged_reservation_cannot_decrement_counter(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + victim_key = "global_router:victim:model:itpm:00-00" + await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + + # A caller forges a reservation pointing at another deployment's counter. + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "standard_logging_object": { + "model_id": "m", + "hidden_params": {"litellm_model_name": "model"}, + "metadata": {}, + "total_tokens": 2, + }, + } + # The router sanitizes the request kwargs before the call runs. + set_io_token_rate_limit_request_kwargs(kwargs) + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + await check.async_log_success_event(kwargs, response, None, None) + + # The forged reservation was stripped, so the victim counter is untouched. + assert await dual_cache.async_get_cache(key=victim_key) == 100 + + @pytest.mark.asyncio + async def test_otpm_reservation_error_rolls_back_itpm(self): + from litellm.utils import get_utc_datetime + + class _OtpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":otpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _OtpmFailCache() + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 1000, + "otpm": 1000, + }, + "model_info": {"id": "io-rollback-id"}, + "model_name": "opus", + } + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + ) + + with pytest.raises(RuntimeError): + await async_io_token_pre_call_check(dual_cache, deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A transient OTPM error must release the ITPM reservation, not leak it. + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + + @pytest.mark.asyncio + async def test_reconcile_clears_stash_even_when_increment_errors(self): + class _ItpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":itpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _ItpmFailCache() + metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + kwargs = {"metadata": metadata} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + with pytest.raises(RuntimeError): + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + # The stash is cleared even though reconciliation raised, so a duplicate + # success event can't re-process it. + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + @pytest.mark.asyncio + async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): + import logging + + check = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, + "model_info": {}, + } + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + + warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] + # id-less deployments are not collapsed onto a single dedup key. + assert len(warnings) == 2 + + @pytest.mark.asyncio + async def test_missing_deployment_id_skips_io_reservation(self): + dual_cache = DualCache() + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, + "model_info": {}, # no id -> cannot build a per-deployment cache key + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + result = await async_io_token_pre_call_check(dual_cache, deployment) + + assert result is deployment + # No reservation is stashed, so nothing lands in a shared None:None bucket. + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_uses_reservation_minute_key(self): + dual_cache = DualCache() + # Reservation was made on a fixed minute key; a call that finishes in a + # later minute must reconcile against that same key, never a key built + # from the response-time minute. + itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + @pytest.mark.asyncio + async def test_reconcile_missing_usage_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + ) + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_total_tokens_only_keeps_reservation(self): + """ + A response usage object with only total_tokens (no prompt/completion + breakdown) can't be split into input/output, so it must be treated the + same as missing usage: keep the reservation instead of resolving to + (0, 0) and refunding it in full. + """ + dual_cache = DualCache() + itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = {"type": "message", "usage": {"total_tokens": 13}} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + + def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": {"total_tokens": 4}, + } + response = {"type": "message", "role": "assistant", "content": []} + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 10 + + @pytest.mark.asyncio + async def test_reconcile_falls_back_to_standard_logging_object(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "prompt_tokens": 4, + "completion_tokens": 0, + "total_tokens": 4, + }, + } + response = {"type": "message", "role": "assistant", "content": []} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + def test_sync_reconcile_anthropic_dict_usage(self): + dual_cache = DualCache() + itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" + dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 6, + OTPM_RESERVED_KEY: 4, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = { + "type": "message", + "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + } + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=otpm_key) == 2 + + @pytest.mark.asyncio + async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): + import logging + + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + deployment = { + "litellm_params": { + "model": deployment_name, + "itpm": 100, + "rpm": 1, + }, + "model_info": {"id": model_id}, + "model_name": "opus", + } + + minute = get_utc_datetime().strftime("%H-%M") + rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + + request_kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + assert any("both limit types are enforced" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): + """ + A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on + success, otherwise the tpm_key the pre-call check reads is never written + and the tpm_limit can never be enforced. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + await check.async_log_success_event(kwargs, response, None, None) + + # ITPM reconciled down from the 5-token reservation to actual usage (3). + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. + assert await dual_cache.async_get_cache(key=tpm_key) == 7 + + def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-sync-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 3 + assert dual_cache.get_cache(key=tpm_key) == 7 + + @pytest.mark.asyncio + async def test_failure_refunds_itpm_reservation(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + kwargs = { + "standard_logging_object": { + "model_id": "io-refund-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": dict(reservation), + }, + "metadata": dict(reservation), + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + +class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + async def test_model_group_info_aggregates_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "opus", + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 20, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + info = router.get_model_group_info("opus") + assert info is not None + assert info.itpm == 100 + assert info.otpm == 20 + + +class TestContextSlotRetention: + def test_setter_stores_kwargs_only_for_io_limited_deployments(self): + """ + The context slot pins the entire request kwargs (messages included) + for the lifetime of the surrounding asyncio context, and pooled + resources created mid-request (e.g. redis connections) capture that + context, extending the pin far past the request. Only ITPM/OTPM + pre-call checks read the slot, so the setter must store None for + deployments without io token limits and still clear reservation + sentinels from kwargs either way. + """ + kwargs = { + "messages": [{"role": "user", "content": "x" * 1000}], + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + } + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + assert ITPM_CACHE_KEY not in kwargs["metadata"] + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + assert get_io_token_rate_limit_request_kwargs() is kwargs + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_does_not_pin_kwargs_without_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("plain") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_pins_kwargs_for_io_limited_deployment(self): + router = Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("limited") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is kwargs diff --git a/tests/unit/types/llms/test_types_llms_bedrock.py b/tests/unit/types/llms/test_types_llms_bedrock.py new file mode 100644 index 00000000000..a5ad882e775 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_bedrock.py @@ -0,0 +1,46 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams + + +def test_model_validate_keeps_auth_params_and_ignores_request_params(): + auth_params = AwsAuthParams.model_validate( + { + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", + "aws_session_name": "litellm-session", + "aws_external_id": "litellm-external-id", + "aws_region_name": "us-west-2", + "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "temperature": 0.1, + "messages": [{"role": "user", "content": "hi"}], + } + ) + + assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" + assert auth_params.aws_session_name == "litellm-session" + assert auth_params.aws_external_id == "litellm-external-id" + assert auth_params.aws_access_key_id is None + assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) + assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("aws_role_name", 1234), + ("aws_session_name", ["litellm-session"]), + ("aws_external_id", {"id": "x"}), + ], +) +def test_model_validate_rejects_non_string_credentials(field, value): + with pytest.raises(ValidationError): + AwsAuthParams.model_validate({field: value}) + + +def test_frozen_struct_rejects_field_assignment(): + auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") + + with pytest.raises(ValidationError): + auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py new file mode 100644 index 00000000000..64ec09838e8 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -0,0 +1,591 @@ +import asyncio +from typing import Optional +from unittest.mock import AsyncMock, patch + +import pytest + +import json + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +@pytest.mark.parametrize("stream", (False, True)) +def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: + from typing import Final + + from litellm.types.llms.openai import ( + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ) + from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + ) + + reasoning_item: Final = ChatCompletionReasoningItem( + type="reasoning", + id="rs_123", + encrypted_content="encrypted", + summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], + ) + response: Final = ( + ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) + if stream + else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) + ) + message_key: Final = "delta" if stream else "message" + assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + restored: Final = type(response).model_validate_json(response.model_dump_json()) + assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + +def test_generic_event(): + from litellm.types.llms.openai import GenericEvent + + event = {"type": "test", "test": "test"} + event = GenericEvent(**event) + assert event.type == "test" + assert event.test == "test" + + +def test_output_item_added_event(): + from litellm.types.llms.openai import OutputItemAddedEvent + + event = { + "type": "response.output_item.added", + "sequence_number": 4, + "output_index": 1, + "item": None, + } + event = OutputItemAddedEvent(**event) + assert event.type == "response.output_item.added" + assert event.sequence_number == 4 + assert event.output_index == 1 + assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" + + +class TestAssistantMessageImageUrlContent: + """ + Regression tests for image_url blocks in assistant message content. + + Bug: ChatCompletionAssistantMessage.content did not include + ChatCompletionImageObject in its union, so Pydantic v2 silently dropped + image_url blocks (content → []) when serialising via AllMessageValues. + This affects users who store conversation history as JSON (e.g. in a DB) + and read it back typed as list[AllMessageValues]. + """ + + ASSISTANT_MESSAGE_WITH_IMAGE = { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is the image you requested:"}, + { + "type": "image_url", + "image_url": { + "url": ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" + "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + }, + }, + ], + } + + def test_assistant_message_image_url_preserved_single(self): + """ + TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive + validate_python → dump_python without being dropped or raising an error. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import ChatCompletionAssistantMessage + + adapter = TypeAdapter(ChatCompletionAssistantMessage) + validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) + dumped = adapter.dump_python(validated) + + raw_content = dumped.get("content") + # Pydantic may return a lazy SerializationIterator for Iterable fields; + # convert to list to consume it — this must not raise ValidationError. + content_blocks = list(raw_content) if raw_content is not None else [] + + assert ( + len(content_blocks) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" + types = [b.get("type") for b in content_blocks if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped; blocks: {content_blocks}" + + def test_assistant_message_image_url_preserved_in_all_message_values(self): + """ + TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an + assistant message must not be silently dropped during dump_python(mode='json'). + + This is the primary failing path: conversation history stored as JSON in a + database and read back typed as list[AllMessageValues]. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import AllMessageValues + + conversation = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a LiteLLM costume", + }, + self.ASSISTANT_MESSAGE_WITH_IMAGE, + ] + + adapter = TypeAdapter(List[AllMessageValues]) + validated = adapter.validate_python(conversation) + dumped = adapter.dump_python(validated, mode="json") + + assistant = next((m for m in dumped if m.get("role") == "assistant"), None) + assert assistant is not None, "Assistant message missing after serialisation" + + content = assistant.get("content", []) + assert isinstance( + content, list + ), f"content should be a list, got {type(content)}" + assert ( + len(content) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" + types = [b.get("type") for b in content if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" + + +class TestResponsesAPIReasoningNullFields: + """ + Tests for issue #16824: reasoning output items should not include null + status/content/encrypted_content fields. + + When a provider returns reasoning items without these fields, LiteLLM's + Pydantic parsing adds them as Optional defaults (None). Serializing them + as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on + status=null). + + The fix uses a field_serializer on ResponsesAPIResponse.output that + mirrors the request-side filtering in + OpenAIResponsesAPIConfig._handle_reasoning_item(). + """ + + def _make_response(self, output): + from litellm.types.llms.openai import ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_test", + created_at=1741476542, + model="gpt-5-mini", + object="response", + status="completed", + output=output, + ) + + def test_reasoning_item_null_fields_removed_model_dump(self): + """Null status/content/encrypted_content should be absent from model_dump.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_null_fields_removed_model_dump_json(self): + """Null fields should also be absent from model_dump_json.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + parsed = json.loads(response.model_dump_json()) + reasoning = parsed["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_non_null_values_preserved(self): + """Non-null values on reasoning items should be kept.""" + response = self._make_response( + output=[ + { + "id": "rs_abc", + "type": "reasoning", + "summary": [], + "status": "completed", + "encrypted_content": "gAAAA...", + } + ] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["status"] == "completed" + assert reasoning["encrypted_content"] == "gAAAA..." + + def test_message_item_not_affected(self): + """Non-reasoning output items should keep all their fields.""" + response = self._make_response( + output=[ + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Hello!", + "annotations": [], + } + ], + } + ] + ) + dumped = response.model_dump() + message = dumped["output"][0] + assert message["status"] == "completed" + assert message["type"] == "message" + assert len(message["content"]) == 1 + + def test_mixed_output_reasoning_and_message(self): + """Reasoning items cleaned, message items untouched in same response.""" + response = self._make_response( + output=[ + {"id": "rs_abc", "type": "reasoning", "summary": []}, + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Answer", + "annotations": [], + } + ], + }, + ] + ) + dumped = response.model_dump() + reasoning = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "reasoning" + ][0] + message = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "message" + ][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert message["status"] == "completed" + assert len(message["content"]) == 1 + + def test_reasoning_core_fields_preserved(self): + """id, type, summary should always be present on reasoning items.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["id"] == "rs_abc" + assert reasoning["type"] == "reasoning" + assert reasoning["summary"] == ["thinking..."] + + def test_top_level_null_fields_unaffected(self): + """Top-level response fields with None should not be affected.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + assert "error" in dumped + assert dumped["error"] is None + assert "instructions" in dumped + assert dumped["instructions"] is None + + +def test_normalize_fine_tuning_job_dict_maps_azure_pending(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + out = _normalize_fine_tuning_job_dict( + {"organization_id": None, "result_files": None, "status": "pending"}, + is_azure=True, + ) + assert out["organization_id"] == "" + assert out["result_files"] == [] + assert out["status"] == "queued" + + +def test_normalize_fine_tuning_job_dict_openai_unchanged(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + data = {"organization_id": None, "result_files": None, "status": "pending"} + out = _normalize_fine_tuning_job_dict(data, is_azure=False) + assert out is data + + +def test_openai_file_object_accepts_pending_status(): + from litellm.types.llms.openai import OpenAIFileObject + + file_obj = OpenAIFileObject( + id="file-123", + bytes=1024, + created_at=1677610602, + filename="train.jsonl", + object="file", + purpose="fine-tune", + status="pending", + ) + assert file_obj.status == "pending" + + +class TestOpenAIFileObjectBatchGuardrailSerialization: + """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" + + @staticmethod + def _file_object(**overrides): + from litellm.types.llms.openai import OpenAIFileObject + + return OpenAIFileObject( + id="file-123", + object="file", + bytes=1024, + created_at=1677610602, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + **overrides, + ) + + @staticmethod + def _report(): + from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport + + return BatchGuardrailReport( + submitted_records=3, + modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), + ) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_absent_when_unset(self, mode): + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_present_when_set(self, mode): + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) + assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 + + def test_nested_nulls_of_a_set_report_survive(self): + """`exclude_none=True` was rejected as the fix because it would strip these.""" + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") + assert dumped["litellm_batch_guardrail"]["modified_records"] == [ + {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} + ] + + def test_by_alias_dump_also_omits_the_key(self): + """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) + + def test_other_optional_fields_still_serialize_as_null(self): + dumped = self._file_object().model_dump(mode="json") + assert dumped["expires_at"] is None + assert dumped["status_details"] is None + + def test_round_trip_of_a_set_report_is_lossless(self): + from litellm.types.llms.openai import OpenAIFileObject + + original = self._file_object(litellm_batch_guardrail=self._report()) + assert OpenAIFileObject(**original.model_dump()) == original + + def test_serialization_json_schema_still_describes_the_model(self): + """A return annotation on the wrap serializer would collapse this to a bare object.""" + from litellm.types.llms.openai import OpenAIFileObject + + schema = OpenAIFileObject.model_json_schema(mode="serialization") + assert "litellm_batch_guardrail" in schema["properties"] + + def test_key_omitted_inside_a_file_list_page(self): + from litellm.types.llms.openai import FileListPage + + page = FileListPage(object="list", data=[self._file_object()], has_more=False) + assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] + + +def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: + import httpx + + return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) + + +def test_httpx_binary_response_content_hidden_params_are_per_instance(): + first = _binary_content(b"first") + second = _binary_content(b"second") + + first._hidden_params["response_cost"] = 0.5 + + assert second._hidden_params == {} + + +def test_set_response_cost_none_leaves_hidden_params_empty(): + binary_response = _binary_content(b"audio") + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params + + binary_response.set_response_cost(0.25) + + assert binary_response._hidden_params["response_cost"] == 0.25 + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params diff --git a/tests/unit/types/proxy/policy_engine/__init__.py b/tests/unit/types/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/policy_engine/test_pipeline_types.py b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py new file mode 100644 index 00000000000..2e5986d3ef8 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py @@ -0,0 +1,168 @@ +""" +Tests for pipeline type definitions. +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.pipeline_types import ( + GuardrailPipeline, + PipelineExecutionResult, + PipelineStep, + PipelineStepResult, +) +from litellm.types.proxy.policy_engine.policy_types import ( + Policy, + PolicyGuardrails, +) + + +def test_pipeline_step_defaults(): + step = PipelineStep(guardrail="my-guard") + assert step.on_fail == "block" + assert step.on_pass == "allow" + assert step.on_error is None + assert step.pass_data is False + assert step.modify_response_message is None + + +def test_pipeline_step_valid_actions(): + step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") + assert step.on_fail == "next" + assert step.on_pass == "next" + + +def test_pipeline_step_all_action_types(): + for action in ("allow", "block", "next", "modify_response"): + step = PipelineStep( + guardrail="g", on_fail=action, on_pass=action, on_error=action + ) + assert step.on_fail == action + assert step.on_pass == action + assert step.on_error == action + + +def test_pipeline_step_invalid_action_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_fail="invalid_action") + + +def test_pipeline_step_invalid_on_pass_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_pass="skip") + + +def test_pipeline_step_on_error_valid(): + step = PipelineStep( + guardrail="g", on_error="next", on_fail="block", on_pass="allow" + ) + assert step.on_error == "next" + + +def test_pipeline_step_invalid_on_error_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_error="invalid") + + +def test_pipeline_requires_at_least_one_step(): + with pytest.raises(ValidationError): + GuardrailPipeline(mode="pre_call", steps=[]) + + +def test_pipeline_invalid_mode_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="during_call", + steps=[PipelineStep(guardrail="g")], + ) + + +def test_pipeline_valid_modes(): + for mode in ("pre_call", "post_call"): + pipeline = GuardrailPipeline( + mode=mode, + steps=[PipelineStep(guardrail="g")], + ) + assert pipeline.mode == mode + + +def test_pipeline_with_multiple_steps(): + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), + PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), + ], + ) + assert len(pipeline.steps) == 2 + assert pipeline.steps[0].guardrail == "g1" + assert pipeline.steps[1].guardrail == "g2" + + +def test_policy_with_pipeline_parses(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1", "g2"]), + pipeline=GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next"), + PipelineStep(guardrail="g2"), + ], + ), + ) + assert policy.pipeline is not None + assert len(policy.pipeline.steps) == 2 + + +def test_policy_without_pipeline(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1"]), + ) + assert policy.pipeline is None + + +def test_pipeline_step_result(): + result = PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + error_detail="Content policy violation", + duration_seconds=0.05, + ) + assert result.outcome == "fail" + assert result.action_taken == "next" + + +def test_pipeline_execution_result(): + result = PipelineExecutionResult( + terminal_action="block", + step_results=[ + PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + ), + PipelineStepResult( + guardrail_name="g2", + outcome="fail", + action_taken="block", + ), + ], + error_message="Content blocked", + ) + assert result.terminal_action == "block" + assert len(result.step_results) == 2 + + +def test_pipeline_step_extra_fields_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="g", unknown_field="value") + + +def test_pipeline_extra_fields_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="g")], + unknown="value", + ) diff --git a/tests/unit/types/proxy/policy_engine/test_policy_types.py b/tests/unit/types/proxy/policy_engine/test_policy_types.py new file mode 100644 index 00000000000..bcd6d39aa4d --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_policy_types.py @@ -0,0 +1,15 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_accepts_int32_priority(priority: int): + assert PolicyAttachment(policy="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachment(policy="p", priority=priority) diff --git a/tests/unit/types/proxy/policy_engine/test_resolver_types.py b/tests/unit/types/proxy/policy_engine/test_resolver_types.py new file mode 100644 index 00000000000..f31b9d7e873 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_resolver_types.py @@ -0,0 +1,115 @@ +""" +Tests for pipeline field on policy CRUD types (resolver_types.py). +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyAttachmentCreateRequest, + PolicyCreateRequest, + PolicyDBResponse, + PolicyUpdateRequest, +) + + +def test_policy_create_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert req.pipeline is not None + assert req.pipeline["mode"] == "pre_call" + assert len(req.pipeline["steps"]) == 2 + + +def test_policy_create_request_without_pipeline(): + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1"], + ) + assert req.pipeline is None + + +def test_policy_update_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyUpdateRequest(pipeline=pipeline_data) + assert req.pipeline is not None + assert req.pipeline["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert resp.pipeline is not None + assert resp.pipeline["mode"] == "pre_call" + dumped = resp.model_dump() + assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_without_pipeline(): + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + ) + assert resp.pipeline is None + dumped = resp.model_dump() + assert dumped["pipeline"] is None + + +def test_policy_create_request_roundtrip(): + pipeline_data = { + "mode": "post_call", + "steps": [ + { + "guardrail": "g1", + "on_fail": "modify_response", + "on_pass": "next", + "pass_data": True, + "modify_response_message": "custom msg", + }, + ], + } + req = PolicyCreateRequest( + policy_name="roundtrip-test", + guardrails_add=["g1"], + pipeline=pipeline_data, + ) + dumped = req.model_dump() + restored = PolicyCreateRequest(**dumped) + assert restored.pipeline == pipeline_data + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_create_request_accepts_int32_priority(priority: int): + assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/unit/videos/__init__.py b/tests/unit/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py new file mode 100644 index 00000000000..22e1e5c05eb --- /dev/null +++ b/tests/unit/videos/test_main.py @@ -0,0 +1,455 @@ +""" +Dispatch-contract tests for litellm/videos/main.py + +Each public video operation is a pair: a sync `video_*` worker (decorated with +@client) that resolves the provider, fetches the provider config, logs, and then +forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async +`avideo_*` wrapper that delegates to the sync worker in an executor. + +This file locks the contract of that layer so a regression fails loudly: + + 1. DISPATCH - the one correct handler fired and every sibling video handler + asserted NOT called. A copy-paste that calls the wrong handler + (e.g. remix -> edit) flips this. + 2. RESULT - the handler's return value is propagated by identity. + 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not + passed (status/content/remix/edit/extension), or defaults to + "openai" (list/create_character/get_character). This is the exact + surface of the historical "content defaulted to openai" bug. + 4. PAYLOAD - the provider config object and the operation's identifying args + (video_id/prompt/name/...) reach the handler; _is_async is False + on the sync path. + 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. + 6. UNSUPPORTED - a None provider config raises before any handler fires. + 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets + async_call=True, and pre-resolves the provider where it must. + +Seams mocked: the http handler (network), the provider-config registry lookup, +get_llm_provider, and the video-generation optional-param builders. The id decode +helper runs for real against genuinely-encoded ids, so the provider assertions +reflect production. +""" + +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Any, Dict +from unittest.mock import MagicMock, patch + +import pytest + + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.videos.main import CharacterObject, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider +from litellm.videos import main as videos_main + +# A real model-encoded video id: decodes (for real) to provider "azure". Used to +# prove the sync workers derive custom_llm_provider from the id, not a hardcode. +AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") + +# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly +# one fired and the other eight did not. +SYNC_HANDLERS = ( + "video_generation_handler", + "video_content_handler", + "video_remix_handler", + "video_create_character_handler", + "video_get_character_handler", + "video_edit_handler", + "video_extension_handler", + "video_list_handler", + "video_status_handler", +) + +GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} + + +@dataclass +class Seams: + handler: MagicMock + get_config: MagicMock + config: MagicMock + + def kwargs_of(self, handler_name: str) -> Dict[str, Any]: + method = getattr(self.handler, handler_name) + assert method.call_count == 1 + return dict(method.call_args.kwargs) + + def assert_only(self, handler_name: str) -> None: + for name in SYNC_HANDLERS: + method = getattr(self.handler, name) + if name == handler_name: + method.assert_called_once() + else: + method.assert_not_called() + + +@pytest.fixture +def seams(): + handler = MagicMock(spec=BaseLLMHTTPHandler) + config = MagicMock(name="provider_video_config") + get_config = MagicMock(return_value=config) + + with ExitStack() as stack: + stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) + stack.enter_context( + patch.object( + videos_main.ProviderConfigManager, + "get_provider_video_config", + get_config, + ) + ) + # video_generation resolves model+provider through get_llm_provider and + # builds optional params; mock those so the dispatch payload is deterministic. + stack.enter_context( + patch.object( + videos_main, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_requested_video_generation_optional_param", + MagicMock(return_value={"seconds": "8"}), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_optional_params_video_generation", + MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), + ) + ) + yield Seams(handler=handler, get_config=get_config, config=config) + + +# =========================================================================== # +# Dispatch contract - one rich test per sync worker. +# =========================================================================== # + + +def test_video_generation__dispatch(seams): + result = videos_main.video_generation(prompt="a sunset", model="sora-2") + + seams.assert_only("video_generation_handler") + assert result is seams.handler.video_generation_handler.return_value + kw = seams.kwargs_of("video_generation_handler") + assert kw["model"] == "sora-2" + assert kw["prompt"] == "a sunset" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_generation_provider_config"] is seams.config + assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS + assert kw["_is_async"] is False + + +def test_video_status__dispatch_and_provider_from_id(seams): + result = videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.assert_only("video_status_handler") + assert result is seams.handler.video_status_handler.return_value + kw = seams.kwargs_of("video_status_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai + assert kw["video_status_provider_config"] is seams.config + assert kw["_is_async"] is False + # provider config requested for the decoded provider, not a hardcode. + assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE + + +def test_video_content__dispatch_and_provider_from_id(seams): + result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") + + seams.assert_only("video_content_handler") + assert result is seams.handler.video_content_handler.return_value + kw = seams.kwargs_of("video_content_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" + assert kw["variant"] == "thumbnail" + assert kw["video_content_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_content__plain_id_defaults_to_openai(seams): + videos_main.video_content(video_id="video_plain") + + assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" + + +def test_video_remix__dispatch_and_provider_from_id(seams): + result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") + + seams.assert_only("video_remix_handler") + assert result is seams.handler.video_remix_handler.return_value + kw = seams.kwargs_of("video_remix_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "new colors" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_remix_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_edit__dispatch_and_provider_from_id(seams): + result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") + + seams.assert_only("video_edit_handler") + assert result is seams.handler.video_edit_handler.return_value + kw = seams.kwargs_of("video_edit_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "brighter" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_extension__dispatch_and_provider_from_id(seams): + result = videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" + ) + + seams.assert_only("video_extension_handler") + assert result is seams.handler.video_extension_handler.return_value + kw = seams.kwargs_of("video_extension_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "continue" + assert kw["seconds"] == "5" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_list__dispatch_defaults_to_openai(seams): + result = videos_main.video_list(after="cur", limit=5, order="desc") + + seams.assert_only("video_list_handler") + assert result is seams.handler.video_list_handler.return_value + kw = seams.kwargs_of("video_list_handler") + assert kw["after"] == "cur" + assert kw["limit"] == 5 + assert kw["order"] == "desc" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_list_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_create_character__dispatch_defaults_to_openai(seams): + video = MagicMock(name="video_upload") + result = videos_main.video_create_character(name="hero", video=video) + + seams.assert_only("video_create_character_handler") + assert result is seams.handler.video_create_character_handler.return_value + kw = seams.kwargs_of("video_create_character_handler") + assert kw["name"] == "hero" + assert kw["video"] is video + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_get_character__dispatch_defaults_to_openai(seams): + result = videos_main.video_get_character(character_id="char_1") + + seams.assert_only("video_get_character_handler") + assert result is seams.handler.video_get_character_handler.return_value + kw = seams.kwargs_of("video_get_character_handler") + assert kw["character_id"] == "char_1" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_explicit_provider_beats_decoded_id(seams): + """An explicit custom_llm_provider wins over the one encoded in the id.""" + videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") + + assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" + + +# =========================================================================== # +# mock_response short-circuit - returns a typed object, no handler call. +# =========================================================================== # + + +def test_generation__mock_response_short_circuits(seams): + resp = videos_main.video_generation( + prompt="x", + model="sora-2", + mock_response={"id": "v1", "object": "video", "status": "queued"}, + ) + + assert isinstance(resp, VideoObject) + assert resp.id == "v1" + seams.handler.video_generation_handler.assert_not_called() + + +def test_list__mock_response_short_circuits(seams): + resp = videos_main.video_list( + mock_response=[{"id": "v1", "object": "video", "status": "completed"}] + ) + + assert isinstance(resp, list) + assert resp[0].id == "v1" + seams.handler.video_list_handler.assert_not_called() + + +def test_get_character__mock_response_short_circuits(seams): + resp = videos_main.video_get_character( + character_id="char_1", + mock_response={ + "id": "char_1", + "object": "character", + "created_at": 1, + "name": "hero", + }, + ) + + assert isinstance(resp, CharacterObject) + assert resp.id == "char_1" + seams.handler.video_get_character_handler.assert_not_called() + + +# =========================================================================== # +# Unsupported provider - a None provider config raises before any dispatch. +# =========================================================================== # + + +def test_unsupported_provider_raises_without_dispatch(seams): + seams.get_config.return_value = None + + with pytest.raises(litellm.APIConnectionError): + videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.handler.video_status_handler.assert_not_called() + + +# =========================================================================== # +# Async-wrapper delegation - representative coverage. +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_avideo_generation__delegates_with_async_flag(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with ( + patch.object( + videos_main, "video_generation", MagicMock(return_value=sentinel) + ) as sync, + patch.object( + litellm, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ), + ): + result = await videos_main.avideo_generation(prompt="x", model="sora-2") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_avideo_status__delegates_untouched(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with patch.object( + videos_main, "video_status", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_status(video_id="video_plain") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["video_id"] == "video_plain" + + +@pytest.mark.asyncio +async def test_avideo_content__pre_decodes_provider_before_delegating(): + """avideo_content resolves the provider from the encoded id itself before + handing off, so the sync worker receives the decoded provider, not None.""" + sentinel = b"mp4-bytes" + with patch.object( + videos_main, "video_content", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "azure" + + +# =========================================================================== # +# Credential passthrough - DB/YAML model-config credentials the router injects +# via kwargs must reach the provider call for EVERY video handler, carried in +# litellm_params. Distinct per-field values catch a cross-wired field. +# =========================================================================== # + +DB_YAML_CREDS = { + "api_key": "sk-db-credential", + "api_base": "https://db-resource.test", + "api_version": "2024-12-31", + "vertex_project": "db-project-xyz", +} + +CREDENTIAL_OPERATIONS = [ + ( + "video_generation_handler", + lambda: videos_main.video_generation( + prompt="p", model="sora-2", **DB_YAML_CREDS + ), + ), + ( + "video_status_handler", + lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_content_handler", + lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_remix_handler", + lambda: videos_main.video_remix( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_edit_handler", + lambda: videos_main.video_edit( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_extension_handler", + lambda: videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS + ), + ), + ( + "video_list_handler", + lambda: videos_main.video_list(**DB_YAML_CREDS), + ), + ( + "video_create_character_handler", + lambda: videos_main.video_create_character( + name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS + ), + ), + ( + "video_get_character_handler", + lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), + ), +] + + +@pytest.mark.parametrize( + "handler_name,invoke", + CREDENTIAL_OPERATIONS, + ids=[op[0] for op in CREDENTIAL_OPERATIONS], +) +def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): + invoke() + + litellm_params = seams.kwargs_of(handler_name)["litellm_params"] + assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] + assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] + assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] + assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/unit/videos/test_utils.py b/tests/unit/videos/test_utils.py new file mode 100644 index 00000000000..728644cdda5 --- /dev/null +++ b/tests/unit/videos/test_utils.py @@ -0,0 +1,181 @@ +""" +Pure-logic contract tests for litellm/videos/main.py's request utils +(litellm/videos/utils.py: VideoGenerationRequestUtils). + +These lock the exact param-shaping behavior so a mutation that drops a filter, +flips a precedence, or stops removing a key fails loudly. The only seam is the +provider config's map_openai_params (a provider boundary); filter_out_litellm_params +runs for real, so the "litellm-internal params get stripped" assertions reflect +production. Every test asserts the exact resulting dict, never "ran without error". +""" + +from unittest.mock import MagicMock + + + +import litellm +from litellm.videos.utils import VideoGenerationRequestUtils + +get_requested = ( + VideoGenerationRequestUtils.get_requested_video_generation_optional_param +) +get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation + + +# =========================================================================== # +# get_requested_video_generation_optional_param +# +# Receives the caller's full local_vars; must return only the API-bound optional +# params. filter_out_litellm_params strips known internal keys for real; the +# values used below were chosen against the live set: seconds/size/user/foo_param/ +# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. +# =========================================================================== # + + +def test_requested__drops_none_and_excluded_keys(): + result = get_requested( + { + "seconds": "8", + "size": None, # None -> dropped + "prompt": "a sunset", # excluded + "model": "sora-2", # excluded + "user": "u1", + } + ) + assert result == {"seconds": "8", "user": "u1"} + + +def test_requested__strips_litellm_internal_params(): + result = get_requested( + { + "seconds": "8", + "api_key": "sk-secret", + "metadata": {"x": 1}, + "litellm_call_id": "id-123", + } + ) + assert result == {"seconds": "8"} + + +def test_requested__timeout_always_removed(): + # timeout is NOT a litellm-internal param, so only the explicit pop removes it. + result = get_requested({"seconds": "8", "timeout": 30}) + assert result == {"seconds": "8"} + + +def test_requested__nested_kwargs_merge_and_override_base(): + result = get_requested( + {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} + ) + # nested kwargs win over the top-level base params on collision. + assert result == {"seconds": "override", "size": "720x1280"} + + +def test_requested__non_dict_kwargs_treated_as_empty(): + result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) + assert result == {"seconds": "8"} + + +def test_requested__none_input_returns_empty(): + assert get_requested(None) == {} + + +def test_requested__top_level_extra_body_spread_and_preserved(): + result = get_requested( + {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} + ) + # extra_body keys are both spread at top level AND kept under "extra_body". + assert result == { + "seconds": "8", + "vertex_project": "proj", + "foo_param": "bar", + "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, + } + + +def test_requested__extra_body_kwargs_overrides_top_level(): + result = get_requested( + { + "extra_body": {"a": "top", "b": "top_b"}, + "kwargs": {"extra_body": {"a": "kw"}}, + } + ) + # kwargs' extra_body wins over the top-level extra_body on collision; the + # non-colliding top-level key survives. + assert result == { + "a": "kw", + "b": "top_b", + "extra_body": {"a": "kw", "b": "top_b"}, + } + + +def test_requested__extra_body_strips_litellm_internal_params(): + result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) + # api_key filtered out of extra_body; only foo_param remains (and is spread). + assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} + + +def test_requested__empty_extra_body_not_added(): + result = get_requested({"seconds": "8", "extra_body": {}}) + assert result == {"seconds": "8"} + assert "extra_body" not in result + + +# =========================================================================== # +# get_optional_params_video_generation +# +# Delegates mapping to the provider config (the seam) then folds extra_body in. +# =========================================================================== # + + +def _config(map_return): + config = MagicMock() + config.map_openai_params.return_value = map_return + return config + + +def test_optional__delegates_to_map_openai_params_with_drop_params(): + config = _config({"seconds": "8"}) + optional_params = {"seconds": "8"} + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params=optional_params, + ) + + assert result == {"seconds": "8"} + config.map_openai_params.assert_called_once_with( + video_create_optional_params=optional_params, + model="sora-2", + drop_params=litellm.drop_params, + ) + + +def test_optional__extra_body_overrides_mapped_and_is_removed(): + # mapped output carries a leftover extra_body that must be popped; the input + # extra_body overrides a colliding mapped key and is spread in. + config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={ + "extra_body": {"size": "override", "extra": "x"} + }, + ) + + assert result == {"seconds": "8", "size": "override", "extra": "x"} + assert "extra_body" not in result + + +def test_optional__non_dict_extra_body_ignored(): + config = _config({"seconds": "8"}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={"seconds": "8", "extra_body": None}, + ) + + assert result == {"seconds": "8"} From 924ad6e57118b7e42f2f483c9a028c41f90f25f3 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:22 +0000 Subject: [PATCH 067/127] test: remove phase 16 legacy test files from tests/test_litellm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../rust_bridge/ocr/test_route_host.py | 85 -- .../rust_bridge/responses/test_route_host.py | 57 - .../test_litellm/sandbox/test_e2b_sandbox.py | 318 ----- .../sandbox/test_opensandbox_sandbox.py | 647 ---------- .../sandbox/test_sandbox_tools.py | 181 --- tests/test_litellm/skills/test_skills_main.py | 57 - .../test_enforce_model_rate_limits.py | 468 -------- .../test_router/test_io_token_rate_limits.py | 1069 ----------------- .../types/llms/test_types_llms_bedrock.py | 46 - .../types/llms/test_types_llms_openai.py | 591 --------- .../policy_engine/test_pipeline_types.py | 168 --- .../proxy/policy_engine/test_policy_types.py | 15 - .../policy_engine/test_resolver_types.py | 115 -- tests/test_litellm/videos/test_main.py | 455 ------- tests/test_litellm/videos/test_utils.py | 193 --- 15 files changed, 4465 deletions(-) delete mode 100644 tests/test_litellm/rust_bridge/ocr/test_route_host.py delete mode 100644 tests/test_litellm/rust_bridge/responses/test_route_host.py delete mode 100644 tests/test_litellm/sandbox/test_e2b_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_opensandbox_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_sandbox_tools.py delete mode 100644 tests/test_litellm/skills/test_skills_main.py delete mode 100644 tests/test_litellm/test_router/test_enforce_model_rate_limits.py delete mode 100644 tests/test_litellm/test_router/test_io_token_rate_limits.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_bedrock.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_openai.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_policy_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py delete mode 100644 tests/test_litellm/videos/test_main.py delete mode 100644 tests/test_litellm/videos/test_utils.py diff --git a/tests/test_litellm/rust_bridge/ocr/test_route_host.py b/tests/test_litellm/rust_bridge/ocr/test_route_host.py deleted file mode 100644 index 699492e4424..00000000000 --- a/tests/test_litellm/rust_bridge/ocr/test_route_host.py +++ /dev/null @@ -1,85 +0,0 @@ -from typing import Final - -import pytest - -import litellm -from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure -from litellm.rust_bridge.ocr.route_host import response as build_ocr_response -from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest - -REQUEST: Final = LiteLLMOcrRequest( - model="mistral/mistral-ocr-latest", - document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, - api_key="test-key", - api_base=None, - timeout=None, - custom_llm_provider=None, - extra_headers=None, - kwargs={"req_format": "markdown"}, -) - - -class RustUpstreamError(Exception): - def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: - super().__init__(status, body) - self.headers: Final = list(headers) - - -class RustFormatError(Exception): - ocr_request_format_error: Final = True - - -def test_rust_ocr_response_retains_provider_native_response(): - provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} - response = build_ocr_response( - { - "pages": [], - "model": "prebuilt-layout", - "document_annotation": None, - "usage_info": {"pages_processed": 0}, - "object": "ocr", - "provider_native_response": provider_response, - } - ) - - assert response.get_provider_native_response() == provider_response - assert response.model_dump().get("provider_native_response") is None - - -def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: - error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.RateLimitError) - assert public_error.status_code == 429 - assert public_error.response.headers["retry-after"] == "7" - assert public_error.response.text == '{"message": "slow down"}' - assert public_error.__context__ is error - assert public_error.llm_provider == "mistral" - - -def test_map_failure_maps_upstream_401_to_authentication_error() -> None: - error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.AuthenticationError) - assert public_error.status_code == 401 - assert public_error.response.text == '{"message": "Unauthorized"}' - assert public_error.__context__ is error - - -def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: - error: Final = RuntimeError("bridge exploded") - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert not isinstance(public_error, UpstreamFailure) - assert isinstance(public_error, litellm.APIConnectionError) - assert "bridge exploded" in str(public_error) - - -def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: - with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): - raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/test_litellm/rust_bridge/responses/test_route_host.py b/tests/test_litellm/rust_bridge/responses/test_route_host.py deleted file mode 100644 index 49bf19e7d8a..00000000000 --- a/tests/test_litellm/rust_bridge/responses/test_route_host.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import MappingProxyType -from typing import Final - -import pytest -from pydantic import ValidationError - -from litellm.rust_bridge.responses.route_host import arguments, response -from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.types.llms.openai import ResponsesAPIResponse - - -def test_response_validates_into_the_public_responses_model() -> None: - built: Final = response( - MappingProxyType( - { - "id": "resp_native", - "object": "response", - "created_at": 1, - "model": "gpt-4o", - "status": "completed", - "output": [ - { - "type": "message", - "id": "msg_native", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "native", "annotations": []}], - } - ], - } - ) - ) - - assert isinstance(built, ResponsesAPIResponse) - assert built.id == "resp_native" - assert built.output[0].content[0].text == "native" - - -def test_response_rejects_a_payload_missing_required_fields() -> None: - with pytest.raises(ValidationError): - response(MappingProxyType({"object": "response"})) - - -def test_arguments_are_the_public_kwargs_view() -> None: - kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) - request: Final = LiteLLMResponsesRequest( - model="gpt-4o", - input="hi", - stream=None, - api_key=None, - api_base=None, - custom_llm_provider="openai", - extra_headers=None, - kwargs=kwargs, - ) - - assert arguments(request) is kwargs diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py deleted file mode 100644 index e01b9120416..00000000000 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -Tests for the e2b code execution sandbox primitive. - -Unit tests inject a fake async HTTP client (dependency injection, no -monkeypatching) and assert request shapes and result mapping. Real-network -integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. -""" - -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.e2b.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - E2BSandboxConfig, -) - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - """Records outbound requests and returns canned responses keyed by URL.""" - - def __init__( - self, - *, - create_json=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "sandboxID": "sbx_123", - "domain": "e2b.app", - "envdAccessToken": "tok_abc", - } - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json)) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/execute"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None)) - if not (200 <= self.delete_status < 300): - raise httpx.HTTPStatusError( - f"status {self.delete_status}", - request=httpx.Request("DELETE", url), - response=httpx.Response(self.delete_status), - ) - return FakeResponse(status_code=self.delete_status) - - -# ---------- pure parser ---------- - - -def test_parse_lines_stdout_and_count(): - lines = [ - json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), - json.dumps({"type": "number_of_executions", "execution_count": 1}), - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "6\n" - assert result.execution_count == 1 - assert result.error is None - - -def test_parse_lines_error_surfaces_name_and_traceback(): - lines = [ - json.dumps( - { - "type": "error", - "name": "ZeroDivisionError", - "value": "division by zero", - "traceback": "Traceback (most recent call last): ...", - } - ) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.error["name"] == "ZeroDivisionError" - assert "Traceback" in result.error["traceback"] - - -def test_parse_lines_result_carries_png(): - lines = [ - json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.results and result.results[0]["png"] == "BASE64DATA" - assert "type" not in result.results[0] - - -# ---------- request shapes ---------- - - -@pytest.mark.asyncio -async def test_template_flows_into_create_request_as_templateID(): - client = FakeHTTPClient() - cfg = E2BSandboxConfig() - handle = await cfg.acreate_sandbox( - template="my-custom-template", api_key="e2b_key", client=client - ) - - method, url, headers, body = client.calls[0] - assert method == "POST" - assert url.endswith("/sandboxes") - assert body["templateID"] == "my-custom-template" # not "template" - assert body["secure"] is True - assert headers["X-API-Key"] == "e2b_key" - assert handle.id == "sbx_123" - assert handle._hidden_params["envd_access_token"] == "tok_abc" - - -@pytest.mark.asyncio -async def test_create_defaults_template_when_omitted(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) - _, _, _, body = client.calls[0] - assert body["templateID"] == "code-interpreter-v1" - - -@pytest.mark.asyncio -async def test_run_code_targets_jupyter_host_with_access_token(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] - ) - handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok_run"} - - result = await E2BSandboxConfig().arun_code( - container=handle, code="print(6*7)", client=client - ) - - method, url, headers, body = client.calls[0] - assert url == "https://49999-sbx_xyz.e2b.app/execute" - assert headers["X-Access-Token"] == "tok_run" - assert body["code"] == "print(6*7)" - assert result.stdout.strip() == "42" - - -@pytest.mark.asyncio -async def test_delete_issues_delete_to_sandbox_id(): - client = FakeHTTPClient(delete_status=204) - handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - method, url, headers, _ = client.calls[0] - assert method == "DELETE" - assert url.endswith("/sandboxes/sbx_del") - assert ok is True - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - assert ok is False - - -# ---------- ephemeral teardown ---------- - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="e2b", code="1/0", api_key="e2b_key", client=client - ) - - methods = [c[0] for c in client.calls] - urls = [c[1] for c in client.calls] - assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete - assert urls[0].endswith("/sandboxes") - assert urls[1].endswith("/execute") - assert urls[2].endswith("/sandboxes/sbx_123") - - -# ---------- correctness guards ---------- - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - with pytest.raises(httpx.HTTPStatusError): - await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - -@pytest.mark.asyncio -async def test_create_preserves_explicit_zero_timeout(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - timeout=0, api_key="e2b_key", client=client - ) - _, _, _, body = client.calls[0] - assert body["timeout"] == 0 - - -@pytest.mark.asyncio -async def test_run_code_rejects_bare_id_without_access_token(): - client = FakeHTTPClient() - with pytest.raises(ValueError, match="access token"): - await E2BSandboxConfig().arun_code( - container="sbx_no_token", code="print(1)", client=client - ) - assert client.calls == [] # never reached the network - - -def test_parse_lines_skips_non_json_lines(): - lines = [ - "not-json-heartbeat", - json.dumps({"type": "stdout", "text": "ok\n"}), - "", - "{partial", - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "ok\n" - assert result.error is None - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - big_line = "x" * (MAX_OUTPUT_BYTES + 1) - client = FakeHTTPClient(execute_lines=[big_line]) - handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok"} - with pytest.raises(ValueError, match="exceeded"): - await E2BSandboxConfig().arun_code( - container=handle, code="print('x'*999)", client=client - ) - - -# ---------- public entrypoints ---------- - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] - ) - container = await litellm.acreate_sandbox( - provider="e2b", api_key="e2b_key", client=client - ) - assert container.id == "sbx_123" - - result = await litellm.arun_code( - provider="e2b", - container=container, - api_key="e2b_key", - code="print(6*7)", - client=client, - ) - assert result.stdout.strip() == "42" - - assert ( - await litellm.adelete_sandbox( - provider="e2b", container=container, api_key="e2b_key", client=client - ) - is True - ) - - -@pytest.mark.asyncio -async def test_unsupported_provider_raises(): - with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): - await litellm.acreate_sandbox(provider="not-a-provider") - - -# ---------- api_base override ---------- - - -@pytest.mark.asyncio -async def test_create_uses_api_base_override(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - api_base="http://my-sandbox:8080", api_key="k", client=client - ) - _, url, _, _ = client.calls[0] - assert url == "http://my-sandbox:8080/sandboxes" - - -@pytest.mark.asyncio -async def test_create_defaults_to_e2b_api_base(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) - _, url, _, _ = client.calls[0] - assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py deleted file mode 100644 index 2928dea100e..00000000000 --- a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py +++ /dev/null @@ -1,647 +0,0 @@ -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.opensandbox.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - OPEN_SANDBOX_DEFAULT_TEMPLATE, - OpenSandboxSandboxConfig, -) -from litellm.utils import ProviderConfigManager - -TEST_API_BASE = "https://sandbox.test/v1" - - -def http_status_error(status_code, url="http://test"): - return httpx.HTTPStatusError( - f"status {status_code}", - request=httpx.Request("GET", url), - response=httpx.Response(status_code), - ) - - -def sse(data): - return f"data: {json.dumps(data)}" - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - def raise_for_status(self): - if self.status_code >= 400: - raise http_status_error(self.status_code) - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - def __init__( - self, - *, - create_json=None, - sandbox_states=None, - endpoint_json=None, - endpoint_responses=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - self.sandbox_states = list( - sandbox_states - or [ - { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - ] - ) - self.endpoint_json = endpoint_json or { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - } - self.endpoint_responses = ( - list(endpoint_responses) if endpoint_responses is not None else None - ) - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json, {"stream": stream})) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/code"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def get(self, url, headers=None, params=None, **kwargs): - self.calls.append(("GET", url, headers, None, params)) - if "/endpoints/44772" in url: - if self.endpoint_responses is not None and self.endpoint_responses: - response = self.endpoint_responses.pop(0) - if isinstance(response, Exception): - raise response - if isinstance(response, FakeResponse): - return response - return FakeResponse(json_data=response) - return FakeResponse(json_data=self.endpoint_json) - if "/sandboxes/" in url: - state = self.sandbox_states.pop(0) - return FakeResponse(json_data=state) - raise AssertionError(f"unexpected GET {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None, None)) - if not (200 <= self.delete_status < 300): - raise http_status_error(self.delete_status, url) - return FakeResponse(status_code=self.delete_status) - - -def test_parse_sse_lines_maps_output_result_count_and_error(): - lines = [ - sse({"type": "stdout", "text": "hello\n"}), - sse({"type": "stderr", "text": "warn\n"}), - sse({"type": "result", "results": {"text/plain": "4"}}), - sse({"type": "execution_count", "execution_count": 7}), - sse( - { - "type": "error", - "error": { - "ename": "ValueError", - "evalue": "bad", - "traceback": ["Traceback"], - }, - } - ), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "hello\n" - assert result.stderr == "warn\n" - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 7 - assert result.error == { - "name": "ValueError", - "value": "bad", - "traceback": ["Traceback"], - } - - -def test_parse_sse_lines_skips_non_json_and_control_lines(): - lines = [ - "event: message", - "not-json", - "", - sse({"type": "stdout", "text": "ok\n"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "ok\n" - assert result.error is None - - -def test_parse_sse_lines_maps_fallback_shapes(): - lines = [ - "data:", - sse(["not-a-dict"]), - sse({"code": "BadRequest", "message": "nope"}), - sse({"type": "result", "text/plain": "4"}), - sse({"type": "error", "name": "RuntimeError", "text": "boom"}), - sse({"type": "execution_count", "execution_count": "8"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 8 - assert result.error == { - "name": "BadRequest", - "value": "nope", - "traceback": [], - } - fallback_error = OpenSandboxSandboxConfig._parse_lines( - [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] - ) - assert fallback_error.error == { - "name": "RuntimeError", - "value": "boom", - "traceback": [], - } - empty_string_error = OpenSandboxSandboxConfig._parse_lines( - [ - sse( - { - "type": "error", - "error": { - "ename": "", - "name": "FallbackName", - "evalue": "", - "value": "fallback value", - "traceback": [], - }, - } - ) - ] - ) - assert empty_string_error.error == { - "name": "", - "value": "", - "traceback": [], - } - - -def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): - def fake_secret(key): - if key == "OPEN_SANDBOX_API_KEY": - return "env-key" - if key == "OPEN_SANDBOX_API_BASE": - return TEST_API_BASE - return None - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - fake_secret, - ) - config = OpenSandboxSandboxConfig() - handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") - - assert config.validate_environment() == "env-key" - assert config.validate_environment(api_key="") == "" - assert config._api_key(api_key=None, handle=handle) == "env-key" - - handle._hidden_params = {"api_key": "stored-key"} - assert config._api_key(api_key=None, handle=handle) == "stored-key" - assert config._http(None) is not None - - body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy={"egress": [{"domain": "example.com"}]}, - secure_access=True, - ) - assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} - assert body["secureAccess"] is True - - other_body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy=None, - secure_access=False, - ) - assert body["resourceLimits"] is not other_body["resourceLimits"] - - assert config._sandbox_state(None) is None - assert config._sandbox_state({"status": "Running"}) is None - assert config._as_str_dict(None) == {} - assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( - "http://execd.local" - ) - assert config._api_base(None) == TEST_API_BASE - assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" - assert config._as_int("9") == 9 - assert config._as_int("nope") is None - assert config._as_int(None) is None - assert isinstance( - ProviderConfigManager.get_provider_sandbox_config("opensandbox"), - OpenSandboxSandboxConfig, - ) - - -def test_api_base_requires_kwarg_or_env(monkeypatch): - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - lambda key: None, - ) - - with pytest.raises(ValueError, match="api_base is required"): - OpenSandboxSandboxConfig._api_base(None) - - -@pytest.mark.asyncio -async def test_create_posts_default_body_and_omits_empty_api_key(): - client = FakeHTTPClient() - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - method, url, headers, body, _ = client.calls[0] - assert method == "POST" - assert url == f"{TEST_API_BASE}/sandboxes" - assert "OPEN-SANDBOX-API-KEY" not in headers - assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} - assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] - assert body["timeout"] == 300 - assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert handle.id == "osb_123" - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - - -@pytest.mark.asyncio -async def test_create_can_opt_into_internet_access(): - client = FakeHTTPClient() - - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - allow_internet_access=True, - client=client, - ) - - _, _, _, body, _ = client.calls[0] - assert "networkPolicy" not in body - - -@pytest.mark.asyncio -async def test_create_custom_options_poll_and_endpoint_resolution(): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - }, - sandbox_states=[ - { - "id": "osb_pending", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - } - ], - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - template="custom/image:latest", - timeout=600, - allow_internet_access=False, - api_key="osb-key", - api_base="https://sandbox.example/v1", - metadata={"suite": "unit"}, - env_vars={"PYTHONUNBUFFERED": "1"}, - resource_limits={"cpu": "500m", "memory": "512Mi"}, - resource_requests={"cpu": "250m", "memory": "256Mi"}, - entrypoint=["/bin/sh", "-lc", "sleep 3600"], - use_server_proxy=True, - client=client, - ) - - _, create_url, create_headers, body, _ = client.calls[0] - _, poll_url, poll_headers, _, _ = client.calls[1] - _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] - - assert create_url == "https://sandbox.example/v1/sandboxes" - assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert body["image"] == {"uri": "custom/image:latest"} - assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] - assert body["metadata"] == {"suite": "unit"} - assert body["env"] == {"PYTHONUNBUFFERED": "1"} - assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} - assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" - assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") - assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_params == {"use_server_proxy": True} - assert handle.id == "osb_pending" - - -@pytest.mark.asyncio -async def test_create_waits_across_pending_state(monkeypatch): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - }, - sandbox_states=[ - {"id": "osb_pending", "status": {"state": "Pending"}}, - {"id": "osb_pending", "status": {"state": "Running"}}, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - assert handle.id == "osb_pending" - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_for_terminal_state(): - client = FakeHTTPClient( - create_json={"id": "osb_failed", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_failed", "status": {"state": "Failed"}}, - ], - ) - - with pytest.raises(ValueError, match="entered Failed"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_create_times_out_waiting_for_running(): - client = FakeHTTPClient( - create_json={"id": "osb_slow", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_slow", "status": {"state": "Pending"}}, - ], - ) - - with pytest.raises(TimeoutError, match="was not Running"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=0, - poll_interval=0, - client=client, - ) - - -@pytest.mark.asyncio -async def test_create_waits_for_endpoint_resolution(monkeypatch): - client = FakeHTTPClient( - endpoint_responses=[ - http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), - { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - }, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - assert len(endpoint_calls) == 2 - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_when_endpoint_is_missing(): - client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) - - with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client - ) - - -@pytest.mark.asyncio -async def test_create_reraises_non_404_endpoint_error(): - client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_run_code_resolves_bare_id_and_posts_sse_request(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - result = await OpenSandboxSandboxConfig().arun_code( - container="osb_bare", - code="print(6*7)", - language="python", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - endpoint_call = client.calls[0] - run_call = client.calls[1] - assert endpoint_call[0] == "GET" - assert ( - endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" - ) - assert run_call[0] == "POST" - assert run_call[1] == "http://execd.local:44772/code" - assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" - assert run_call[3] == { - "code": "print(6*7)", - "context": {"language": "python"}, - } - assert run_call[4] == {"stream": True} - assert result.stdout == "42\n" - - -@pytest.mark.asyncio -async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): - client = FakeHTTPClient() - handle = ContainerHandle( - id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" - ) - handle._hidden_params = { - "execd_endpoint": "execd.example/route/44772", - "execd_headers": {}, - } - - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print(1)", client=client - ) - - assert client.calls[0][1] == "https://execd.example/route/44772/code" - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) - handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") - handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} - - with pytest.raises(ValueError, match="exceeded"): - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print('x')", client=client - ) - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - - ok = await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_gone", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - assert ok is False - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_err", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - container = await litellm.acreate_sandbox( - provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client - ) - result = await litellm.arun_code( - provider="opensandbox", - container=container, - code="print(6*7)", - api_key="", - client=client, - ) - ok = await litellm.adelete_sandbox( - provider="opensandbox", - container=container, - api_key="", - client=client, - ) - - assert container.id == "osb_123" - assert result.stdout == "42\n" - assert ok is True - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="opensandbox", - code="1/0", - api_key="", - api_base=TEST_API_BASE, - client=client, - ) - - assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] - assert client.calls[0][1].endswith("/sandboxes") - assert client.calls[1][1].endswith("/endpoints/44772") - assert client.calls[2][1].endswith("/code") - assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/test_litellm/sandbox/test_sandbox_tools.py b/tests/test_litellm/sandbox/test_sandbox_tools.py deleted file mode 100644 index 06136534b13..00000000000 --- a/tests/test_litellm/sandbox/test_sandbox_tools.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Unit tests for the sandbox-tool registry.""" - -from litellm.sandbox import sandbox_tools - - -def _reset(): - sandbox_tools.clear_sandbox_tools() - - -def test_register_resolves_provider_key_and_base(): - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved == { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - } - _reset() - - -def test_register_clears_stale_entries_on_reload(): - """A tool removed from the config must not survive a re-registration.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "old", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("old") is not None - - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "new", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("new") is not None - assert ( - sandbox_tools.resolve_sandbox_tool("old") is None - ), "stale tool must be gone after the config is reloaded" - _reset() - - -def test_register_empty_list_clears_removed_tools(): - """Reloading a config with sandbox_tools removed (the proxy passes an empty - list) must drop previously registered credentials from the process.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None - - sandbox_tools.register_sandbox_tools([]) - - assert ( - sandbox_tools.resolve_sandbox_tool("e2b_default") is None - ), "removing sandbox_tools from config must clear stale credentials" - _reset() - - -def test_register_resolves_secret_from_env(monkeypatch): - _reset() - monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "os.environ/MY_SANDBOX_KEY", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved is not None - assert resolved["api_key"] == "sk-from-env" - assert resolved["api_base"] is None - _reset() - - -def test_resolve_unknown_returns_none(): - _reset() - assert sandbox_tools.resolve_sandbox_tool("nope") is None - - -def test_register_skips_malformed_entries_without_crashing(): - """A single malformed entry (missing sandbox_tool_name, or not a dict) must - not crash registration during proxy startup/hot-reload; valid entries in the - same list must still register.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name - "not-a-dict", # wrong type - {"sandbox_tool_name": "", "litellm_params": {}}, # empty name - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("good") is not None - assert sandbox_tools.resolve_sandbox_tool("") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_skips_entry_missing_sandbox_provider(): - """An entry with a name but no sandbox_provider must be skipped at - registration so it cannot later resolve and call acreate_sandbox(provider=None), - which fails with a cryptic runtime error instead of a clear startup warning.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, - { - "sandbox_tool_name": "null_provider", - "litellm_params": {"sandbox_provider": None}, - }, - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("no_provider") is None - assert sandbox_tools.resolve_sandbox_tool("null_provider") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_swaps_registry_atomically(): - """register_sandbox_tools must replace the registry in one rebind so a - concurrent resolve never observes a half-populated or transiently empty - registry between clearing and repopulating.""" - _reset() - sandbox_tools.register_sandbox_tools( - [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] - ) - before = sandbox_tools._SANDBOX_TOOL_REGISTRY - - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, - {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, - ] - ) - after = sandbox_tools._SANDBOX_TOOL_REGISTRY - - assert after is not before, "the registry must be replaced, not mutated in place" - assert set(after) == {"b", "c"} - assert "a" not in after - _reset() diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py deleted file mode 100644 index e1c66c8d9ea..00000000000 --- a/tests/test_litellm/skills/test_skills_main.py +++ /dev/null @@ -1,57 +0,0 @@ -from unittest.mock import MagicMock - -import litellm.skills.main as skills_main -from litellm.types.utils import LlmProviders - - -def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( - monkeypatch, -) -> None: - """The REST /v1/skills form endpoint passes description/instructions as top-level - kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch - branch of create_skill() dropped both, so every LiteLLM-hosted skill was created - with description=None and instructions=None regardless of what the caller sent.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Document Translator", - description="Converts files from one language into another", - instructions="Take an uploaded document and produce it in the target language", - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Converts files from one language into another" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( - "Take an uploaded document and produce it in the target language" - ) - - -def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: - """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under - extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Warehouse SQL Analyst", - extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Runs SQL against the inventory database" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" - - -def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) - - assert handler.create_skill_handler.call_args.kwargs["description"] is None - assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py deleted file mode 100644 index 7577064b7f9..00000000000 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ /dev/null @@ -1,468 +0,0 @@ -""" -Tests for enforce_model_rate_limits feature. - -This feature allows users to enforce TPM/RPM limits set on model deployments -regardless of the routing strategy being used. -""" - -import asyncio -from datetime import timedelta -from unittest.mock import AsyncMock, MagicMock - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) - -TPM_DEPLOYMENT = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "replica-test-id"}, - "model_name": "test-model", -} - - -def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: - dual_cache = DualCache(redis_cache=redis_cache) - check = ModelRateLimitingCheck(dual_cache=dual_cache) - now = litellm.utils.get_utc_datetime() - for minute in (now, now + timedelta(minutes=1)): - tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) - dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) - return dual_cache - - -class TestModelRateLimitingCheck: - """Test the ModelRateLimitingCheck class directly.""" - - def test_get_deployment_limits_from_top_level(self): - """Test extracting limits from top-level deployment config.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "tpm": 1000, - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 1000 - assert rpm == 10 - - def test_get_deployment_limits_from_litellm_params(self): - """Test extracting limits from litellm_params.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 2000 - assert rpm == 20 - - def test_get_deployment_limits_from_model_info(self): - """Test extracting limits from model_info.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 3000 - assert rpm == 30 - - def test_get_deployment_limits_none_when_not_set(self): - """Test that None is returned when limits are not set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm is None - assert rpm is None - - def test_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - assert "current usage=11" in str(exc_info.value) - - def test_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 6 - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = 1000 # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 1000 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.parametrize( - "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] - ) - def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): - redis_cache = MagicMock() - redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() - redis_cache.increment_cache.return_value = 2 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert check.pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - def test_log_success_event_increments_cache(self): - """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - check.log_success_event(kwargs, None, None, None) - - # Verify increment_cache was called - mock_cache.increment_cache.assert_called_once() - _, kwarg_params = mock_cache.increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestModelRateLimitingCheckAsync: - """Test async methods of ModelRateLimitingCheck.""" - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=6) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=1000) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] - ) - async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.async_get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( - self, - ): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment = AsyncMock(return_value=2) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert await check.async_pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_log_success_event_increments_cache(self): - """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - await check.async_log_success_event(kwargs, None, None, None) - - # Verify async_increment_cache was called - mock_cache.async_increment_cache.assert_called_once() - _, kwarg_params = mock_cache.async_increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestRouterWithEnforceModelRateLimits: - """Test Router integration with enforce_model_rate_limits.""" - - def test_router_initializes_with_enforce_model_rate_limits(self): - """Test that Router properly initializes the ModelRateLimitingCheck.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - router = Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Check that the callback was added - assert router.optional_callbacks is not None - assert len(router.optional_callbacks) == 1 - assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) - - def test_router_optional_callbacks_contains_model_rate_limiting(self): - """Test that ModelRateLimitingCheck is in the callbacks list.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Find the ModelRateLimitingCheck in litellm.callbacks - found = False - for callback in litellm.callbacks: - if isinstance(callback, ModelRateLimitingCheck): - found = True - break - - assert found, "ModelRateLimitingCheck should be in litellm.callbacks" - - -class TestModelRateLimitConcurrency: - """Test that RPM rate limiting is atomic under concurrent requests.""" - - @pytest.mark.asyncio - async def test_concurrent_requests_respect_rpm_limit(self): - """ - Fire 4 concurrent async requests with RPM limit of 2. - Exactly 2 should succeed and 2 should raise RateLimitError. - - This test validates the atomic increment-first pattern: - the old check-then-increment pattern would let 3+ through - due to a race condition on the local cache read. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - - deployment = { - "rpm": 2, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "concurrent-test-id"}, - "model_name": "test-model", - } - - async def attempt_request(): - return await check.async_pre_call_check(deployment) - - results = await asyncio.gather( - *[attempt_request() for _ in range(4)], - return_exceptions=True, - ) - - successes = [r for r in results if not isinstance(r, Exception)] - failures = [r for r in results if isinstance(r, litellm.RateLimitError)] - - assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" - assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py deleted file mode 100644 index 3cef1c7bb63..00000000000 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ /dev/null @@ -1,1069 +0,0 @@ -""" -Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). -""" - -import asyncio - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, - ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, - OTPM_RESERVED_KEY, - _reservation_value, - _resolve_max_tokens, - async_io_token_pre_call_check, - async_io_token_reconcile_success, - build_io_token_rate_limit_headers, - deployment_has_io_token_limits, - get_io_token_rate_limit_request_kwargs, - io_token_reconcile_success, - io_token_refund_failure, - refund_stale_reservation_before_retry, - set_io_token_rate_limit_request_kwargs, -) -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) -from litellm.types.utils import ModelResponse, Usage - - -class TestIOTokenRateLimitHelpers: - def test_deployment_has_io_token_limits(self): - assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) - assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) - - def test_reservation_value_minimal_when_estimate_fails(self): - # A failed/empty estimate (0) must reserve a minimal slot, not the - # entire limit - otherwise one request whose estimate failed fills - # the whole bucket and blocks every concurrent request until it - # completes and reconciles. - assert _reservation_value(0, 100) == 1 - assert _reservation_value(0, 1) == 1 - # A real non-zero estimate is reserved as-is. - assert _reservation_value(42, 100) == 42 - - def test_resolve_max_tokens_respects_explicit_zero(self): - deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} - # An explicit max_tokens=0 is honored, not replaced by the model default. - assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 - # max_completion_tokens is the fallback only when max_tokens is absent. - assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 - assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 - - def test_build_io_token_rate_limit_headers(self): - headers = build_io_token_rate_limit_headers( - itpm_limit=200, - otpm_limit=40, - current_itpm=15, - current_otpm=4, - ) - assert headers["x-ratelimit-limit-input-tokens"] == 200 - assert headers["x-ratelimit-remaining-input-tokens"] == 185 - assert headers["x-ratelimit-limit-output-tokens"] == 40 - assert headers["x-ratelimit-remaining-output-tokens"] == 36 - - -class TestModelRateLimitingCheckIOTokens: - @pytest.mark.asyncio - async def test_itpm_reservation_and_reconcile(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-test-id"}, - "model_name": "opus", - } - - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - - kwargs = { - "standard_logging_object": { - "model_id": "io-test-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[ - { - "message": {"role": "assistant", "content": "hi"}, - "index": 0, - "finish_reason": "stop", - } - ], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - await check.async_log_success_event(kwargs, response, None, None) - - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - # ITPM tracks input tokens only (billable prompt tokens), not output. - assert current_itpm == 5 - assert current_otpm == 3 - - @pytest.mark.asyncio - async def test_itpm_limit_raises_429(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 5, - }, - "model_info": {"id": "io-limit-id"}, - "model_name": "opus", - } - - # ITPM enforces input tokens only; the prompt alone must exceed the limit, - # a large max_tokens must not contribute to the ITPM reservation. - set_io_token_rate_limit_request_kwargs( - { - "messages": [ - { - "role": "user", - "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", - } - ], - "max_tokens": 10, - "metadata": {}, - } - ) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "ITPM limit=5" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - otpm_limit = 10 - max_tokens = 4 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": otpm_limit, - }, - "model_info": {"id": "io-otpm-race-id"}, - "model_name": "opus", - } - - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": max_tokens, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - - # Atomic reservation must never let concurrent requests overshoot the limit. - assert current_otpm is not None - assert current_otpm <= otpm_limit - assert successes == otpm_limit // max_tokens - assert current_otpm == successes * max_tokens - - @pytest.mark.asyncio - async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): - """ - When input-token estimation yields 0 (no messages/prompt/input field, - unsupported model, tokenizer error), the reservation must be a - minimal 1 token, not the entire itpm limit. Otherwise the first - request whose estimate fails fills the whole bucket and every - concurrent request is rejected until it completes - effectively - serializing traffic to the deployment. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_limit = 5 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": itpm_limit, - }, - "model_info": {"id": "io-itpm-estimate-fail-id"}, - "model_name": "opus", - } - - # No messages/prompt/input field -> _estimate_input_tokens returns 0. - set_io_token_rate_limit_request_kwargs( - { - "max_tokens": 5, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - - # A minimal 1-token reservation per request lets itpm_limit concurrent - # requests through, instead of a single request starving the rest. - assert current_itpm is not None - assert current_itpm <= itpm_limit - assert successes == itpm_limit - - @pytest.mark.asyncio - async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - # Production kwargs commonly carry litellm_params.metadata; the stashed - # reservation lives in the top-level metadata and must still be found. - kwargs = { - "standard_logging_object": { - "model_id": "io-lp-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, - "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - @pytest.mark.asyncio - async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - }, - "model_info": {"id": "io-zero-est-id"}, - "model_name": "opus", - } - - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A failed/zero estimate reserves a minimal 1 token, not the full - # itpm limit, so it doesn't starve concurrent requests. - assert await dual_cache.async_get_cache(key=itpm_key) == 1 - - kwargs = { - "standard_logging_object": { - "model_id": "io-zero-est-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 7 - - @pytest.mark.asyncio - async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): - """ - A zero/failed estimate reserves a minimal 1 token rather than the - full itpm limit, so up to itpm_limit such calls are allowed - concurrently instead of the first one claiming the entire bucket. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 2, - }, - "model_info": {"id": "io-zero-cap-id"}, - "model_name": "opus", - } - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - await check.async_pre_call_check(deployment) - - # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - @pytest.mark.asyncio - async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": 5, - }, - "model_info": {"id": "io-zero-output-id"}, - "model_name": "opus", - } - zero_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 0, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) - await check.async_pre_call_check(deployment) - - zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 - assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 - - normal_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) - await check.async_pre_call_check(deployment) - - normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 - - def test_sync_io_pre_call_reserves_and_reconciles(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-sync-id"}, - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - check.pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - kwargs = { - "standard_logging_object": { - "model_id": "io-sync-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 5 - assert dual_cache.get_cache(key=otpm_key) == 3 - - @pytest.mark.asyncio - async def test_reconcile_runs_via_success_event_without_model_id(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - # standard_logging_object has no model_id (only the TPM path needs it); - # IO reconciliation must still run off the stashed cache key. - kwargs = { - "standard_logging_object": { - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - - @pytest.mark.asyncio - async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - - # Shared request metadata carrying the first (IO) deployment's reservation. - metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} - fail_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "io-first", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - } - await check.async_log_failure_event(fail_kwargs, None, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - # Retry succeeds on a non-IO fallback deployment reusing the same metadata. - retry_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "non-io-second", - "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, - "metadata": {}, - "total_tokens": 12, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), - ) - await check.async_log_success_event(retry_kwargs, response, None, None) - - # The first deployment's ITPM counter is not driven negative... - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - # ...and the non-IO deployment's TPM usage is tracked normally. - tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 - - @pytest.mark.asyncio - async def test_stale_reservation_refunded_before_retry_overwrites_it(self): - """ - A retry reuses the same mutable kwargs dict for the next deployment. - If deployment A's failure event hasn't run yet (e.g. it was scheduled - as a background task) when the retry calls - set_io_token_rate_limit_request_kwargs for deployment B, the router - must first synchronously refund + clear A's reservation via - refund_stale_reservation_before_retry - otherwise A's counter stays - elevated by the reservation until its TTL expires, and the - now-orphaned sentinels must not leak into B's accounting either. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - minute = get_utc_datetime().strftime("%H-%M") - itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) - - # Deployment A's still-unreconciled reservation, stashed on the shared - # kwargs dict the retry loop reuses. - shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} - - # Router calls this before overwriting kwargs for deployment B's attempt - - # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. - refund_stale_reservation_before_retry(dual_cache, shared_kwargs) - - # A's reservation is refunded immediately, not left stranded for a - # background failure task that may run arbitrarily later (or never, - # if the sentinels get cleared out from under it first). - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] - assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] - - # A's own (now-late) failure event finds nothing left to refund and - # is a safe no-op, since the sentinels were already cleared above. - io_token_refund_failure(dual_cache, shared_kwargs) - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - - # The retry proceeds to stash deployment B's own reservation on the - # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) - itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" - shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 - shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b - await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) - assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 - - @pytest.mark.asyncio - async def test_client_supplied_reservation_keys_are_stripped(self): - # metadata is caller-controlled; the server-only reservation sentinels - # must be removed before the router captures the request kwargs. - forged = { - "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, - "litellm_metadata": {OTPM_RESERVED_KEY: 7}, - "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, - } - set_io_token_rate_limit_request_kwargs(forged) - stored = get_io_token_rate_limit_request_kwargs() - - assert ITPM_RESERVED_KEY not in stored["metadata"] - assert ITPM_CACHE_KEY not in stored["metadata"] - assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] - assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] - - @pytest.mark.asyncio - async def test_forged_reservation_cannot_decrement_counter(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - victim_key = "global_router:victim:model:itpm:00-00" - await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) - - # A caller forges a reservation pointing at another deployment's counter. - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, - "standard_logging_object": { - "model_id": "m", - "hidden_params": {"litellm_model_name": "model"}, - "metadata": {}, - "total_tokens": 2, - }, - } - # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), - ) - await check.async_log_success_event(kwargs, response, None, None) - - # The forged reservation was stripped, so the victim counter is untouched. - assert await dual_cache.async_get_cache(key=victim_key) == 100 - - @pytest.mark.asyncio - async def test_otpm_reservation_error_rolls_back_itpm(self): - from litellm.utils import get_utc_datetime - - class _OtpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":otpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _OtpmFailCache() - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 1000, - "otpm": 1000, - }, - "model_info": {"id": "io-rollback-id"}, - "model_name": "opus", - } - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - ) - - with pytest.raises(RuntimeError): - await async_io_token_pre_call_check(dual_cache, deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A transient OTPM error must release the ITPM reservation, not leak it. - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - - @pytest.mark.asyncio - async def test_reconcile_clears_stash_even_when_increment_errors(self): - class _ItpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":itpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _ItpmFailCache() - metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} - kwargs = {"metadata": metadata} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - with pytest.raises(RuntimeError): - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - # The stash is cleared even though reconciliation raised, so a duplicate - # success event can't re-process it. - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - @pytest.mark.asyncio - async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): - import logging - - check = ModelRateLimitingCheck(dual_cache=DualCache()) - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, - "model_info": {}, - } - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - - warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] - # id-less deployments are not collapsed onto a single dedup key. - assert len(warnings) == 2 - - @pytest.mark.asyncio - async def test_missing_deployment_id_skips_io_reservation(self): - dual_cache = DualCache() - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, - "model_info": {}, # no id -> cannot build a per-deployment cache key - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - result = await async_io_token_pre_call_check(dual_cache, deployment) - - assert result is deployment - # No reservation is stashed, so nothing lands in a shared None:None bucket. - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_uses_reservation_minute_key(self): - dual_cache = DualCache() - # Reservation was made on a fixed minute key; a call that finishes in a - # later minute must reconcile against that same key, never a key built - # from the response-time minute. - itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), - ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - @pytest.mark.asyncio - async def test_reconcile_missing_usage_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - ) - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_total_tokens_only_keeps_reservation(self): - """ - A response usage object with only total_tokens (no prompt/completion - breakdown) can't be split into input/output, so it must be treated the - same as missing usage: keep the reservation instead of resolving to - (0, 0) and refunding it in full. - """ - dual_cache = DualCache() - itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = {"type": "message", "usage": {"total_tokens": 13}} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - - def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": {"total_tokens": 4}, - } - response = {"type": "message", "role": "assistant", "content": []} - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 10 - - @pytest.mark.asyncio - async def test_reconcile_falls_back_to_standard_logging_object(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "prompt_tokens": 4, - "completion_tokens": 0, - "total_tokens": 4, - }, - } - response = {"type": "message", "role": "assistant", "content": []} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - def test_sync_reconcile_anthropic_dict_usage(self): - dual_cache = DualCache() - itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=6, ttl=60) - dual_cache.set_cache(key=otpm_key, value=4, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 6, - OTPM_RESERVED_KEY: 4, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = { - "type": "message", - "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, - } - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 2 - assert dual_cache.get_cache(key=otpm_key) == 2 - - @pytest.mark.asyncio - async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): - import logging - - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - deployment = { - "litellm_params": { - "model": deployment_name, - "itpm": 100, - "rpm": 1, - }, - "model_info": {"id": model_id}, - "model_name": "opus", - } - - minute = get_utc_datetime().strftime("%H-%M") - rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) - - request_kwargs = { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - assert await dual_cache.async_get_cache(key=rpm_key) == 6 - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - assert any("both limit types are enforced" in record.message for record in caplog.records) - - @pytest.mark.asyncio - async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): - """ - A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on - success, otherwise the tpm_key the pre-call check reads is never written - and the tpm_limit can never be enforced. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - await check.async_log_success_event(kwargs, response, None, None) - - # ITPM reconciled down from the 5-token reservation to actual usage (3). - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 - - def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-sync-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - dual_cache.set_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 - - @pytest.mark.asyncio - async def test_failure_refunds_itpm_reservation(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} - kwargs = { - "standard_logging_object": { - "model_id": "io-refund-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": dict(reservation), - }, - "metadata": dict(reservation), - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - -class TestRouterIOTokenIntegration: - @pytest.mark.asyncio - async def test_model_group_info_aggregates_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "opus", - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 20, - }, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - info = router.get_model_group_info("opus") - assert info is not None - assert info.itpm == 100 - assert info.otpm == 20 - - -class TestContextSlotRetention: - def test_setter_stores_kwargs_only_for_io_limited_deployments(self): - """ - The context slot pins the entire request kwargs (messages included) - for the lifetime of the surrounding asyncio context, and pooled - resources created mid-request (e.g. redis connections) capture that - context, extending the pin far past the request. Only ITPM/OTPM - pre-call checks read the slot, so the setter must store None for - deployments without io token limits and still clear reservation - sentinels from kwargs either way. - """ - kwargs = { - "messages": [{"role": "user", "content": "x" * 1000}], - "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, - } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - assert ITPM_CACHE_KEY not in kwargs["metadata"] - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_does_not_pin_kwargs_without_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, - } - ] - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("plain") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_pins_kwargs_for_io_limited_deployment(self): - router = Router( - model_list=[ - { - "model_name": "limited", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("limited") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - -@pytest.mark.asyncio -async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): - from litellm.utils import get_utc_datetime - from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( - assert_loop_stayed_free, - timed_with_loop_lags, - warm_tokenizer, - ) - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - warm_tokenizer("anthropic/claude-fable-5") - deployment = { - "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, - "model_info": {"id": "io-loop-id"}, - "model_name": "claude", - } - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) - - _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) - - minute = get_utc_datetime().strftime("%H-%M") - reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") - assert reserved > 100_000 - assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/types/llms/test_types_llms_bedrock.py b/tests/test_litellm/types/llms/test_types_llms_bedrock.py deleted file mode 100644 index a5ad882e775..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_bedrock.py +++ /dev/null @@ -1,46 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams - - -def test_model_validate_keeps_auth_params_and_ignores_request_params(): - auth_params = AwsAuthParams.model_validate( - { - "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", - "aws_session_name": "litellm-session", - "aws_external_id": "litellm-external-id", - "aws_region_name": "us-west-2", - "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", - "model": "anthropic.claude-haiku-4-5-20251001-v1:0", - "temperature": 0.1, - "messages": [{"role": "user", "content": "hi"}], - } - ) - - assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" - assert auth_params.aws_session_name == "litellm-session" - assert auth_params.aws_external_id == "litellm-external-id" - assert auth_params.aws_access_key_id is None - assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) - assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("aws_role_name", 1234), - ("aws_session_name", ["litellm-session"]), - ("aws_external_id", {"id": "x"}), - ], -) -def test_model_validate_rejects_non_string_credentials(field, value): - with pytest.raises(ValidationError): - AwsAuthParams.model_validate({field: value}) - - -def test_frozen_struct_rejects_field_assignment(): - auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") - - with pytest.raises(ValidationError): - auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py deleted file mode 100644 index 64ec09838e8..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ /dev/null @@ -1,591 +0,0 @@ -import asyncio -from typing import Optional -from unittest.mock import AsyncMock, patch - -import pytest - -import json - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent - - -@pytest.mark.parametrize("stream", (False, True)) -def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: - from typing import Final - - from litellm.types.llms.openai import ( - ChatCompletionReasoningItem, - ChatCompletionReasoningSummaryTextBlock, - ) - from litellm.types.utils import ( - Choices, - Delta, - Message, - ModelResponse, - ModelResponseStream, - StreamingChoices, - ) - - reasoning_item: Final = ChatCompletionReasoningItem( - type="reasoning", - id="rs_123", - encrypted_content="encrypted", - summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], - ) - response: Final = ( - ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) - if stream - else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) - ) - message_key: Final = "delta" if stream else "message" - assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - restored: Final = type(response).model_validate_json(response.model_dump_json()) - assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - -def test_generic_event(): - from litellm.types.llms.openai import GenericEvent - - event = {"type": "test", "test": "test"} - event = GenericEvent(**event) - assert event.type == "test" - assert event.test == "test" - - -def test_output_item_added_event(): - from litellm.types.llms.openai import OutputItemAddedEvent - - event = { - "type": "response.output_item.added", - "sequence_number": 4, - "output_index": 1, - "item": None, - } - event = OutputItemAddedEvent(**event) - assert event.type == "response.output_item.added" - assert event.sequence_number == 4 - assert event.output_index == 1 - assert event.item is None - - -class TestResponsesAPIResponseOutputText: - """Tests for the output_text property on ResponsesAPIResponse""" - - def test_output_text_with_single_message(self): - """Test output_text with a single message containing text output""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_123", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Hello, world!", - } - ], - } - ], - ) - - assert response.output_text == "Hello, world!" - - def test_output_text_with_multiple_messages(self): - """Test output_text with multiple messages aggregates all text""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "First part. ", - } - ], - }, - { - "type": "message", - "id": "msg_2", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Second part.", - } - ], - }, - ], - ) - - assert response.output_text == "First part. Second part." - - def test_output_text_with_no_text_content(self): - """Test output_text returns empty string when no output_text content exists""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - } - ], - ) - - assert response.output_text == "" - - def test_output_text_with_mixed_content(self): - """Test output_text only aggregates output_text type content""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "The weather is sunny. ", - }, - { - "type": "refusal", - "refusal": "I cannot do that.", - }, - ], - }, - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - }, - ], - ) - - assert response.output_text == "The weather is sunny. " - - def test_output_text_with_empty_output(self): - """Test output_text returns empty string with empty output list""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[], - ) - - assert response.output_text == "" - - -class TestAssistantMessageImageUrlContent: - """ - Regression tests for image_url blocks in assistant message content. - - Bug: ChatCompletionAssistantMessage.content did not include - ChatCompletionImageObject in its union, so Pydantic v2 silently dropped - image_url blocks (content → []) when serialising via AllMessageValues. - This affects users who store conversation history as JSON (e.g. in a DB) - and read it back typed as list[AllMessageValues]. - """ - - ASSISTANT_MESSAGE_WITH_IMAGE = { - "role": "assistant", - "content": [ - {"type": "text", "text": "Here is the image you requested:"}, - { - "type": "image_url", - "image_url": { - "url": ( - "data:image/png;base64," - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" - "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - ) - }, - }, - ], - } - - def test_assistant_message_image_url_preserved_single(self): - """ - TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive - validate_python → dump_python without being dropped or raising an error. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import ChatCompletionAssistantMessage - - adapter = TypeAdapter(ChatCompletionAssistantMessage) - validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) - dumped = adapter.dump_python(validated) - - raw_content = dumped.get("content") - # Pydantic may return a lazy SerializationIterator for Iterable fields; - # convert to list to consume it — this must not raise ValidationError. - content_blocks = list(raw_content) if raw_content is not None else [] - - assert ( - len(content_blocks) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" - types = [b.get("type") for b in content_blocks if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped; blocks: {content_blocks}" - - def test_assistant_message_image_url_preserved_in_all_message_values(self): - """ - TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an - assistant message must not be silently dropped during dump_python(mode='json'). - - This is the primary failing path: conversation history stored as JSON in a - database and read back typed as list[AllMessageValues]. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import AllMessageValues - - conversation = [ - { - "role": "user", - "content": "Generate an image of a banana wearing a LiteLLM costume", - }, - self.ASSISTANT_MESSAGE_WITH_IMAGE, - ] - - adapter = TypeAdapter(List[AllMessageValues]) - validated = adapter.validate_python(conversation) - dumped = adapter.dump_python(validated, mode="json") - - assistant = next((m for m in dumped if m.get("role") == "assistant"), None) - assert assistant is not None, "Assistant message missing after serialisation" - - content = assistant.get("content", []) - assert isinstance( - content, list - ), f"content should be a list, got {type(content)}" - assert ( - len(content) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" - types = [b.get("type") for b in content if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" - - -class TestResponsesAPIReasoningNullFields: - """ - Tests for issue #16824: reasoning output items should not include null - status/content/encrypted_content fields. - - When a provider returns reasoning items without these fields, LiteLLM's - Pydantic parsing adds them as Optional defaults (None). Serializing them - as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on - status=null). - - The fix uses a field_serializer on ResponsesAPIResponse.output that - mirrors the request-side filtering in - OpenAIResponsesAPIConfig._handle_reasoning_item(). - """ - - def _make_response(self, output): - from litellm.types.llms.openai import ResponsesAPIResponse - - return ResponsesAPIResponse( - id="resp_test", - created_at=1741476542, - model="gpt-5-mini", - object="response", - status="completed", - output=output, - ) - - def test_reasoning_item_null_fields_removed_model_dump(self): - """Null status/content/encrypted_content should be absent from model_dump.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_null_fields_removed_model_dump_json(self): - """Null fields should also be absent from model_dump_json.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - parsed = json.loads(response.model_dump_json()) - reasoning = parsed["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_non_null_values_preserved(self): - """Non-null values on reasoning items should be kept.""" - response = self._make_response( - output=[ - { - "id": "rs_abc", - "type": "reasoning", - "summary": [], - "status": "completed", - "encrypted_content": "gAAAA...", - } - ] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["status"] == "completed" - assert reasoning["encrypted_content"] == "gAAAA..." - - def test_message_item_not_affected(self): - """Non-reasoning output items should keep all their fields.""" - response = self._make_response( - output=[ - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Hello!", - "annotations": [], - } - ], - } - ] - ) - dumped = response.model_dump() - message = dumped["output"][0] - assert message["status"] == "completed" - assert message["type"] == "message" - assert len(message["content"]) == 1 - - def test_mixed_output_reasoning_and_message(self): - """Reasoning items cleaned, message items untouched in same response.""" - response = self._make_response( - output=[ - {"id": "rs_abc", "type": "reasoning", "summary": []}, - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Answer", - "annotations": [], - } - ], - }, - ] - ) - dumped = response.model_dump() - reasoning = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "reasoning" - ][0] - message = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "message" - ][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert message["status"] == "completed" - assert len(message["content"]) == 1 - - def test_reasoning_core_fields_preserved(self): - """id, type, summary should always be present on reasoning items.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["id"] == "rs_abc" - assert reasoning["type"] == "reasoning" - assert reasoning["summary"] == ["thinking..."] - - def test_top_level_null_fields_unaffected(self): - """Top-level response fields with None should not be affected.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - assert "error" in dumped - assert dumped["error"] is None - assert "instructions" in dumped - assert dumped["instructions"] is None - - -def test_normalize_fine_tuning_job_dict_maps_azure_pending(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - out = _normalize_fine_tuning_job_dict( - {"organization_id": None, "result_files": None, "status": "pending"}, - is_azure=True, - ) - assert out["organization_id"] == "" - assert out["result_files"] == [] - assert out["status"] == "queued" - - -def test_normalize_fine_tuning_job_dict_openai_unchanged(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - data = {"organization_id": None, "result_files": None, "status": "pending"} - out = _normalize_fine_tuning_job_dict(data, is_azure=False) - assert out is data - - -def test_openai_file_object_accepts_pending_status(): - from litellm.types.llms.openai import OpenAIFileObject - - file_obj = OpenAIFileObject( - id="file-123", - bytes=1024, - created_at=1677610602, - filename="train.jsonl", - object="file", - purpose="fine-tune", - status="pending", - ) - assert file_obj.status == "pending" - - -class TestOpenAIFileObjectBatchGuardrailSerialization: - """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" - - @staticmethod - def _file_object(**overrides): - from litellm.types.llms.openai import OpenAIFileObject - - return OpenAIFileObject( - id="file-123", - object="file", - bytes=1024, - created_at=1677610602, - filename="batch.jsonl", - purpose="batch", - status="uploaded", - **overrides, - ) - - @staticmethod - def _report(): - from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport - - return BatchGuardrailReport( - submitted_records=3, - modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), - ) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_absent_when_unset(self, mode): - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_present_when_set(self, mode): - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) - assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 - - def test_nested_nulls_of_a_set_report_survive(self): - """`exclude_none=True` was rejected as the fix because it would strip these.""" - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") - assert dumped["litellm_batch_guardrail"]["modified_records"] == [ - {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} - ] - - def test_by_alias_dump_also_omits_the_key(self): - """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) - - def test_other_optional_fields_still_serialize_as_null(self): - dumped = self._file_object().model_dump(mode="json") - assert dumped["expires_at"] is None - assert dumped["status_details"] is None - - def test_round_trip_of_a_set_report_is_lossless(self): - from litellm.types.llms.openai import OpenAIFileObject - - original = self._file_object(litellm_batch_guardrail=self._report()) - assert OpenAIFileObject(**original.model_dump()) == original - - def test_serialization_json_schema_still_describes_the_model(self): - """A return annotation on the wrap serializer would collapse this to a bare object.""" - from litellm.types.llms.openai import OpenAIFileObject - - schema = OpenAIFileObject.model_json_schema(mode="serialization") - assert "litellm_batch_guardrail" in schema["properties"] - - def test_key_omitted_inside_a_file_list_page(self): - from litellm.types.llms.openai import FileListPage - - page = FileListPage(object="list", data=[self._file_object()], has_more=False) - assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] - - -def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: - import httpx - - return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) - - -def test_httpx_binary_response_content_hidden_params_are_per_instance(): - first = _binary_content(b"first") - second = _binary_content(b"second") - - first._hidden_params["response_cost"] = 0.5 - - assert second._hidden_params == {} - - -def test_set_response_cost_none_leaves_hidden_params_empty(): - binary_response = _binary_content(b"audio") - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params - - binary_response.set_response_cost(0.25) - - assert binary_response._hidden_params["response_cost"] == 0.25 - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params diff --git a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py deleted file mode 100644 index 2e5986d3ef8..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -Tests for pipeline type definitions. -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.pipeline_types import ( - GuardrailPipeline, - PipelineExecutionResult, - PipelineStep, - PipelineStepResult, -) -from litellm.types.proxy.policy_engine.policy_types import ( - Policy, - PolicyGuardrails, -) - - -def test_pipeline_step_defaults(): - step = PipelineStep(guardrail="my-guard") - assert step.on_fail == "block" - assert step.on_pass == "allow" - assert step.on_error is None - assert step.pass_data is False - assert step.modify_response_message is None - - -def test_pipeline_step_valid_actions(): - step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") - assert step.on_fail == "next" - assert step.on_pass == "next" - - -def test_pipeline_step_all_action_types(): - for action in ("allow", "block", "next", "modify_response"): - step = PipelineStep( - guardrail="g", on_fail=action, on_pass=action, on_error=action - ) - assert step.on_fail == action - assert step.on_pass == action - assert step.on_error == action - - -def test_pipeline_step_invalid_action_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_fail="invalid_action") - - -def test_pipeline_step_invalid_on_pass_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_pass="skip") - - -def test_pipeline_step_on_error_valid(): - step = PipelineStep( - guardrail="g", on_error="next", on_fail="block", on_pass="allow" - ) - assert step.on_error == "next" - - -def test_pipeline_step_invalid_on_error_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_error="invalid") - - -def test_pipeline_requires_at_least_one_step(): - with pytest.raises(ValidationError): - GuardrailPipeline(mode="pre_call", steps=[]) - - -def test_pipeline_invalid_mode_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="during_call", - steps=[PipelineStep(guardrail="g")], - ) - - -def test_pipeline_valid_modes(): - for mode in ("pre_call", "post_call"): - pipeline = GuardrailPipeline( - mode=mode, - steps=[PipelineStep(guardrail="g")], - ) - assert pipeline.mode == mode - - -def test_pipeline_with_multiple_steps(): - pipeline = GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), - PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), - ], - ) - assert len(pipeline.steps) == 2 - assert pipeline.steps[0].guardrail == "g1" - assert pipeline.steps[1].guardrail == "g2" - - -def test_policy_with_pipeline_parses(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1", "g2"]), - pipeline=GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next"), - PipelineStep(guardrail="g2"), - ], - ), - ) - assert policy.pipeline is not None - assert len(policy.pipeline.steps) == 2 - - -def test_policy_without_pipeline(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1"]), - ) - assert policy.pipeline is None - - -def test_pipeline_step_result(): - result = PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - error_detail="Content policy violation", - duration_seconds=0.05, - ) - assert result.outcome == "fail" - assert result.action_taken == "next" - - -def test_pipeline_execution_result(): - result = PipelineExecutionResult( - terminal_action="block", - step_results=[ - PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - ), - PipelineStepResult( - guardrail_name="g2", - outcome="fail", - action_taken="block", - ), - ], - error_message="Content blocked", - ) - assert result.terminal_action == "block" - assert len(result.step_results) == 2 - - -def test_pipeline_step_extra_fields_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="g", unknown_field="value") - - -def test_pipeline_extra_fields_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="pre_call", - steps=[PipelineStep(guardrail="g")], - unknown="value", - ) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py b/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py deleted file mode 100644 index bcd6d39aa4d..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py +++ /dev/null @@ -1,15 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_accepts_int32_priority(priority: int): - assert PolicyAttachment(policy="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachment(policy="p", priority=priority) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py b/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py deleted file mode 100644 index f31b9d7e873..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Tests for pipeline field on policy CRUD types (resolver_types.py). -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.resolver_types import ( - PolicyAttachmentCreateRequest, - PolicyCreateRequest, - PolicyDBResponse, - PolicyUpdateRequest, -) - - -def test_policy_create_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert req.pipeline is not None - assert req.pipeline["mode"] == "pre_call" - assert len(req.pipeline["steps"]) == 2 - - -def test_policy_create_request_without_pipeline(): - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1"], - ) - assert req.pipeline is None - - -def test_policy_update_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyUpdateRequest(pipeline=pipeline_data) - assert req.pipeline is not None - assert req.pipeline["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert resp.pipeline is not None - assert resp.pipeline["mode"] == "pre_call" - dumped = resp.model_dump() - assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_without_pipeline(): - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - ) - assert resp.pipeline is None - dumped = resp.model_dump() - assert dumped["pipeline"] is None - - -def test_policy_create_request_roundtrip(): - pipeline_data = { - "mode": "post_call", - "steps": [ - { - "guardrail": "g1", - "on_fail": "modify_response", - "on_pass": "next", - "pass_data": True, - "modify_response_message": "custom msg", - }, - ], - } - req = PolicyCreateRequest( - policy_name="roundtrip-test", - guardrails_add=["g1"], - pipeline=pipeline_data, - ) - dumped = req.model_dump() - restored = PolicyCreateRequest(**dumped) - assert restored.pipeline == pipeline_data - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_create_request_accepts_int32_priority(priority: int): - assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/test_litellm/videos/test_main.py b/tests/test_litellm/videos/test_main.py deleted file mode 100644 index 22e1e5c05eb..00000000000 --- a/tests/test_litellm/videos/test_main.py +++ /dev/null @@ -1,455 +0,0 @@ -""" -Dispatch-contract tests for litellm/videos/main.py - -Each public video operation is a pair: a sync `video_*` worker (decorated with -@client) that resolves the provider, fetches the provider config, logs, and then -forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async -`avideo_*` wrapper that delegates to the sync worker in an executor. - -This file locks the contract of that layer so a regression fails loudly: - - 1. DISPATCH - the one correct handler fired and every sibling video handler - asserted NOT called. A copy-paste that calls the wrong handler - (e.g. remix -> edit) flips this. - 2. RESULT - the handler's return value is propagated by identity. - 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not - passed (status/content/remix/edit/extension), or defaults to - "openai" (list/create_character/get_character). This is the exact - surface of the historical "content defaulted to openai" bug. - 4. PAYLOAD - the provider config object and the operation's identifying args - (video_id/prompt/name/...) reach the handler; _is_async is False - on the sync path. - 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. - 6. UNSUPPORTED - a None provider config raises before any handler fires. - 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets - async_call=True, and pre-resolves the provider where it must. - -Seams mocked: the http handler (network), the provider-config registry lookup, -get_llm_provider, and the video-generation optional-param builders. The id decode -helper runs for real against genuinely-encoded ids, so the provider assertions -reflect production. -""" - -from contextlib import ExitStack -from dataclasses import dataclass -from typing import Any, Dict -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.types.videos.main import CharacterObject, VideoObject -from litellm.types.videos.utils import encode_video_id_with_provider -from litellm.videos import main as videos_main - -# A real model-encoded video id: decodes (for real) to provider "azure". Used to -# prove the sync workers derive custom_llm_provider from the id, not a hardcode. -AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") - -# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly -# one fired and the other eight did not. -SYNC_HANDLERS = ( - "video_generation_handler", - "video_content_handler", - "video_remix_handler", - "video_create_character_handler", - "video_get_character_handler", - "video_edit_handler", - "video_extension_handler", - "video_list_handler", - "video_status_handler", -) - -GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} - - -@dataclass -class Seams: - handler: MagicMock - get_config: MagicMock - config: MagicMock - - def kwargs_of(self, handler_name: str) -> Dict[str, Any]: - method = getattr(self.handler, handler_name) - assert method.call_count == 1 - return dict(method.call_args.kwargs) - - def assert_only(self, handler_name: str) -> None: - for name in SYNC_HANDLERS: - method = getattr(self.handler, name) - if name == handler_name: - method.assert_called_once() - else: - method.assert_not_called() - - -@pytest.fixture -def seams(): - handler = MagicMock(spec=BaseLLMHTTPHandler) - config = MagicMock(name="provider_video_config") - get_config = MagicMock(return_value=config) - - with ExitStack() as stack: - stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) - stack.enter_context( - patch.object( - videos_main.ProviderConfigManager, - "get_provider_video_config", - get_config, - ) - ) - # video_generation resolves model+provider through get_llm_provider and - # builds optional params; mock those so the dispatch payload is deterministic. - stack.enter_context( - patch.object( - videos_main, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_requested_video_generation_optional_param", - MagicMock(return_value={"seconds": "8"}), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_optional_params_video_generation", - MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), - ) - ) - yield Seams(handler=handler, get_config=get_config, config=config) - - -# =========================================================================== # -# Dispatch contract - one rich test per sync worker. -# =========================================================================== # - - -def test_video_generation__dispatch(seams): - result = videos_main.video_generation(prompt="a sunset", model="sora-2") - - seams.assert_only("video_generation_handler") - assert result is seams.handler.video_generation_handler.return_value - kw = seams.kwargs_of("video_generation_handler") - assert kw["model"] == "sora-2" - assert kw["prompt"] == "a sunset" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_generation_provider_config"] is seams.config - assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS - assert kw["_is_async"] is False - - -def test_video_status__dispatch_and_provider_from_id(seams): - result = videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.assert_only("video_status_handler") - assert result is seams.handler.video_status_handler.return_value - kw = seams.kwargs_of("video_status_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai - assert kw["video_status_provider_config"] is seams.config - assert kw["_is_async"] is False - # provider config requested for the decoded provider, not a hardcode. - assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE - - -def test_video_content__dispatch_and_provider_from_id(seams): - result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") - - seams.assert_only("video_content_handler") - assert result is seams.handler.video_content_handler.return_value - kw = seams.kwargs_of("video_content_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" - assert kw["variant"] == "thumbnail" - assert kw["video_content_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_content__plain_id_defaults_to_openai(seams): - videos_main.video_content(video_id="video_plain") - - assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" - - -def test_video_remix__dispatch_and_provider_from_id(seams): - result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") - - seams.assert_only("video_remix_handler") - assert result is seams.handler.video_remix_handler.return_value - kw = seams.kwargs_of("video_remix_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "new colors" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_remix_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_edit__dispatch_and_provider_from_id(seams): - result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") - - seams.assert_only("video_edit_handler") - assert result is seams.handler.video_edit_handler.return_value - kw = seams.kwargs_of("video_edit_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "brighter" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) - - seams.assert_only("video_extension_handler") - assert result is seams.handler.video_extension_handler.return_value - kw = seams.kwargs_of("video_extension_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "continue" - assert kw["seconds"] == "5" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_list__dispatch_defaults_to_openai(seams): - result = videos_main.video_list(after="cur", limit=5, order="desc") - - seams.assert_only("video_list_handler") - assert result is seams.handler.video_list_handler.return_value - kw = seams.kwargs_of("video_list_handler") - assert kw["after"] == "cur" - assert kw["limit"] == 5 - assert kw["order"] == "desc" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_list_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_create_character__dispatch_defaults_to_openai(seams): - video = MagicMock(name="video_upload") - result = videos_main.video_create_character(name="hero", video=video) - - seams.assert_only("video_create_character_handler") - assert result is seams.handler.video_create_character_handler.return_value - kw = seams.kwargs_of("video_create_character_handler") - assert kw["name"] == "hero" - assert kw["video"] is video - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_get_character__dispatch_defaults_to_openai(seams): - result = videos_main.video_get_character(character_id="char_1") - - seams.assert_only("video_get_character_handler") - assert result is seams.handler.video_get_character_handler.return_value - kw = seams.kwargs_of("video_get_character_handler") - assert kw["character_id"] == "char_1" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_explicit_provider_beats_decoded_id(seams): - """An explicit custom_llm_provider wins over the one encoded in the id.""" - videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") - - assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" - - -# =========================================================================== # -# mock_response short-circuit - returns a typed object, no handler call. -# =========================================================================== # - - -def test_generation__mock_response_short_circuits(seams): - resp = videos_main.video_generation( - prompt="x", - model="sora-2", - mock_response={"id": "v1", "object": "video", "status": "queued"}, - ) - - assert isinstance(resp, VideoObject) - assert resp.id == "v1" - seams.handler.video_generation_handler.assert_not_called() - - -def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) - - assert isinstance(resp, list) - assert resp[0].id == "v1" - seams.handler.video_list_handler.assert_not_called() - - -def test_get_character__mock_response_short_circuits(seams): - resp = videos_main.video_get_character( - character_id="char_1", - mock_response={ - "id": "char_1", - "object": "character", - "created_at": 1, - "name": "hero", - }, - ) - - assert isinstance(resp, CharacterObject) - assert resp.id == "char_1" - seams.handler.video_get_character_handler.assert_not_called() - - -# =========================================================================== # -# Unsupported provider - a None provider config raises before any dispatch. -# =========================================================================== # - - -def test_unsupported_provider_raises_without_dispatch(seams): - seams.get_config.return_value = None - - with pytest.raises(litellm.APIConnectionError): - videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.handler.video_status_handler.assert_not_called() - - -# =========================================================================== # -# Async-wrapper delegation - representative coverage. -# =========================================================================== # - - -@pytest.mark.asyncio -async def test_avideo_generation__delegates_with_async_flag(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, - patch.object( - litellm, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ), - ): - result = await videos_main.avideo_generation(prompt="x", model="sora-2") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "openai" - - -@pytest.mark.asyncio -async def test_avideo_status__delegates_untouched(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_status(video_id="video_plain") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["video_id"] == "video_plain" - - -@pytest.mark.asyncio -async def test_avideo_content__pre_decodes_provider_before_delegating(): - """avideo_content resolves the provider from the encoded id itself before - handing off, so the sync worker receives the decoded provider, not None.""" - sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "azure" - - -# =========================================================================== # -# Credential passthrough - DB/YAML model-config credentials the router injects -# via kwargs must reach the provider call for EVERY video handler, carried in -# litellm_params. Distinct per-field values catch a cross-wired field. -# =========================================================================== # - -DB_YAML_CREDS = { - "api_key": "sk-db-credential", - "api_base": "https://db-resource.test", - "api_version": "2024-12-31", - "vertex_project": "db-project-xyz", -} - -CREDENTIAL_OPERATIONS = [ - ( - "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), - ), - ( - "video_status_handler", - lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_content_handler", - lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), - ), - ( - "video_list_handler", - lambda: videos_main.video_list(**DB_YAML_CREDS), - ), - ( - "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), - ), - ( - "video_get_character_handler", - lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), - ), -] - - -@pytest.mark.parametrize( - "handler_name,invoke", - CREDENTIAL_OPERATIONS, - ids=[op[0] for op in CREDENTIAL_OPERATIONS], -) -def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): - invoke() - - litellm_params = seams.kwargs_of(handler_name)["litellm_params"] - assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] - assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] - assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] - assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/test_litellm/videos/test_utils.py b/tests/test_litellm/videos/test_utils.py deleted file mode 100644 index 57fb549c23d..00000000000 --- a/tests/test_litellm/videos/test_utils.py +++ /dev/null @@ -1,193 +0,0 @@ -""" -Pure-logic contract tests for litellm/videos/main.py's request utils -(litellm/videos/utils.py: VideoGenerationRequestUtils). - -These lock the exact param-shaping behavior so a mutation that drops a filter, -flips a precedence, or stops removing a key fails loudly. The only seam is the -provider config's map_openai_params (a provider boundary); filter_out_litellm_params -runs for real, so the "litellm-internal params get stripped" assertions reflect -production. Every test asserts the exact resulting dict, never "ran without error". -""" - -from unittest.mock import MagicMock - - - -import litellm -from litellm.videos.utils import VideoGenerationRequestUtils - -get_requested = ( - VideoGenerationRequestUtils.get_requested_video_generation_optional_param -) -get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation - - -# =========================================================================== # -# get_requested_video_generation_optional_param -# -# Receives the caller's full local_vars; must return only the API-bound optional -# params. filter_out_litellm_params strips known internal keys for real; the -# values used below were chosen against the live set: seconds/size/user/foo_param/ -# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. -# =========================================================================== # - - -def test_requested__drops_none_and_excluded_keys(): - result = get_requested( - { - "seconds": "8", - "size": None, # None -> dropped - "prompt": "a sunset", # excluded - "model": "sora-2", # excluded - "user": "u1", - } - ) - assert result == {"seconds": "8", "user": "u1"} - - -def test_requested__strips_litellm_internal_params(): - result = get_requested( - { - "seconds": "8", - "api_key": "sk-secret", - "metadata": {"x": 1}, - "litellm_call_id": "id-123", - } - ) - assert result == {"seconds": "8"} - - -def test_requested__timeout_always_removed(): - # timeout is NOT a litellm-internal param, so only the explicit pop removes it. - result = get_requested({"seconds": "8", "timeout": 30}) - assert result == {"seconds": "8"} - - -def test_requested__nested_kwargs_merge_and_override_base(): - result = get_requested( - {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} - ) - # nested kwargs win over the top-level base params on collision. - assert result == {"seconds": "override", "size": "720x1280"} - - -def test_requested__non_dict_kwargs_treated_as_empty(): - result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) - assert result == {"seconds": "8"} - - -def test_requested__none_input_returns_empty(): - assert get_requested(None) == {} - - -def test_requested__top_level_extra_body_spread_and_preserved(): - result = get_requested( - {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} - ) - # extra_body keys are both spread at top level AND kept under "extra_body". - assert result == { - "seconds": "8", - "vertex_project": "proj", - "foo_param": "bar", - "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, - } - - -def test_requested__extra_body_kwargs_overrides_top_level(): - result = get_requested( - { - "extra_body": {"a": "top", "b": "top_b"}, - "kwargs": {"extra_body": {"a": "kw"}}, - } - ) - # kwargs' extra_body wins over the top-level extra_body on collision; the - # non-colliding top-level key survives. - assert result == { - "a": "kw", - "b": "top_b", - "extra_body": {"a": "kw", "b": "top_b"}, - } - - -def test_requested__extra_body_strips_litellm_internal_params(): - result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) - # api_key filtered out of extra_body; only foo_param remains (and is spread). - assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} - - -def test_requested__empty_extra_body_not_added(): - result = get_requested({"seconds": "8", "extra_body": {}}) - assert result == {"seconds": "8"} - assert "extra_body" not in result - - -# =========================================================================== # -# get_optional_params_video_generation -# -# Delegates mapping to the provider config (the seam) then folds extra_body in. -# =========================================================================== # - - -def _config(map_return): - config = MagicMock() - config.map_openai_params.return_value = map_return - return config - - -def test_optional__delegates_to_map_openai_params_with_drop_params(): - config = _config({"seconds": "8"}) - optional_params = {"seconds": "8"} - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params=optional_params, - ) - - assert result == {"seconds": "8"} - config.map_openai_params.assert_called_once_with( - video_create_optional_params=optional_params, - model="sora-2", - drop_params=litellm.drop_params, - ) - - -def test_optional__extra_body_overrides_mapped_and_is_removed(): - # mapped output carries a leftover extra_body that must be popped; the input - # extra_body overrides a colliding mapped key and is spread in. - config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={ - "extra_body": {"size": "override", "extra": "x"} - }, - ) - - assert result == {"seconds": "8", "size": "override", "extra": "x"} - assert "extra_body" not in result - - -def test_optional__no_extra_body_returns_mapped_unchanged(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8"}, - ) - - assert result == {"seconds": "8"} - - -def test_optional__non_dict_extra_body_ignored(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8", "extra_body": None}, - ) - - assert result == {"seconds": "8"} From 72abd11b4ce43fb3762b008d97e60c95bc76e313 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:32 +0000 Subject: [PATCH 068/127] test: assert the responses bridge forwards aws_region_name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 tests/unit/completion_extras/test_responses_bridge_provider_propagation.py diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py new file mode 100644 index 00000000000..09ef1889818 --- /dev/null +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -0,0 +1,59 @@ +from datetime import datetime +from unittest.mock import patch + +import pytest + +from litellm.completion_extras.litellm_responses_transformation.handler import ( + ResponsesToCompletionBridgeHandler, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.utils import ModelResponse + +MODEL = "openai.gpt-5.5" +REGION = "us-east-2" + + +def _bedrock_mantle_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="acompletion", + model=MODEL, + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": MODEL, + "custom_llm_provider": "bedrock_mantle", + "messages": messages, + "optional_params": {}, + "litellm_params": { + "aws_region_name": REGION, + "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", + "custom_llm_provider": "bedrock_mantle", + }, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +@pytest.mark.asyncio +async def test_acompletion_forwards_aws_region_name_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model=MODEL) + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_bedrock_mantle_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["aws_region_name"] == REGION + assert _fake_aresponses.kwargs["custom_llm_provider"] == "bedrock_mantle" From a69ca90ea5380bde22ad36ac9276604bae28652f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:40 +0000 Subject: [PATCH 069/127] test: migrate phase 14 wave 2 provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...sonx_audio_transcription_transformation.py | 337 ---------- .../test_litellm/llms/watsonx/test_watsonx.py | 577 ------------------ .../llms/xai/xai_responses/__init__.py | 1 - .../xai/xai_responses/test_transformation.py | 105 ---- .../count_tokens/__init__.py | 0 .../test_count_tokens_location.py | 2 + .../test_count_tokens_no_vertexai_sdk.py | 0 ...ai_partner_models_llama3_transformation.py | 0 ...i_partner_models_mistral_transformation.py | 15 + .../vertex_ai/vertex_gemma_models/__init__.py | 0 .../test_vertex_gemma_transformation.py | 0 tests/unit/llms/vertex_ai/videos/__init__.py | 0 .../test_vertex_video_transformation.py | 0 ...est_volcengine_responses_transformation.py | 24 - tests/unit/llms/voyage/rerank/__init__.py | 0 .../test_voyage_rerank_transformation.py | 0 .../test_voyage_contextual_embedding.py | 0 .../test_voyage_multimodal_embedding.py | 0 tests/unit/llms/watsonx/__init__.py | 0 .../watsonx/audio_transcription/__init__.py | 0 ...sonx_audio_transcription_transformation.py | 85 +++ .../test_watsonx_embedding_transformation.py | 0 ...test_watsonx_passthrough_transformation.py | 0 tests/unit/llms/watsonx/rerank/__init__.py | 0 .../watsonx/rerank/test_watsonx_rerank.py | 0 tests/unit/llms/watsonx/test_watsonx.py | 74 +++ .../llms/watsonx/test_watsonx_common_utils.py | 0 .../test_xai_responses_transformation.py | 0 tests/unit/llms/you_com/__init__.py | 0 .../llms/you_com/test_you_com_search.py | 0 .../llms/zai/test_zai_provider.py | 0 31 files changed, 176 insertions(+), 1044 deletions(-) delete mode 100644 tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py delete mode 100644 tests/test_litellm/llms/watsonx/test_watsonx.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/__init__.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/test_transformation.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py (98%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py (60%) create mode 100644 tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/videos/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/videos/test_vertex_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/volcengine/responses/test_volcengine_responses_transformation.py (94%) create mode 100644 tests/unit/llms/voyage/rerank/__init__.py rename tests/{test_litellm => unit}/llms/voyage/rerank/test_voyage_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_contextual_embedding.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_multimodal_embedding.py (100%) create mode 100644 tests/unit/llms/watsonx/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py rename tests/{test_litellm => unit}/llms/watsonx/embed/test_watsonx_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py (100%) create mode 100644 tests/unit/llms/watsonx/rerank/__init__.py rename tests/{test_litellm => unit}/llms/watsonx/rerank/test_watsonx_rerank.py (100%) create mode 100644 tests/unit/llms/watsonx/test_watsonx.py rename tests/{test_litellm => unit}/llms/watsonx/test_watsonx_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/xai/responses/test_xai_responses_transformation.py (100%) create mode 100644 tests/unit/llms/you_com/__init__.py rename tests/{test_litellm => unit}/llms/you_com/test_you_com_search.py (100%) rename tests/{test_litellm => unit}/llms/zai/test_zai_provider.py (100%) diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py deleted file mode 100644 index e269e782061..00000000000 --- a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -Tests for IBM WatsonX Audio Transcription. - -Validates that litellm.transcription transforms requests correctly for WatsonX. -""" - -import json -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.watsonx.audio_transcription.transformation import ( - IBMWatsonXAudioTranscriptionConfig, -) -from litellm.types.utils import TranscriptionResponse - - -class TestWatsonXAudioTranscription: - """Tests for WatsonX audio transcription via litellm.transcription.""" - - @pytest.mark.asyncio - async def test_watsonx_transcription_url_and_headers(self): - """ - Test that litellm.transcription sends request to correct WatsonX URL with proper headers. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["url"] = str(kwargs.get("url", args[0] if args else None)) - captured_request["headers"] = kwargs.get("headers", {}) - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - # Validate URL contains WatsonX audio transcription endpoint - assert "/ml/v1/audio/transcriptions" in captured_request["url"] - assert "version=" in captured_request["url"] - # project_id should NOT be in URL (it should be in form data instead) - assert "project_id=test-project-123" not in captured_request["url"] - - # Validate headers contain WatsonX auth - assert "Authorization" in captured_request["headers"] - assert ( - "Bearer test-bearer-token" in captured_request["headers"]["Authorization"] - ) - - # Validate Content-Type is NOT set (httpx sets multipart/form-data automatically) - assert "Content-Type" not in captured_request["headers"] - - # Validate project_id is in form data, not URL - assert captured_request["data"].get("project_id") == "test-project-123" - - # Validate file is in files dict - assert "file" in captured_request["files"] - - @pytest.mark.asyncio - async def test_watsonx_transcription_request_body(self): - """ - Test that litellm.transcription sends correct request body for WatsonX. - - Validates that: - - Request uses multipart/form-data (data + files) - - Model name has watsonx/ prefix removed - - project_id is in form data, not URL - - Audio file is in files dict - - OpenAI params are included in form data - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - language="en", - temperature=0.5, - ) - except Exception: - pass # We just want to capture the request - - # Validate form data contains expected fields - data = captured_request.get("data", {}) - - print("JSON DUMPS captured_request:") - print(json.dumps(captured_request, indent=4, default=str)) - - # Model name should NOT have watsonx/ prefix - assert data.get("model") == "whisper-large-v3-turbo" - - # project_id should be in form data - assert data.get("project_id") == "test-project-123" - - # OpenAI params should be in form data - assert data.get("language") == "en" - assert data.get("temperature") == 0.5 - # response_format should NOT be set by default - only send what user specifies - assert "response_format" not in data - - # Validate file is in files dict (multipart/form-data) - files = captured_request.get("files", {}) - assert "file" in files - assert isinstance( - files["file"], tuple - ) # Should be (filename, content, content_type) - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_project_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "project_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_space_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - space_id="test-space_id-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "space_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - def test_transform_audio_transcription_response_removes_model_field(self): - """ - Test that transform_audio_transcription_response removes the 'model' field - from WatsonX response before creating TranscriptionResponse. - - This test ensures that when WatsonX returns a response with a 'model' field, - it is removed before creating the TranscriptionResponse object, since - TranscriptionResponse doesn't accept a 'model' parameter. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response with 'model' field (as WatsonX may return) - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "model": "whisper-large-v3-turbo", # This field should be removed - "duration": 5.5, - } - mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' - - # This should not raise a TypeError - model field should be removed - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 - - # Verify the model field is NOT in the serialized result - # Check via model_dump() or dict() to ensure it's not in the output - try: - result_dict = result.model_dump() - except AttributeError: - # Fallback for pydantic v1 - result_dict = result.dict() - - # The 'model' field should not be in the result - assert "model" not in result_dict, "Model field should be removed from response" - - def test_transform_audio_transcription_response_without_model_field(self): - """ - Test that transform_audio_transcription_response works correctly - when WatsonX response doesn't include a 'model' field. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response without 'model' field - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "duration": 5.5, - } - mock_response.text = ( - '{"text": "Hello, this is a test transcription.", "duration": 5.5}' - ) - - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py deleted file mode 100644 index 285afffefc0..00000000000 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ /dev/null @@ -1,577 +0,0 @@ -import json - -from typing import Optional -from unittest.mock import Mock, patch - -import pytest - -import litellm -from litellm import completion -from litellm.llms.custom_httpx.http_handler import HTTPHandler - - -@pytest.fixture -def watsonx_chat_completion_call(): - def _call( - model="watsonx/my-test-model", - messages=None, - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if messages is None: - messages = [{"role": "user", "content": "Hello, how are you?"}] - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() # No-op to simulate no exception - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_deployment_model_id_not_in_payload( - monkeypatch, watsonx_chat_completion_call -): - """Test that deployment models do not include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/deployment/test-deployment-id" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data or json_data["model_id"] is None - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data or json_data["project_id"] is None - - -def test_watsonx_regular_model_includes_model_id( - monkeypatch, watsonx_chat_completion_call -): - """Test that regular models include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/regular-model" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -@pytest.fixture -def watsonx_completion_call(): - def _call( - model="watsonx_text/my-test-model", - prompt="Hello, how are you?", - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_completion_deployment_model_id_not_in_payload( - monkeypatch, watsonx_completion_call -): - """Test that deployment models do not include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/deployment/test-deployment-id" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data - - -def test_watsonx_completion_regular_model_includes_model_id( - monkeypatch, watsonx_completion_call -): - """Test that regular models include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/regular-model" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): - """ - Test that gpt-oss-120b model transforms messages to proper format instead of simple concatenation. - - This test calls litellm.completion (sync) and verifies what gets sent in the final POST request body. - Input messages should be transformed using the HuggingFace chat template from openai/gpt-oss-120b, - not just concatenated as "You are chatgpt Hi there". - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - # Test with gpt-oss model using watsonx_text provider (text generation endpoint) - model = "watsonx_text/openai/gpt-oss-120b" - - # Input messages - messages = [ - {"role": "system", "content": "You are chatgpt"}, - {"role": "user", "content": "Hi there"}, - ] - - client = HTTPHandler() - - # Mock HuggingFace template fetch to make test deterministic and avoid network flakiness. - # The test verifies that prompt transformation occurs (not simple concatenation), not the exact - # HuggingFace template format. Using a mock template that produces the correct format is sufficient. - # - # Mock template that produces gpt-oss-120b-like format. - # Note: This is a simplified version of the actual template. The real template is more complex - # (adds metadata, handles tools, thinking messages, etc.), but this captures the key aspects: - # - Converts system role to developer (matching real template behavior) - # - Uses the same tag structure (<|start|>, <|message|>, <|end|>) - # - Preserves message content - mock_tokenizer_config = { - "status": "success", - "tokenizer": { - "chat_template": "{% for message in messages %}{% if message['role'] == 'system' %}<|start|>developer<|message|>{% else %}<|start|>{{ message['role'] }}<|message|>{% endif %}{{ message['content'] }}<|end|>{% endfor %}", - "bos_token": None, - "eos_token": None, - }, - } - - # Isolate known_tokenizer_config so parallel tests don't interfere. - # monkeypatch.setitem restores the original value on teardown. - hf_model = "openai/gpt-oss-120b" - monkeypatch.setitem(litellm.known_tokenizer_config, hf_model, mock_tokenizer_config) - - # Mock IAM token generation to avoid real HTTP calls. - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the POST was called - assert ( - mock_post.call_count == 1 - ), f"POST should have been called exactly once, got {mock_post.call_count}" - - # Get the request body - call_args = mock_post.call_args - assert "data" in call_args.kwargs, "call_args.kwargs should contain 'data'" - json_data = json.loads(call_args.kwargs["data"]) - - # Verify the transformed input is in the request - assert "input" in json_data, "Request should have 'input' field" - transformed_prompt = json_data["input"] - - # Verify it's NOT simple concatenation - simple_concat = "You are chatgpt Hi there" - assert transformed_prompt != simple_concat, ( - f"Prompt should not be simple concatenation.\n" - f"Expected: Chat template with <|start|> tags\n" - f"Got: {transformed_prompt}" - ) - - # Verify it contains proper chat template formatting - assert "<|start|>" in transformed_prompt, "Prompt should contain <|start|> tag" - assert "<|message|>" in transformed_prompt, "Prompt should contain <|message|> tag" - assert "<|end|>" in transformed_prompt, "Prompt should contain <|end|> tag" - assert ( - "You are chatgpt" in transformed_prompt - ), "Prompt should contain system message content" - assert ( - "Hi there" in transformed_prompt - ), "Prompt should contain user message content" - - -@pytest.mark.asyncio -@pytest.mark.xdist_group("watsonx_heavy") -async def test_watsonx_gpt_oss_uses_async_http_handler(): - """ - Test that verifies async HTTP client is used when fetching HuggingFace templates. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( - _aget_chat_template_file, - ) - - # Mock the async HTTP client - mock_async_client = MagicMock() - mock_get = AsyncMock() - mock_async_client.get = mock_get - - # Create mock response for chat template file - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.content = b"test template content" - mock_get.return_value = mock_response - - # Test the async function directly - with patch( - "litellm.litellm_core_utils.prompt_templates.huggingface_template_handler.get_async_httpx_client", - return_value=mock_async_client, - ): - result = await _aget_chat_template_file(hf_model_name="test/model") - - # Verify async HTTP client was called - assert mock_get.called, "Async HTTP client's get method should be called" - assert mock_get.await_count > 0, "Async HTTP client's get should be awaited" - - # Verify it was called with HuggingFace URL - call_args = mock_get.call_args - assert call_args is not None, "get should have been called with arguments" - called_url = call_args.kwargs.get("url", "") - assert ( - "huggingface.co/test/model" in called_url - ), f"Should call HuggingFace API for test/model, got: {called_url}" - assert result["status"] == "success", "Should return success status" - - -@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) -async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( - monkeypatch, tokenizer_config_cached -): - import httpx - - from litellm._uuid import uuid - from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - hf_model = f"openai/gpt-oss-{uuid.uuid4()}" - chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" - if tokenizer_config_cached: - cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} - monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" - else: - monkeypatch.setattr(litellm, "known_tokenizer_config", {}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" - hf_fetched = [] - captured = {} - - def forbid_sync_client(): - raise AssertionError("sync HuggingFace fetch ran on the request path") - - async def serve_hf_file(url, **kwargs): - hf_fetched.append(url) - if url.endswith(".jinja"): - return httpx.Response(200, content=chat_template.encode()) - return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) - - monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) - monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) - - def handle(request): - captured["body"] = json.loads(request.content) - return httpx.Response( - 200, - json={ - "model_id": hf_model, - "results": [ - { - "generated_text": "Hi", - "generated_token_count": 1, - "input_token_count": 1, - "stop_reason": "eos_token", - } - ], - }, - ) - - client = AsyncHTTPHandler() - client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) - - response = await litellm.acompletion( - model=f"watsonx_text/{hf_model}", - messages=[{"role": "user", "content": "Hi there"}], - api_base="https://test-api.watsonx.ai", - project_id="test-project-id", - token="test-token", - client=client, - ) - - assert response.choices[0].message.content == "Hi" - assert hf_fetched == [expected_fetch] - assert captured["body"]["input"] == "<|user|>Hi there" - - -def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch): - """ - Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/openai/gpt-oss-120b" - messages = [{"role": "user", "content": "Test message"}] - - client = HTTPHandler() - - # Mock the token generation call - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - # Call litellm.completion with the new parameter - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - reasoning_effort="low", - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the parameter is in the final request payload - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the JSON data sent in the POST request - request_kwargs = mock_post.call_args.kwargs - json_data = json.loads(request_kwargs["data"]) - - print("\nRequest payload sent to WatsonX API:") - print(json.dumps(json_data, indent=2)) - - # Check for the parameter at the top level of the payload - assert ( - "reasoning_effort" in json_data - ), "'reasoning_effort' should be at the top level of the payload." - assert ( - json_data["reasoning_effort"] == "low" - ), "The value of 'reasoning_effort' should be 'low'." - - -def test_watsonx_zen_api_key_from_client(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key can be passed from client code and is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - zen_api_key = "U1ZDLWQo=" - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - zen_api_key=zen_api_key, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) - - -def test_watsonx_zen_api_key_from_env(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key from environment variable is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - zen_api_key = "U1ZDLWxpdG--===" - monkeypatch.setenv("WATSONX_ZENAPIKEY", zen_api_key) - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) diff --git a/tests/test_litellm/llms/xai/xai_responses/__init__.py b/tests/test_litellm/llms/xai/xai_responses/__init__.py deleted file mode 100644 index 330e9f5a560..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# XAI Responses API tests diff --git a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py b/tests/test_litellm/llms/xai/xai_responses/test_transformation.py deleted file mode 100644 index 3ea3fe631bd..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -Tests for XAI Responses API transformation - -Tests the XAIResponsesAPIConfig class that handles XAI-specific -transformations for the Responses API. - -Source: litellm/llms/xai/responses/transformation.py -""" - - - -import pytest -from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager -from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams - - -class TestXAIResponsesAPITransformation: - """Test XAI Responses API configuration and transformations""" - - def test_xai_provider_config_registration(self): - """Test that XAI provider returns XAIResponsesAPIConfig""" - config = ProviderConfigManager.get_provider_responses_api_config( - model="xai/grok-4-fast", - provider=LlmProviders.XAI, - ) - - assert config is not None, "Config should not be None for XAI provider" - assert isinstance( - config, XAIResponsesAPIConfig - ), f"Expected XAIResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.XAI - ), "custom_llm_provider should be XAI" - - def test_code_interpreter_container_field_removed(self): - """Test that container field is removed from code_interpreter tools""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - tools=[{"type": "code_interpreter", "container": {"type": "auto"}}] - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert "tools" in result - assert len(result["tools"]) == 1 - assert result["tools"][0]["type"] == "code_interpreter" - assert ( - "container" not in result["tools"][0] - ), "Container field should be removed" - - def test_instructions_parameter_forwarded(self): - """xAI supports 'instructions' on /v1/responses, so it must survive param mapping""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - instructions="You are a helpful assistant.", temperature=0.7 - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert result.get("instructions") == "You are a helpful assistant." - assert result.get("temperature") == 0.7, "Other params should be preserved" - - def test_supported_params_includes_instructions(self): - """A system message bridged to 'instructions' must not be rejected for xAI""" - config = XAIResponsesAPIConfig() - supported = config.get_supported_openai_params("grok-4-fast") - - assert "instructions" in supported, "instructions should be supported" - assert "tools" in supported, "tools should be supported" - assert "temperature" in supported, "temperature should be supported" - assert "model" in supported, "model should be supported" - - def test_xai_responses_endpoint_url(self): - """Test that get_complete_url returns correct XAI endpoint""" - config = XAIResponsesAPIConfig() - - # Test with default XAI API base - url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.x.ai/v1/responses" - ), f"Expected XAI responses endpoint, got {url}" - - # Test with custom api_base - custom_url = config.get_complete_url( - api_base="https://custom.x.ai/v1", litellm_params={} - ) - assert ( - custom_url == "https://custom.x.ai/v1/responses" - ), f"Expected custom endpoint, got {custom_url}" - - # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.x.ai/v1/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.x.ai/v1/responses" - ), "Should handle trailing slash" diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py similarity index 98% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py index 4b710175a48..e2fb81bc240 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py @@ -98,6 +98,8 @@ class TestCountTokensLocationResolution: self, counter, monkeypatch ): """Claude models without any location should default to us-east5.""" + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) captured = {} async def fake_ensure_access_token( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py similarity index 60% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index f7df4507651..15df8e47af3 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,21 @@ +import pytest + import litellm +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model_cost_map): assert "reasoning_effort" in litellm.get_supported_openai_params( model="mistral-medium-3", custom_llm_provider="mistral" diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py rename to tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py diff --git a/tests/unit/llms/vertex_ai/videos/__init__.py b/tests/unit/llms/vertex_ai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py rename to tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py similarity index 94% rename from tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py rename to tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py index d42bf7b7a1c..5c8d67ecc70 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -137,30 +137,6 @@ class TestVolcengineResponsesAPITransformation: with pytest.raises(ValueError, match='Volcengine API key is required\\. Set ARK_API_KEY /'): config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) - def test_unsupported_params_are_dropped_with_extra_body(self): - """Unknown fields (including extra_body) should be dropped before send.""" - config = VolcEngineResponsesAPIConfig() - - request = config.transform_responses_api_request( - model="volcengine/demo-model", - input="hi", - response_api_optional_request_params={ - "unsupported_custom_param": 0.1, - "temperature": 0.2, - "metadata": {"k": "v"}, - "extra_body": {"unsupported_custom_param": 1, "temperature": 0.3}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - - assert "unsupported_custom_param" not in request - assert "metadata" not in request - assert request["temperature"] == 0.2 - assert "extra_body" in request - assert "unsupported_custom_param" not in request["extra_body"] - assert request["extra_body"]["temperature"] == 0.3 - def test_valid_thinking_caching_and_expire_at_pass(self): """Documented params should pass through without validation errors.""" config = VolcEngineResponsesAPIConfig() diff --git a/tests/unit/llms/voyage/rerank/__init__.py b/tests/unit/llms/voyage/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py rename to tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/unit/llms/voyage/test_voyage_contextual_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py rename to tests/unit/llms/voyage/test_voyage_contextual_embedding.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/unit/llms/voyage/test_voyage_multimodal_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py rename to tests/unit/llms/voyage/test_voyage_multimodal_embedding.py diff --git a/tests/unit/llms/watsonx/__init__.py b/tests/unit/llms/watsonx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/__init__.py b/tests/unit/llms/watsonx/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py new file mode 100644 index 00000000000..efe592f515e --- /dev/null +++ b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py @@ -0,0 +1,85 @@ +""" +Tests for IBM WatsonX Audio Transcription. + +Validates the WatsonX transcription response transformation. +""" + +from unittest.mock import MagicMock + +from litellm.llms.watsonx.audio_transcription.transformation import ( + IBMWatsonXAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse + + +class TestWatsonXAudioTranscription: + def test_transform_audio_transcription_response_removes_model_field(self): + """ + Test that transform_audio_transcription_response removes the 'model' field + from WatsonX response before creating TranscriptionResponse. + + This test ensures that when WatsonX returns a response with a 'model' field, + it is removed before creating the TranscriptionResponse object, since + TranscriptionResponse doesn't accept a 'model' parameter. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response with 'model' field (as WatsonX may return) + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "model": "whisper-large-v3-turbo", # This field should be removed + "duration": 5.5, + } + mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' + + # This should not raise a TypeError - model field should be removed + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 + + # Verify the model field is NOT in the serialized result + # Check via model_dump() or dict() to ensure it's not in the output + try: + result_dict = result.model_dump() + except AttributeError: + # Fallback for pydantic v1 + result_dict = result.dict() + + # The 'model' field should not be in the result + assert "model" not in result_dict, "Model field should be removed from response" + + def test_transform_audio_transcription_response_without_model_field(self): + """ + Test that transform_audio_transcription_response works correctly + when WatsonX response doesn't include a 'model' field. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response without 'model' field + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "duration": 5.5, + } + mock_response.text = ( + '{"text": "Hello, this is a test transcription.", "duration": 5.5}' + ) + + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py b/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py rename to tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py diff --git a/tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py b/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py rename to tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py diff --git a/tests/unit/llms/watsonx/rerank/__init__.py b/tests/unit/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py similarity index 100% rename from tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py rename to tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py diff --git a/tests/unit/llms/watsonx/test_watsonx.py b/tests/unit/llms/watsonx/test_watsonx.py new file mode 100644 index 00000000000..077539c9acd --- /dev/null +++ b/tests/unit/llms/watsonx/test_watsonx.py @@ -0,0 +1,74 @@ +import json +from unittest.mock import Mock + +import pytest + +import litellm + + +@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) +async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( + monkeypatch, tokenizer_config_cached +): + import httpx + + from litellm._uuid import uuid + from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + hf_model = f"openai/gpt-oss-{uuid.uuid4()}" + chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" + if tokenizer_config_cached: + cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} + monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" + else: + monkeypatch.setattr(litellm, "known_tokenizer_config", {}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" + hf_fetched = [] + captured = {} + + def forbid_sync_client(): + raise AssertionError("sync HuggingFace fetch ran on the request path") + + async def serve_hf_file(url, **kwargs): + hf_fetched.append(url) + if url.endswith(".jinja"): + return httpx.Response(200, content=chat_template.encode()) + return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) + + monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) + monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "model_id": hf_model, + "results": [ + { + "generated_text": "Hi", + "generated_token_count": 1, + "input_token_count": 1, + "stop_reason": "eos_token", + } + ], + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model=f"watsonx_text/{hf_model}", + messages=[{"role": "user", "content": "Hi there"}], + api_base="https://test-api.watsonx.ai", + project_id="test-project-id", + token="test-token", + client=client, + ) + + assert response.choices[0].message.content == "Hi" + assert hf_fetched == [expected_fetch] + assert captured["body"]["input"] == "<|user|>Hi there" diff --git a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py b/tests/unit/llms/watsonx/test_watsonx_common_utils.py similarity index 100% rename from tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py rename to tests/unit/llms/watsonx/test_watsonx_common_utils.py diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/unit/llms/xai/responses/test_xai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py rename to tests/unit/llms/xai/responses/test_xai_responses_transformation.py diff --git a/tests/unit/llms/you_com/__init__.py b/tests/unit/llms/you_com/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/you_com/test_you_com_search.py b/tests/unit/llms/you_com/test_you_com_search.py similarity index 100% rename from tests/test_litellm/llms/you_com/test_you_com_search.py rename to tests/unit/llms/you_com/test_you_com_search.py diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/unit/llms/zai/test_zai_provider.py similarity index 100% rename from tests/test_litellm/llms/zai/test_zai_provider.py rename to tests/unit/llms/zai/test_zai_provider.py From ef6237b9f2bcafba3e7e535a283dd542d67adc70 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:02:53 +0000 Subject: [PATCH 070/127] test: migrate phase 12 legacy llm provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ ...est_openrouter_responses_transformation.py | 0 .../parallel_ai/test_parallel_ai_search.py | 0 .../test_parallel_ai_search_gateway.py | 0 .../llms/parasail/test_parasail.py | 0 .../test_perplexity_chat_transformation.py | 0 ...est_perplexity_embedding_transformation.py | 0 ...est_perplexity_responses_transformation.py | 0 .../test_publicai_chat_transformation.py | 0 .../chat/test_ragflow_chat_transformation.py | 0 .../test_recraft_image_edit_transformation.py | 0 .../test_recraft_image_gen_transformation.py | 0 .../test_text_to_speech_transformation.py | 0 .../test_runway_video_transformation.py | 0 .../test_s3_vectors_transformation.py | 4 ---- .../llms/sap/test_sap_fetch_creds.py | 0 ...eway_audio_transcription_transformation.py | 0 .../test_snowflake_native_endpoints.py | 0 .../test_soniox_provider_registration.py | 0 .../test_stability_image_generation.py | 19 +------------------ .../chat/test_tencent_chat_transformation.py | 0 21 files changed, 15 insertions(+), 22 deletions(-) rename tests/{test_litellm => unit}/llms/openrouter/responses/test_openrouter_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search_gateway.py (100%) rename tests/{test_litellm => unit}/llms/parasail/test_parasail.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/chat/test_perplexity_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/embedding/test_perplexity_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/responses/test_perplexity_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/publicai/test_publicai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/ragflow/chat/test_ragflow_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_edit/test_recraft_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_generation/test_recraft_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/test_text_to_speech_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/videos/test_runway_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py (99%) rename tests/{test_litellm => unit}/llms/sap/test_sap_fetch_creds.py (100%) rename tests/{test_litellm => unit}/llms/scaleway/test_scaleway_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/snowflake/test_snowflake_native_endpoints.py (100%) rename tests/{test_litellm => unit}/llms/soniox/test_soniox_provider_registration.py (100%) rename tests/{test_litellm => unit}/llms/stability/image_generation/test_stability_image_generation.py (93%) rename tests/{test_litellm => unit}/llms/tencent/chat/test_tencent_chat_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 017e63ed1b8..fa91437964f 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,6 +1,7 @@ import os from typing import Final +import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -9,6 +10,19 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) diff --git a/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py rename to tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py diff --git a/tests/test_litellm/llms/parasail/test_parasail.py b/tests/unit/llms/parasail/test_parasail.py similarity index 100% rename from tests/test_litellm/llms/parasail/test_parasail.py rename to tests/unit/llms/parasail/test_parasail.py diff --git a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py rename to tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py rename to tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py diff --git a/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py rename to tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/unit/llms/publicai/test_publicai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py rename to tests/unit/llms/publicai/test_publicai_chat_transformation.py diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py rename to tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py rename to tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py rename to tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py diff --git a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py b/tests/unit/llms/runwayml/test_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py rename to tests/unit/llms/runwayml/test_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py rename to tests/unit/llms/runwayml/videos/test_runway_video_transformation.py diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py similarity index 99% rename from tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py rename to tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 781e92ea7d9..c39887d86ce 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -55,10 +55,6 @@ def _search_kwargs(**overrides): class TestS3VectorsVectorStoreConfig: - def test_init(self): - config = S3VectorsVectorStoreConfig() - assert config is not None - def test_get_supported_openai_params(self): config = S3VectorsVectorStoreConfig() params = config.get_supported_openai_params("test-model") diff --git a/tests/test_litellm/llms/sap/test_sap_fetch_creds.py b/tests/unit/llms/sap/test_sap_fetch_creds.py similarity index 100% rename from tests/test_litellm/llms/sap/test_sap_fetch_creds.py rename to tests/unit/llms/sap/test_sap_fetch_creds.py diff --git a/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py b/tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py rename to tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py similarity index 100% rename from tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py rename to tests/unit/llms/snowflake/test_snowflake_native_endpoints.py diff --git a/tests/test_litellm/llms/soniox/test_soniox_provider_registration.py b/tests/unit/llms/soniox/test_soniox_provider_registration.py similarity index 100% rename from tests/test_litellm/llms/soniox/test_soniox_provider_registration.py rename to tests/unit/llms/soniox/test_soniox_provider_registration.py diff --git a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py similarity index 93% rename from tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py rename to tests/unit/llms/stability/image_generation/test_stability_image_generation.py index c5b3c8fbdc5..c5a78603f9c 100644 --- a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py @@ -10,10 +10,7 @@ from unittest.mock import MagicMock import httpx import pytest -from litellm.llms.stability.image_generation import ( - StabilityImageGenerationConfig, - get_stability_image_generation_config, -) +from litellm.llms.stability.image_generation import StabilityImageGenerationConfig from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, STABILITY_GENERATION_MODELS, @@ -266,20 +263,6 @@ class TestStabilityImageGenerationConfig: assert "filtered" in str(exc_info.value).lower() -class TestFactoryFunction: - """Test the factory function""" - - def test_get_stability_image_generation_config(self): - """Test that factory returns correct config type""" - config = get_stability_image_generation_config("stability/sd3") - assert isinstance(config, StabilityImageGenerationConfig) - - def test_factory_returns_config_for_any_model(self): - """Test that factory works for any model name""" - config = get_stability_image_generation_config("stability/custom-model") - assert isinstance(config, StabilityImageGenerationConfig) - - class TestOpenAISizeMapping: """Test the size to aspect ratio mapping""" diff --git a/tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py rename to tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py From a3dd47ea11b420018b5082e431b1f0ea611a6fad Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:05:03 +0000 Subject: [PATCH 071/127] test(unit): clear ambient Azure credentials in entra token tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ .../azure/realtime/test_azure_realtime_handler.py | 2 +- .../test_azure_ai_image_edit_transformation.py | 4 ++-- .../test_mai_image_edit_transformation.py | 2 +- .../test_azure_ai_passthrough_transformation.py | 4 ++-- .../rerank/test_azure_ai_rerank_transformation.py | 2 +- 6 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 6b6f7c43760..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -12,6 +12,14 @@ import litellm.router as litellm_router_module # noqa: E402 # same import-time import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", +) def _allow_loopback_only() -> None: @@ -53,5 +61,11 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + def pytest_sessionfinish() -> None: enable_socket() diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index c1ba286f8c0..73f43ec8d8a 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -707,7 +707,7 @@ async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypat @pytest.mark.asyncio -async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch): +async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambient_azure_credentials): """ The router binds a deployment's `azure_ad_token` to `_arealtime`'s named parameter rather than **kwargs, so it must still reach the handler. diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 39001c1795b..51ba2c34cd7 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -41,7 +41,7 @@ def test_azure_ai_url_generation(): assert complete_url == expected_url -def test_azure_ai_validate_environment_with_entra_token(monkeypatch): +def test_azure_ai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFluxImageEditConfig() @@ -55,7 +55,7 @@ def test_azure_ai_validate_environment_with_entra_token(monkeypatch): assert headers == {"Authorization": "Bearer entra-token"} -def test_flux2_validate_environment_with_entra_token(monkeypatch): +def test_flux2_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFlux2ImageEditConfig() diff --git a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 75e046825a3..2d6f0083194 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -174,7 +174,7 @@ class TestAzureMAIImageEdit: assert image_response.usage.total_tokens == 1024 -def test_mai_validate_environment_with_entra_token(monkeypatch): +def test_mai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) diff --git a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index f00698a6624..f9fd9681db8 100644 --- a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -256,13 +256,13 @@ def test_serverless_host_gets_a_bearer_token(): assert "api-key" not in headers -def test_entra_token_is_used_when_the_deployment_has_no_api_key(): +def test_entra_token_is_used_when_the_deployment_has_no_api_key(no_ambient_azure_credentials): headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) assert headers["Authorization"] == "Bearer entra-token" -def test_no_credentials_at_all_raises(): +def test_no_credentials_at_all_raises(no_ambient_azure_credentials): with pytest.raises(ValueError, match="Missing Azure AI credentials"): _auth_headers(api_key=None, api_base=FOUNDRY_BASE) diff --git a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index 91bf665f18d..3de27199e2e 100644 --- a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -105,7 +105,7 @@ class TestAzureAIRerankConfigValidateEnvironment: assert headers["Authorization"] == "Bearer my-key" - def test_falls_back_to_entra_token(self, monkeypatch): + def test_falls_back_to_entra_token(self, monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "azure_key", None) From e0b92b2b257bd26650a5887343cfc0de3127129d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:07:38 +0000 Subject: [PATCH 072/127] refactor(types): keep a2a response_dict annotation as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/a2a_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index faf3e98a3a7..834c16ba6dc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, object]] = ( + response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response From 6ba4c2e3399bc16a9996c379c020d1b726f8247a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:40:11 +0000 Subject: [PATCH 073/127] refactor(types): keep guardrail metadata helper accepting dicts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/base_llm/guardrail_translation/base_translation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ace5af8124f..89ad67f0485 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Optional["UserAPIKeyAuth"], + user_api_key_dict: Any | None, ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. From e4a58ef91acdebe2a25c66f7dbdf7d4bdc338fd7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:50:59 +0000 Subject: [PATCH 074/127] test(unit): make every tests/unit directory a package so pytest collection is unique Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/levo/__init__.py | 0 tests/unit/integrations/litellm_agent/__init__.py | 0 tests/unit/integrations/mavvrik_focus/__init__.py | 0 tests/unit/integrations/opik/__init__.py | 0 tests/unit/integrations/pointfive/__init__.py | 0 .../vector_store_integrations/__init__.py | 0 tests/unit/litellm_core_utils/__init__.py | 0 .../unit/litellm_core_utils/audio_utils/__init__.py | 0 .../llm_response_utils/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/a2a/__init__.py | 0 tests/unit/llms/a2a/chat/__init__.py | 0 .../llms/a2a/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/batches/__init__.py | 0 tests/unit/llms/base_llm/__init__.py | 0 tests/unit/llms/base_llm/batches/__init__.py | 0 tests/unit/llms/base_llm/realtime/__init__.py | 0 tests/unit/llms/baseten/__init__.py | 0 tests/unit/llms/baseten/chat/__init__.py | 0 tests/unit/llms/bedrock/__init__.py | 0 tests/unit/llms/bedrock/chat/__init__.py | 0 tests/unit/llms/bedrock/chat/agentcore/__init__.py | 0 .../bedrock/chat/invoke_transformations/__init__.py | 0 tests/unit/llms/bedrock/chat/mantle/__init__.py | 0 tests/unit/llms/bedrock/count_tokens/__init__.py | 0 tests/unit/llms/bedrock/files/__init__.py | 0 tests/unit/llms/bedrock/image/__init__.py | 0 tests/unit/llms/bedrock/image_edit/__init__.py | 0 tests/unit/llms/bedrock/invoke_agent/__init__.py | 0 tests/unit/llms/bedrock/passthrough/__init__.py | 0 .../passthrough/guardrail_translation/__init__.py | 0 tests/unit/llms/bedrock/realtime/__init__.py | 0 tests/unit/llms/bedrock/rerank/__init__.py | 0 tests/unit/llms/bedrock/vector_stores/__init__.py | 0 tests/unit/llms/bedrock_mantle/__init__.py | 0 .../unit/llms/bedrock_mantle/passthrough/__init__.py | 0 tests/unit/llms/black_forest_labs/__init__.py | 0 .../llms/black_forest_labs/image_edit/__init__.py | 0 .../black_forest_labs/image_generation/__init__.py | 0 tests/unit/llms/bytez/__init__.py | 0 tests/unit/llms/bytez/chat/__init__.py | 0 tests/unit/llms/cerebras/__init__.py | 0 tests/unit/llms/chat/__init__.py | 0 tests/unit/llms/chatgpt/__init__.py | 0 tests/unit/llms/chatgpt/chat/__init__.py | 0 tests/unit/llms/chatgpt/responses/__init__.py | 0 tests/unit/llms/cloudflare/__init__.py | 0 tests/unit/llms/cohere/__init__.py | 0 tests/unit/llms/cohere/chat/__init__.py | 0 tests/unit/llms/cohere/embed/__init__.py | 0 tests/unit/llms/cohere/ocr/__init__.py | 0 tests/unit/llms/cohere/rerank/__init__.py | 0 tests/unit/llms/crusoe/__init__.py | 0 tests/unit/llms/databricks/__init__.py | 0 tests/unit/llms/databricks/chat/__init__.py | 0 tests/unit/llms/databricks/responses/__init__.py | 0 tests/unit/llms/datarobot/__init__.py | 0 tests/unit/llms/datarobot/chat/__init__.py | 0 tests/unit/llms/deepseek/__init__.py | 0 tests/unit/llms/deepseek/chat/__init__.py | 0 tests/unit/llms/deepseek/messages/__init__.py | 0 tests/unit/llms/docker_model_runner/__init__.py | 0 tests/unit/llms/elevenlabs/__init__.py | 0 tests/unit/llms/fastcrw/__init__.py | 0 tests/unit/llms/fastcrw/search/__init__.py | 0 tests/unit/llms/fireworks_ai/__init__.py | 0 tests/unit/llms/fireworks_ai/chat/__init__.py | 0 tests/unit/llms/fireworks_ai/rerank/__init__.py | 0 tests/unit/llms/fireworks_ai/responses/__init__.py | 0 tests/unit/llms/gemini/__init__.py | 0 .../unit/llms/gemini/audio_transcription/__init__.py | 0 tests/unit/llms/gemini/files/__init__.py | 0 tests/unit/llms/gemini/google_genai/__init__.py | 0 .../google_genai/guardrail_translation/__init__.py | 0 tests/unit/llms/gemini/image_edit/__init__.py | 0 tests/unit/llms/gemini/realtime/__init__.py | 0 tests/unit/llms/gemini/videos/__init__.py | 0 tests/unit/llms/gigachat/__init__.py | 0 tests/unit/llms/gigachat/chat/__init__.py | 0 tests/unit/llms/gigachat/embedding/__init__.py | 0 tests/unit/llms/gigachat/passthrough/__init__.py | 0 tests/unit/llms/github_copilot/__init__.py | 0 tests/unit/llms/github_copilot/embedding/__init__.py | 0 tests/unit/llms/github_copilot/messages/__init__.py | 0 tests/unit/llms/github_copilot/responses/__init__.py | 0 tests/unit/llms/gradient_ai/__init__.py | 0 tests/unit/llms/gradient_ai/chat/__init__.py | 0 tests/unit/llms/groq/__init__.py | 0 tests/unit/llms/groq/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/__init__.py | 0 tests/unit/llms/hosted_vllm/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/embedding/__init__.py | 0 tests/unit/llms/hosted_vllm/image_edit/__init__.py | 0 tests/unit/llms/hosted_vllm/responses/__init__.py | 0 tests/unit/llms/hosted_vllm/videos/__init__.py | 0 tests/unit/llms/huggingface/__init__.py | 0 tests/unit/llms/huggingface/rerank/__init__.py | 0 tests/unit/llms/inception/__init__.py | 0 tests/unit/llms/jina_ai/__init__.py | 0 tests/unit/llms/jina_ai/embedding/__init__.py | 0 tests/unit/llms/langflow/__init__.py | 0 tests/unit/llms/langflow/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/__init__.py | 0 tests/unit/llms/litellm_proxy/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/skills/__init__.py | 0 tests/unit/llms/llamafile/__init__.py | 0 tests/unit/llms/llamafile/chat/__init__.py | 0 tests/unit/llms/meta/__init__.py | 0 tests/unit/llms/meta/realtime/__init__.py | 0 tests/unit/llms/meta_llama/__init__.py | 0 tests/unit/llms/mistral/audio_speech/__init__.py | 0 tests/unit/llms/modelscope/__init__.py | 0 .../llms/modelscope/image_generation/__init__.py | 0 tests/unit/llms/mongodb/__init__.py | 0 tests/unit/llms/mongodb/vector_stores/__init__.py | 0 tests/unit/llms/moonshot/__init__.py | 0 tests/unit/llms/neosantara/__init__.py | 0 tests/unit/llms/nimble/__init__.py | 0 tests/unit/llms/nimble/search/__init__.py | 0 tests/unit/llms/novita/__init__.py | 0 tests/unit/llms/novita/chat/__init__.py | 0 tests/unit/llms/nscale/__init__.py | 0 tests/unit/llms/nscale/chat/__init__.py | 0 tests/unit/llms/nvidia_nim/__init__.py | 0 tests/unit/llms/nvidia_nim/passthrough/__init__.py | 0 tests/unit/llms/nvidia_nim/rerank/__init__.py | 0 tests/unit/llms/nvidia_riva/__init__.py | 0 .../llms/nvidia_riva/audio_transcription/__init__.py | 0 tests/unit/llms/oci/__init__.py | 0 tests/unit/llms/oci/chat/__init__.py | 0 tests/unit/llms/oci/embed/__init__.py | 0 tests/unit/llms/ocr/__init__.py | 0 .../unit/llms/ocr/guardrail_translation/__init__.py | 0 tests/unit/llms/oobabooga/__init__.py | 0 tests/unit/llms/oobabooga/chat/__init__.py | 0 tests/unit/llms/openai/__init__.py | 0 tests/unit/llms/openai/chat/__init__.py | 0 .../openai/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/completion/__init__.py | 0 tests/unit/test_package_layout.py | 12 ++++++++++++ 143 files changed, 12 insertions(+) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/levo/__init__.py create mode 100644 tests/unit/integrations/litellm_agent/__init__.py create mode 100644 tests/unit/integrations/mavvrik_focus/__init__.py create mode 100644 tests/unit/integrations/opik/__init__.py create mode 100644 tests/unit/integrations/pointfive/__init__.py create mode 100644 tests/unit/integrations/vector_store_integrations/__init__.py create mode 100644 tests/unit/litellm_core_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/audio_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/llm_response_utils/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/a2a/__init__.py create mode 100644 tests/unit/llms/a2a/chat/__init__.py create mode 100644 tests/unit/llms/a2a/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/__init__.py create mode 100644 tests/unit/llms/base_llm/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/realtime/__init__.py create mode 100644 tests/unit/llms/baseten/__init__.py create mode 100644 tests/unit/llms/baseten/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/agentcore/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/mantle/__init__.py create mode 100644 tests/unit/llms/bedrock/count_tokens/__init__.py create mode 100644 tests/unit/llms/bedrock/files/__init__.py create mode 100644 tests/unit/llms/bedrock/image/__init__.py create mode 100644 tests/unit/llms/bedrock/image_edit/__init__.py create mode 100644 tests/unit/llms/bedrock/invoke_agent/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/bedrock/realtime/__init__.py create mode 100644 tests/unit/llms/bedrock/rerank/__init__.py create mode 100644 tests/unit/llms/bedrock/vector_stores/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/passthrough/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_edit/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_generation/__init__.py create mode 100644 tests/unit/llms/bytez/__init__.py create mode 100644 tests/unit/llms/bytez/chat/__init__.py create mode 100644 tests/unit/llms/cerebras/__init__.py create mode 100644 tests/unit/llms/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/__init__.py create mode 100644 tests/unit/llms/chatgpt/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/responses/__init__.py create mode 100644 tests/unit/llms/cloudflare/__init__.py create mode 100644 tests/unit/llms/cohere/__init__.py create mode 100644 tests/unit/llms/cohere/chat/__init__.py create mode 100644 tests/unit/llms/cohere/embed/__init__.py create mode 100644 tests/unit/llms/cohere/ocr/__init__.py create mode 100644 tests/unit/llms/cohere/rerank/__init__.py create mode 100644 tests/unit/llms/crusoe/__init__.py create mode 100644 tests/unit/llms/databricks/__init__.py create mode 100644 tests/unit/llms/databricks/chat/__init__.py create mode 100644 tests/unit/llms/databricks/responses/__init__.py create mode 100644 tests/unit/llms/datarobot/__init__.py create mode 100644 tests/unit/llms/datarobot/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/__init__.py create mode 100644 tests/unit/llms/deepseek/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/messages/__init__.py create mode 100644 tests/unit/llms/docker_model_runner/__init__.py create mode 100644 tests/unit/llms/elevenlabs/__init__.py create mode 100644 tests/unit/llms/fastcrw/__init__.py create mode 100644 tests/unit/llms/fastcrw/search/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/chat/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/rerank/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/responses/__init__.py create mode 100644 tests/unit/llms/gemini/__init__.py create mode 100644 tests/unit/llms/gemini/audio_transcription/__init__.py create mode 100644 tests/unit/llms/gemini/files/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/gemini/image_edit/__init__.py create mode 100644 tests/unit/llms/gemini/realtime/__init__.py create mode 100644 tests/unit/llms/gemini/videos/__init__.py create mode 100644 tests/unit/llms/gigachat/__init__.py create mode 100644 tests/unit/llms/gigachat/chat/__init__.py create mode 100644 tests/unit/llms/gigachat/embedding/__init__.py create mode 100644 tests/unit/llms/gigachat/passthrough/__init__.py create mode 100644 tests/unit/llms/github_copilot/__init__.py create mode 100644 tests/unit/llms/github_copilot/embedding/__init__.py create mode 100644 tests/unit/llms/github_copilot/messages/__init__.py create mode 100644 tests/unit/llms/github_copilot/responses/__init__.py create mode 100644 tests/unit/llms/gradient_ai/__init__.py create mode 100644 tests/unit/llms/gradient_ai/chat/__init__.py create mode 100644 tests/unit/llms/groq/__init__.py create mode 100644 tests/unit/llms/groq/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/embedding/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/image_edit/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/responses/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/videos/__init__.py create mode 100644 tests/unit/llms/huggingface/__init__.py create mode 100644 tests/unit/llms/huggingface/rerank/__init__.py create mode 100644 tests/unit/llms/inception/__init__.py create mode 100644 tests/unit/llms/jina_ai/__init__.py create mode 100644 tests/unit/llms/jina_ai/embedding/__init__.py create mode 100644 tests/unit/llms/langflow/__init__.py create mode 100644 tests/unit/llms/langflow/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/skills/__init__.py create mode 100644 tests/unit/llms/llamafile/__init__.py create mode 100644 tests/unit/llms/llamafile/chat/__init__.py create mode 100644 tests/unit/llms/meta/__init__.py create mode 100644 tests/unit/llms/meta/realtime/__init__.py create mode 100644 tests/unit/llms/meta_llama/__init__.py create mode 100644 tests/unit/llms/mistral/audio_speech/__init__.py create mode 100644 tests/unit/llms/modelscope/__init__.py create mode 100644 tests/unit/llms/modelscope/image_generation/__init__.py create mode 100644 tests/unit/llms/mongodb/__init__.py create mode 100644 tests/unit/llms/mongodb/vector_stores/__init__.py create mode 100644 tests/unit/llms/moonshot/__init__.py create mode 100644 tests/unit/llms/neosantara/__init__.py create mode 100644 tests/unit/llms/nimble/__init__.py create mode 100644 tests/unit/llms/nimble/search/__init__.py create mode 100644 tests/unit/llms/novita/__init__.py create mode 100644 tests/unit/llms/novita/chat/__init__.py create mode 100644 tests/unit/llms/nscale/__init__.py create mode 100644 tests/unit/llms/nscale/chat/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/passthrough/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/rerank/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/audio_transcription/__init__.py create mode 100644 tests/unit/llms/oci/__init__.py create mode 100644 tests/unit/llms/oci/chat/__init__.py create mode 100644 tests/unit/llms/oci/embed/__init__.py create mode 100644 tests/unit/llms/ocr/__init__.py create mode 100644 tests/unit/llms/ocr/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/oobabooga/__init__.py create mode 100644 tests/unit/llms/oobabooga/chat/__init__.py create mode 100644 tests/unit/llms/openai/__init__.py create mode 100644 tests/unit/llms/openai/chat/__init__.py create mode 100644 tests/unit/llms/openai/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/completion/__init__.py create mode 100644 tests/unit/test_package_layout.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/levo/__init__.py b/tests/unit/integrations/levo/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/litellm_agent/__init__.py b/tests/unit/integrations/litellm_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/mavvrik_focus/__init__.py b/tests/unit/integrations/mavvrik_focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/opik/__init__.py b/tests/unit/integrations/opik/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/pointfive/__init__.py b/tests/unit/integrations/pointfive/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/vector_store_integrations/__init__.py b/tests/unit/integrations/vector_store_integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/__init__.py b/tests/unit/litellm_core_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/audio_utils/__init__.py b/tests/unit/litellm_core_utils/audio_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/llm_response_utils/__init__.py b/tests/unit/litellm_core_utils/llm_response_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/__init__.py b/tests/unit/llms/a2a/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/__init__.py b/tests/unit/llms/a2a/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py b/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/batches/__init__.py b/tests/unit/llms/anthropic/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/__init__.py b/tests/unit/llms/base_llm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/batches/__init__.py b/tests/unit/llms/base_llm/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/realtime/__init__.py b/tests/unit/llms/base_llm/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/__init__.py b/tests/unit/llms/baseten/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/chat/__init__.py b/tests/unit/llms/baseten/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/__init__.py b/tests/unit/llms/bedrock/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/__init__.py b/tests/unit/llms/bedrock/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/agentcore/__init__.py b/tests/unit/llms/bedrock/chat/agentcore/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py b/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/mantle/__init__.py b/tests/unit/llms/bedrock/chat/mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/count_tokens/__init__.py b/tests/unit/llms/bedrock/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/files/__init__.py b/tests/unit/llms/bedrock/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image/__init__.py b/tests/unit/llms/bedrock/image/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image_edit/__init__.py b/tests/unit/llms/bedrock/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/invoke_agent/__init__.py b/tests/unit/llms/bedrock/invoke_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/__init__.py b/tests/unit/llms/bedrock/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/realtime/__init__.py b/tests/unit/llms/bedrock/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/rerank/__init__.py b/tests/unit/llms/bedrock/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/vector_stores/__init__.py b/tests/unit/llms/bedrock/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/__init__.py b/tests/unit/llms/bedrock_mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/passthrough/__init__.py b/tests/unit/llms/bedrock_mantle/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/__init__.py b/tests/unit/llms/black_forest_labs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_edit/__init__.py b/tests/unit/llms/black_forest_labs/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_generation/__init__.py b/tests/unit/llms/black_forest_labs/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/__init__.py b/tests/unit/llms/bytez/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/chat/__init__.py b/tests/unit/llms/bytez/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cerebras/__init__.py b/tests/unit/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chat/__init__.py b/tests/unit/llms/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/__init__.py b/tests/unit/llms/chatgpt/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/chat/__init__.py b/tests/unit/llms/chatgpt/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/responses/__init__.py b/tests/unit/llms/chatgpt/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cloudflare/__init__.py b/tests/unit/llms/cloudflare/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/__init__.py b/tests/unit/llms/cohere/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/chat/__init__.py b/tests/unit/llms/cohere/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/embed/__init__.py b/tests/unit/llms/cohere/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/ocr/__init__.py b/tests/unit/llms/cohere/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/rerank/__init__.py b/tests/unit/llms/cohere/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/crusoe/__init__.py b/tests/unit/llms/crusoe/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/__init__.py b/tests/unit/llms/databricks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/chat/__init__.py b/tests/unit/llms/databricks/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/responses/__init__.py b/tests/unit/llms/databricks/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/__init__.py b/tests/unit/llms/datarobot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/chat/__init__.py b/tests/unit/llms/datarobot/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/__init__.py b/tests/unit/llms/deepseek/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/chat/__init__.py b/tests/unit/llms/deepseek/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/messages/__init__.py b/tests/unit/llms/deepseek/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/docker_model_runner/__init__.py b/tests/unit/llms/docker_model_runner/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/elevenlabs/__init__.py b/tests/unit/llms/elevenlabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/__init__.py b/tests/unit/llms/fastcrw/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/search/__init__.py b/tests/unit/llms/fastcrw/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/__init__.py b/tests/unit/llms/fireworks_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/chat/__init__.py b/tests/unit/llms/fireworks_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/rerank/__init__.py b/tests/unit/llms/fireworks_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/responses/__init__.py b/tests/unit/llms/fireworks_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/__init__.py b/tests/unit/llms/gemini/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/audio_transcription/__init__.py b/tests/unit/llms/gemini/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/files/__init__.py b/tests/unit/llms/gemini/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/__init__.py b/tests/unit/llms/gemini/google_genai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/image_edit/__init__.py b/tests/unit/llms/gemini/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/realtime/__init__.py b/tests/unit/llms/gemini/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/videos/__init__.py b/tests/unit/llms/gemini/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/__init__.py b/tests/unit/llms/gigachat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/chat/__init__.py b/tests/unit/llms/gigachat/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/embedding/__init__.py b/tests/unit/llms/gigachat/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/passthrough/__init__.py b/tests/unit/llms/gigachat/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/__init__.py b/tests/unit/llms/github_copilot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/embedding/__init__.py b/tests/unit/llms/github_copilot/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/messages/__init__.py b/tests/unit/llms/github_copilot/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/responses/__init__.py b/tests/unit/llms/github_copilot/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/__init__.py b/tests/unit/llms/gradient_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/chat/__init__.py b/tests/unit/llms/gradient_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/__init__.py b/tests/unit/llms/groq/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/chat/__init__.py b/tests/unit/llms/groq/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/__init__.py b/tests/unit/llms/hosted_vllm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/chat/__init__.py b/tests/unit/llms/hosted_vllm/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/embedding/__init__.py b/tests/unit/llms/hosted_vllm/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/image_edit/__init__.py b/tests/unit/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/responses/__init__.py b/tests/unit/llms/hosted_vllm/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/videos/__init__.py b/tests/unit/llms/hosted_vllm/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/__init__.py b/tests/unit/llms/huggingface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/rerank/__init__.py b/tests/unit/llms/huggingface/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/inception/__init__.py b/tests/unit/llms/inception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/__init__.py b/tests/unit/llms/jina_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/embedding/__init__.py b/tests/unit/llms/jina_ai/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/__init__.py b/tests/unit/llms/langflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/chat/__init__.py b/tests/unit/llms/langflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/__init__.py b/tests/unit/llms/litellm_proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/chat/__init__.py b/tests/unit/llms/litellm_proxy/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/skills/__init__.py b/tests/unit/llms/litellm_proxy/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/__init__.py b/tests/unit/llms/llamafile/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/chat/__init__.py b/tests/unit/llms/llamafile/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/__init__.py b/tests/unit/llms/meta/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/realtime/__init__.py b/tests/unit/llms/meta/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta_llama/__init__.py b/tests/unit/llms/meta_llama/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mistral/audio_speech/__init__.py b/tests/unit/llms/mistral/audio_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/__init__.py b/tests/unit/llms/modelscope/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/image_generation/__init__.py b/tests/unit/llms/modelscope/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/__init__.py b/tests/unit/llms/mongodb/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/vector_stores/__init__.py b/tests/unit/llms/mongodb/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/moonshot/__init__.py b/tests/unit/llms/moonshot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/neosantara/__init__.py b/tests/unit/llms/neosantara/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/__init__.py b/tests/unit/llms/nimble/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/search/__init__.py b/tests/unit/llms/nimble/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/__init__.py b/tests/unit/llms/novita/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/chat/__init__.py b/tests/unit/llms/novita/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/__init__.py b/tests/unit/llms/nscale/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/chat/__init__.py b/tests/unit/llms/nscale/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/__init__.py b/tests/unit/llms/nvidia_nim/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/passthrough/__init__.py b/tests/unit/llms/nvidia_nim/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/rerank/__init__.py b/tests/unit/llms/nvidia_nim/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/__init__.py b/tests/unit/llms/nvidia_riva/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py b/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/__init__.py b/tests/unit/llms/oci/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/chat/__init__.py b/tests/unit/llms/oci/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/embed/__init__.py b/tests/unit/llms/oci/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/__init__.py b/tests/unit/llms/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/guardrail_translation/__init__.py b/tests/unit/llms/ocr/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/__init__.py b/tests/unit/llms/oobabooga/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/chat/__init__.py b/tests/unit/llms/oobabooga/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/__init__.py b/tests/unit/llms/openai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/__init__.py b/tests/unit/llms/openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/guardrail_translation/__init__.py b/tests/unit/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/completion/__init__.py b/tests/unit/llms/openai/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_package_layout.py b/tests/unit/test_package_layout.py new file mode 100644 index 00000000000..4ea68fc06ba --- /dev/null +++ b/tests/unit/test_package_layout.py @@ -0,0 +1,12 @@ +import os + +TESTS_UNIT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def test_every_directory_under_tests_unit_is_a_package(): + missing = [] + for root, dirs, _files in os.walk(TESTS_UNIT_DIR): + dirs[:] = [d for d in dirs if d != "__pycache__"] + if not os.path.isfile(os.path.join(root, "__init__.py")): + missing.append(os.path.relpath(root, TESTS_UNIT_DIR)) + assert missing == [] From 5599c59923c5362dccc5c17012ba0c02a2d7254f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:07 +0000 Subject: [PATCH 075/127] test: add package initializers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openai/embeddings/__init__.py | 0 .../unit/llms/openai/embeddings/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/evals/__init__.py | 0 tests/unit/llms/openai/image_generation/__init__.py | 0 tests/unit/llms/openai/speech/__init__.py | 0 tests/unit/llms/openai/transcriptions/__init__.py | 0 tests/unit/llms/openai/vector_store_files/__init__.py | 0 tests/unit/llms/openai/vector_stores/__init__.py | 0 tests/unit/llms/openai/videos/__init__.py | 0 tests/unit/llms/openai_like/__init__.py | 0 tests/unit/llms/openai_like/chat/__init__.py | 0 tests/unit/llms/openai_like/embedding/__init__.py | 0 tests/unit/llms/openai_like/messages/__init__.py | 0 tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/openrouter/chat/__init__.py | 0 tests/unit/llms/openrouter/image_edit/__init__.py | 0 tests/unit/llms/openrouter/image_generation/__init__.py | 0 17 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openai/embeddings/__init__.py create mode 100644 tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/evals/__init__.py create mode 100644 tests/unit/llms/openai/image_generation/__init__.py create mode 100644 tests/unit/llms/openai/speech/__init__.py create mode 100644 tests/unit/llms/openai/transcriptions/__init__.py create mode 100644 tests/unit/llms/openai/vector_store_files/__init__.py create mode 100644 tests/unit/llms/openai/vector_stores/__init__.py create mode 100644 tests/unit/llms/openai/videos/__init__.py create mode 100644 tests/unit/llms/openai_like/__init__.py create mode 100644 tests/unit/llms/openai_like/chat/__init__.py create mode 100644 tests/unit/llms/openai_like/embedding/__init__.py create mode 100644 tests/unit/llms/openai_like/messages/__init__.py create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/openrouter/chat/__init__.py create mode 100644 tests/unit/llms/openrouter/image_edit/__init__.py create mode 100644 tests/unit/llms/openrouter/image_generation/__init__.py diff --git a/tests/unit/llms/openai/embeddings/__init__.py b/tests/unit/llms/openai/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py b/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/evals/__init__.py b/tests/unit/llms/openai/evals/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/image_generation/__init__.py b/tests/unit/llms/openai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/speech/__init__.py b/tests/unit/llms/openai/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/transcriptions/__init__.py b/tests/unit/llms/openai/transcriptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_store_files/__init__.py b/tests/unit/llms/openai/vector_store_files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_stores/__init__.py b/tests/unit/llms/openai/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/videos/__init__.py b/tests/unit/llms/openai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/__init__.py b/tests/unit/llms/openai_like/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/chat/__init__.py b/tests/unit/llms/openai_like/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/embedding/__init__.py b/tests/unit/llms/openai_like/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/messages/__init__.py b/tests/unit/llms/openai_like/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/chat/__init__.py b/tests/unit/llms/openrouter/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_edit/__init__.py b/tests/unit/llms/openrouter/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_generation/__init__.py b/tests/unit/llms/openrouter/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 49cd32affe1cfd597c223cdb62049d52ed55e5b5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:09 +0000 Subject: [PATCH 076/127] test: add __init__.py to every tests/unit directory this migration touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/vertex_ai/__init__.py | 0 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py | 0 tests/unit/llms/volcengine/__init__.py | 0 tests/unit/llms/volcengine/responses/__init__.py | 0 tests/unit/llms/voyage/__init__.py | 0 tests/unit/llms/watsonx/embed/__init__.py | 0 tests/unit/llms/watsonx/passthrough/__init__.py | 0 tests/unit/llms/xai/__init__.py | 0 tests/unit/llms/xai/responses/__init__.py | 0 tests/unit/llms/zai/__init__.py | 0 12 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py create mode 100644 tests/unit/llms/volcengine/__init__.py create mode 100644 tests/unit/llms/volcengine/responses/__init__.py create mode 100644 tests/unit/llms/voyage/__init__.py create mode 100644 tests/unit/llms/watsonx/embed/__init__.py create mode 100644 tests/unit/llms/watsonx/passthrough/__init__.py create mode 100644 tests/unit/llms/xai/__init__.py create mode 100644 tests/unit/llms/xai/responses/__init__.py create mode 100644 tests/unit/llms/zai/__init__.py diff --git a/tests/unit/llms/vertex_ai/__init__.py b/tests/unit/llms/vertex_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/__init__.py b/tests/unit/llms/volcengine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/responses/__init__.py b/tests/unit/llms/volcengine/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/voyage/__init__.py b/tests/unit/llms/voyage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/embed/__init__.py b/tests/unit/llms/watsonx/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/passthrough/__init__.py b/tests/unit/llms/watsonx/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/__init__.py b/tests/unit/llms/xai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/responses/__init__.py b/tests/unit/llms/xai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/zai/__init__.py b/tests/unit/llms/zai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 434af084e7572d2b2f23101f56d7c722189aa33d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:38 +0000 Subject: [PATCH 077/127] test: add __init__.py to every tests/unit directory phase 16 touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/sandbox/__init__.py | 0 tests/unit/skills/__init__.py | 0 tests/unit/test_router/__init__.py | 0 tests/unit/types/__init__.py | 0 tests/unit/types/llms/__init__.py | 0 tests/unit/types/proxy/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/sandbox/__init__.py create mode 100644 tests/unit/skills/__init__.py create mode 100644 tests/unit/test_router/__init__.py create mode 100644 tests/unit/types/__init__.py create mode 100644 tests/unit/types/llms/__init__.py create mode 100644 tests/unit/types/proxy/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/sandbox/__init__.py b/tests/unit/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/skills/__init__.py b/tests/unit/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_router/__init__.py b/tests/unit/test_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/__init__.py b/tests/unit/types/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/llms/__init__.py b/tests/unit/types/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/__init__.py b/tests/unit/types/proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From baf40ea5e896674dd8e75d6ac8b1937fae07e475 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:14 +0000 Subject: [PATCH 078/127] test(unit): add package markers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/experimental_pass_through/__init__.py | 0 tests/unit/llms/anthropic/files/__init__.py | 0 tests/unit/llms/anthropic/messages/__init__.py | 0 tests/unit/llms/apiserpent/__init__.py | 0 tests/unit/llms/azure/__init__.py | 0 tests/unit/llms/azure/image_edit/__init__.py | 0 tests/unit/llms/azure/image_generation/__init__.py | 0 tests/unit/llms/azure/passthrough/__init__.py | 0 tests/unit/llms/azure/realtime/__init__.py | 0 tests/unit/llms/azure/response/__init__.py | 0 tests/unit/llms/azure/search/__init__.py | 0 tests/unit/llms/azure/text_to_speech/__init__.py | 0 tests/unit/llms/azure/vector_stores/__init__.py | 0 tests/unit/llms/azure_ai/__init__.py | 0 tests/unit/llms/azure_ai/chat/__init__.py | 0 tests/unit/llms/azure_ai/embed/__init__.py | 0 tests/unit/llms/azure_ai/image_edit/__init__.py | 0 tests/unit/llms/azure_ai/ocr/__init__.py | 0 tests/unit/llms/azure_ai/passthrough/__init__.py | 0 tests/unit/llms/azure_ai/rerank/__init__.py | 0 tests/unit/llms/azure_ai/responses/__init__.py | 0 24 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/experimental_pass_through/__init__.py create mode 100644 tests/unit/llms/anthropic/files/__init__.py create mode 100644 tests/unit/llms/anthropic/messages/__init__.py create mode 100644 tests/unit/llms/apiserpent/__init__.py create mode 100644 tests/unit/llms/azure/__init__.py create mode 100644 tests/unit/llms/azure/image_edit/__init__.py create mode 100644 tests/unit/llms/azure/image_generation/__init__.py create mode 100644 tests/unit/llms/azure/passthrough/__init__.py create mode 100644 tests/unit/llms/azure/realtime/__init__.py create mode 100644 tests/unit/llms/azure/response/__init__.py create mode 100644 tests/unit/llms/azure/search/__init__.py create mode 100644 tests/unit/llms/azure/text_to_speech/__init__.py create mode 100644 tests/unit/llms/azure/vector_stores/__init__.py create mode 100644 tests/unit/llms/azure_ai/__init__.py create mode 100644 tests/unit/llms/azure_ai/chat/__init__.py create mode 100644 tests/unit/llms/azure_ai/embed/__init__.py create mode 100644 tests/unit/llms/azure_ai/image_edit/__init__.py create mode 100644 tests/unit/llms/azure_ai/ocr/__init__.py create mode 100644 tests/unit/llms/azure_ai/passthrough/__init__.py create mode 100644 tests/unit/llms/azure_ai/rerank/__init__.py create mode 100644 tests/unit/llms/azure_ai/responses/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/experimental_pass_through/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/files/__init__.py b/tests/unit/llms/anthropic/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/messages/__init__.py b/tests/unit/llms/anthropic/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/apiserpent/__init__.py b/tests/unit/llms/apiserpent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/__init__.py b/tests/unit/llms/azure/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_edit/__init__.py b/tests/unit/llms/azure/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_generation/__init__.py b/tests/unit/llms/azure/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/passthrough/__init__.py b/tests/unit/llms/azure/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/realtime/__init__.py b/tests/unit/llms/azure/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/response/__init__.py b/tests/unit/llms/azure/response/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/search/__init__.py b/tests/unit/llms/azure/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/text_to_speech/__init__.py b/tests/unit/llms/azure/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/vector_stores/__init__.py b/tests/unit/llms/azure/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/__init__.py b/tests/unit/llms/azure_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/chat/__init__.py b/tests/unit/llms/azure_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/embed/__init__.py b/tests/unit/llms/azure_ai/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/image_edit/__init__.py b/tests/unit/llms/azure_ai/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/ocr/__init__.py b/tests/unit/llms/azure_ai/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/passthrough/__init__.py b/tests/unit/llms/azure_ai/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/rerank/__init__.py b/tests/unit/llms/azure_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/responses/__init__.py b/tests/unit/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30a422087c9bf28cc3fc94a6d68b241b703325a1 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:49 +0000 Subject: [PATCH 079/127] test: add __init__.py to migrated tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/a2a_protocol/__init__.py | 0 tests/unit/a2a_protocol/providers/__init__.py | 0 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py | 0 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py | 0 tests/unit/anthropic_interface/__init__.py | 0 tests/unit/anthropic_interface/exceptions/__init__.py | 0 tests/unit/batches/__init__.py | 0 tests/unit/chat_completions/__init__.py | 0 tests/unit/completion_extras/__init__.py | 0 tests/unit/compression/__init__.py | 0 tests/unit/endpoints/__init__.py | 0 tests/unit/endpoints/speech/__init__.py | 0 .../unit/endpoints/speech/speech_to_completion_bridge/__init__.py | 0 tests/unit/enterprise/__init__.py | 0 tests/unit/enterprise/enterprise_callbacks/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/compression_interception/__init__.py | 0 tests/unit/integrations/gcs_bucket/__init__.py | 0 tests/unit/integrations/gcs_pubsub/__init__.py | 0 tests/unit/integrations/helicone/__init__.py | 0 20 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/a2a_protocol/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py create mode 100644 tests/unit/anthropic_interface/__init__.py create mode 100644 tests/unit/anthropic_interface/exceptions/__init__.py create mode 100644 tests/unit/batches/__init__.py create mode 100644 tests/unit/chat_completions/__init__.py create mode 100644 tests/unit/completion_extras/__init__.py create mode 100644 tests/unit/compression/__init__.py create mode 100644 tests/unit/endpoints/__init__.py create mode 100644 tests/unit/endpoints/speech/__init__.py create mode 100644 tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py create mode 100644 tests/unit/enterprise/__init__.py create mode 100644 tests/unit/enterprise/enterprise_callbacks/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/compression_interception/__init__.py create mode 100644 tests/unit/integrations/gcs_bucket/__init__.py create mode 100644 tests/unit/integrations/gcs_pubsub/__init__.py create mode 100644 tests/unit/integrations/helicone/__init__.py diff --git a/tests/unit/a2a_protocol/__init__.py b/tests/unit/a2a_protocol/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/__init__.py b/tests/unit/a2a_protocol/providers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/__init__.py b/tests/unit/anthropic_interface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/exceptions/__init__.py b/tests/unit/anthropic_interface/exceptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/batches/__init__.py b/tests/unit/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/chat_completions/__init__.py b/tests/unit/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/completion_extras/__init__.py b/tests/unit/completion_extras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/compression/__init__.py b/tests/unit/compression/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/__init__.py b/tests/unit/endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/__init__.py b/tests/unit/endpoints/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/__init__.py b/tests/unit/enterprise/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/enterprise_callbacks/__init__.py b/tests/unit/enterprise/enterprise_callbacks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/compression_interception/__init__.py b/tests/unit/integrations/compression_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_bucket/__init__.py b/tests/unit/integrations/gcs_bucket/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_pubsub/__init__.py b/tests/unit/integrations/gcs_pubsub/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/helicone/__init__.py b/tests/unit/integrations/helicone/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 78a751c04972d0079ef01bb87fef58a6e472be65 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:55:43 +0000 Subject: [PATCH 080/127] test: migrate phase 15 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/messages/__init__.py | 0 .../messages/test_dispatch.py | 11 +- .../models/test_models.py | 46 ++--- tests/unit/ocr/__init__.py | 0 .../ocr/test_dispatch.py | 0 tests/{test_litellm => unit}/ocr/test_main.py | 0 .../ocr/test_ocr_file_input.py | 48 ++--- tests/unit/passthrough/__init__.py | 0 .../test_async_streaming_error_propagation.py | 27 ++- .../passthrough/test_passthrough_main.py | 77 ++------ ...test_streaming_interrupt_spend_tracking.py | 31 +--- tests/unit/rag/ingestion/__init__.py | 0 .../ingestion/test_s3_vectors_ingestion.py | 12 +- .../realtime_api/test_main.py | 11 +- .../repositories/test_repositories.py | 48 ++--- .../repositories/test_unit_of_work.py | 0 .../complexity_router/test_jev_classifier.py | 0 .../test_deployment_affinity_check.py | 70 ++++---- .../test_encrypted_content_affinity_check.py | 169 ++++++++++-------- .../test_prompt_caching_deployment_check.py | 47 +++-- .../test_responses_api_deployment_check.py | 26 +-- .../test_session_id_affinity.py | 70 ++++---- tests/unit/rust_bridge/__init__.py | 0 .../rust_bridge/chat_completions/__init__.py | 0 .../chat_completions/test_route_host.py | 0 tests/unit/rust_bridge/messages/__init__.py | 0 .../rust_bridge/messages/test_route_host.py | 0 27 files changed, 290 insertions(+), 403 deletions(-) create mode 100644 tests/unit/messages/__init__.py rename tests/{test_litellm => unit}/messages/test_dispatch.py (97%) rename tests/{test_litellm => unit}/models/test_models.py (93%) create mode 100644 tests/unit/ocr/__init__.py rename tests/{test_litellm => unit}/ocr/test_dispatch.py (100%) rename tests/{test_litellm => unit}/ocr/test_main.py (100%) rename tests/{test_litellm => unit}/ocr/test_ocr_file_input.py (92%) create mode 100644 tests/unit/passthrough/__init__.py rename tests/{test_litellm => unit}/passthrough/test_async_streaming_error_propagation.py (92%) rename tests/{test_litellm => unit}/passthrough/test_passthrough_main.py (94%) rename tests/{test_litellm => unit}/passthrough/test_streaming_interrupt_spend_tracking.py (91%) create mode 100644 tests/unit/rag/ingestion/__init__.py rename tests/{test_litellm => unit}/rag/ingestion/test_s3_vectors_ingestion.py (94%) rename tests/{test_litellm => unit}/realtime_api/test_main.py (98%) rename tests/{test_litellm => unit}/repositories/test_repositories.py (98%) rename tests/{test_litellm => unit}/repositories/test_unit_of_work.py (100%) rename tests/{test_litellm => unit}/router_strategy/complexity_router/test_jev_classifier.py (100%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_deployment_affinity_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py (96%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py (94%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_responses_api_deployment_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_session_id_affinity.py (95%) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/chat_completions/__init__.py rename tests/{test_litellm => unit}/rust_bridge/chat_completions/test_route_host.py (100%) create mode 100644 tests/unit/rust_bridge/messages/__init__.py rename tests/{test_litellm => unit}/rust_bridge/messages/test_route_host.py (100%) diff --git a/tests/unit/messages/__init__.py b/tests/unit/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py similarity index 97% rename from tests/test_litellm/messages/test_dispatch.py rename to tests/unit/messages/test_dispatch.py index 2eaf4cd9a50..586b77d9a25 100644 --- a/tests/test_litellm/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -29,9 +29,7 @@ RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),) def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]: - binding: Final[NativeBinding[NativeMessages]] = NativeBinding( - "anthropic_messages_handler", validate=lambda _: None - ) + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("anthropic_messages_handler", validate=lambda _: None) binding.override(native) return binding @@ -99,7 +97,8 @@ async def test_async_python_route_forwards_original_call_shape() -> None: expected: Final = response() async def python( - *call_args: object, **call_kwargs: object # kwargs-ok: records call shape + *call_args: object, + **call_kwargs: object, # kwargs-ok: records call shape ) -> AnthropicMessagesResponse: captured.append((call_args, call_kwargs)) return expected @@ -217,7 +216,9 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] expected: Final = response() - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call + def python( + *call_args: object, **call_kwargs: object + ) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call captured.append((call_args, call_kwargs)) return expected diff --git a/tests/test_litellm/models/test_models.py b/tests/unit/models/test_models.py similarity index 93% rename from tests/test_litellm/models/test_models.py rename to tests/unit/models/test_models.py index 777b4a265ac..b8bf55f1b4a 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -5,7 +5,7 @@ Tests for backend domain models. from datetime import datetime, timezone import pytest -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.autorouter_session import LiteLLM_AutoRouterSession @@ -19,7 +19,6 @@ from litellm.models.credentials import CreateCredentialItem, CredentialItem from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.managed_files import ( LiteLLM_ManagedFileTable, - LiteLLM_ManagedObjectTable, LiteLLM_ManagedVectorStoresTable, ) from litellm.models.mcp_server import LiteLLM_MCPServerTable @@ -41,7 +40,6 @@ from litellm.models.verification_token import ( LiteLLM_DeletedVerificationToken, LiteLLM_VerificationToken, ) -from pydantic import ValidationError class TestBudget: @@ -121,9 +119,7 @@ class TestCredentials: assert item.credential_values is None def test_create_credential_item_requires_values_or_model_id(self): - with pytest.raises( - ValueError, match="Either credential_values or model_id must be set" - ): + with pytest.raises(ValueError, match="Either credential_values or model_id must be set"): CreateCredentialItem(credential_name="bad", credential_info={}) @@ -141,12 +137,8 @@ class TestModel: assert model.team_public_model_name == "my-gpt4" def test_is_blocked(self): - model_blocked = LiteLLM_ProxyModelTable( - model_id="m1", model_name="test", litellm_params={}, blocked=True - ) - model_unblocked = LiteLLM_ProxyModelTable( - model_id="m2", model_name="test", litellm_params={}, blocked=False - ) + model_blocked = LiteLLM_ProxyModelTable(model_id="m1", model_name="test", litellm_params={}, blocked=True) + model_unblocked = LiteLLM_ProxyModelTable(model_id="m2", model_name="test", litellm_params={}, blocked=False) assert model_blocked.is_blocked assert not model_unblocked.is_blocked @@ -188,9 +180,7 @@ class TestModel: assert model.blocked is True def test_team_helpers_none_when_no_model_info(self): - model = LiteLLM_ProxyModelTable( - model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None - ) + model = LiteLLM_ProxyModelTable(model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None) assert model.team_id is None assert model.team_public_model_name is None @@ -292,9 +282,7 @@ class TestTeam: assert team.model_max_budget == {"gpt-4": 5.0} def test_cached_team(self): - cached = LiteLLM_TeamTableCachedObj( - team_id="t1", last_refreshed_at=1234567890.0 - ) + cached = LiteLLM_TeamTableCachedObj(team_id="t1", last_refreshed_at=1234567890.0) assert cached.last_refreshed_at == 1234567890.0 def test_deleted_team(self): @@ -345,9 +333,7 @@ class TestUser: assert "password" not in user.model_dump() assert "password" not in user.model_dump_json() - with_keys = LiteLLM_UserTableWithKeyCount( - user_id="u1", user_email="a@b.c", password=secret, key_count=2 - ) + with_keys = LiteLLM_UserTableWithKeyCount(user_id="u1", user_email="a@b.c", password=secret, key_count=2) assert with_keys.password == secret assert "password" not in with_keys.model_dump() assert "password" not in with_keys.model_dump_json() @@ -479,9 +465,7 @@ class TestEndUserTable: class TestBudgetTableFull: def test_full_adds_server_managed_fields(self): now = datetime.now() - budget = LiteLLM_BudgetTableFull( - budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now - ) + budget = LiteLLM_BudgetTableFull(budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now) assert budget.created_at == now assert budget.budget_reset_at == now assert budget.max_budget == 10.0 @@ -493,9 +477,7 @@ class TestBudgetTableFull: class TestTeamMemberTable: def test_tracks_user_within_team(self): - member = LiteLLM_TeamMemberTable( - user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 - ) + member = LiteLLM_TeamMemberTable(user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0) assert member.user_id == "u1" assert member.team_id == "t1" assert member.spend == 3.0 @@ -585,9 +567,7 @@ class TestSpendLogs: assert log.updated_at == updated_at def test_error_logs_creation(self): - log = LiteLLM_ErrorLogs( - request_id="r1", startTime=None, endTime=None, status_code="500" - ) + log = LiteLLM_ErrorLogs(request_id="r1", startTime=None, endTime=None, status_code="500") assert log.request_id == "r1" assert log.status_code == "500" @@ -603,12 +583,6 @@ class TestManagedTables: assert table.model_mappings == {"gpt-4": "file-abc"} assert table.flat_model_file_ids == ["file-abc"] - def test_managed_object_table_requires_purpose(self): - with pytest.raises(ValidationError): - LiteLLM_ManagedObjectTable( - unified_object_id="o1", model_object_id="m1", file_object={} - ) - def test_managed_vector_stores_table(self): table = LiteLLM_ManagedVectorStoresTable( vector_store_id="vs1", diff --git a/tests/unit/ocr/__init__.py b/tests/unit/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/ocr/test_dispatch.py b/tests/unit/ocr/test_dispatch.py similarity index 100% rename from tests/test_litellm/ocr/test_dispatch.py rename to tests/unit/ocr/test_dispatch.py diff --git a/tests/test_litellm/ocr/test_main.py b/tests/unit/ocr/test_main.py similarity index 100% rename from tests/test_litellm/ocr/test_main.py rename to tests/unit/ocr/test_main.py diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/unit/ocr/test_ocr_file_input.py similarity index 92% rename from tests/test_litellm/ocr/test_ocr_file_input.py rename to tests/unit/ocr/test_ocr_file_input.py index 4ac27d286e1..d67f5280195 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/unit/ocr/test_ocr_file_input.py @@ -73,9 +73,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -95,9 +93,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -112,9 +108,7 @@ class TestConvertFileDocumentToUrlDocument: request handler the value is attacker-controlled, and opening it as a path is an arbitrary local file read on the proxy host.""" with pytest.raises(ValueError, match="does not accept bare str values"): - convert_file_document_to_url_document( - {"type": "file", "file": "/etc/passwd"} - ) + convert_file_document_to_url_document({"type": "file", "file": "/etc/passwd"}) def test_should_convert_pathlib_path(self): """pathlib.Path objects should work the same as string paths.""" @@ -126,9 +120,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -139,9 +131,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes should be converted using a fallback MIME type.""" content = b"raw bytes content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -164,9 +154,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes with an image MIME type should produce type=image_url.""" content = b"raw image content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content, "mime_type": "image/jpeg"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content, "mime_type": "image/jpeg"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/jpeg;base64,") @@ -176,9 +164,7 @@ class TestConvertFileDocumentToUrlDocument: content = b"file-like content" file_obj = BytesIO(content) - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -189,9 +175,7 @@ class TestConvertFileDocumentToUrlDocument: file_obj = BytesIO(content) file_obj.name = "test_image.png" - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -204,9 +188,7 @@ class TestConvertFileDocumentToUrlDocument: def test_should_raise_error_for_nonexistent_pathlib_path(self): """Non-existent pathlib.Path should raise FileNotFoundError.""" with pytest.raises(FileNotFoundError, match="File not found"): - convert_file_document_to_url_document( - {"type": "file", "file": Path("/nonexistent/path/to/file.pdf")} - ) + convert_file_document_to_url_document({"type": "file", "file": Path("/nonexistent/path/to/file.pdf")}) def test_should_raise_error_for_empty_file(self): """Empty file should raise ValueError.""" @@ -215,9 +197,7 @@ class TestConvertFileDocumentToUrlDocument: try: with pytest.raises(ValueError, match="File is empty"): - convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + convert_file_document_to_url_document({"type": "file", "file": tmp_path}) finally: os.unlink(str(tmp_path)) @@ -248,9 +228,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path, "mime_type": "image/png"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path, "mime_type": "image/png"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -477,9 +455,7 @@ class TestProxySecurityGuard: result = await self._parse_multipart(mock_request) assert result["document"]["type"] == "document_url" - assert result["document"]["document_url"].startswith( - "data:application/pdf;base64," - ) + assert result["document"]["document_url"].startswith("data:application/pdf;base64,") assert result["model"] == "mistral/mistral-ocr-latest" diff --git a/tests/unit/passthrough/__init__.py b/tests/unit/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/unit/passthrough/test_async_streaming_error_propagation.py similarity index 92% rename from tests/test_litellm/passthrough/test_async_streaming_error_propagation.py rename to tests/unit/passthrough/test_async_streaming_error_propagation.py index 9f2b436d2d8..cb93183957c 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/unit/passthrough/test_async_streaming_error_propagation.py @@ -21,9 +21,7 @@ def _make_mock_response(status_code: int, body: bytes, headers: dict = None): # def _raise_for_status(): if status_code >= 400: - request = httpx.Request( - "POST", "https://azure.example.com/openai/responses" - ) + request = httpx.Request("POST", "https://azure.example.com/openai/responses") real_response = httpx.Response( status_code=status_code, content=body, @@ -55,16 +53,15 @@ def _make_mock_logging_obj(): async def test_async_streaming_429_raises(): """429 from upstream should raise HTTPStatusError, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "429", "message": "Rate limit exceeded."}} - ).encode() + + error_body = json.dumps({"error": {"code": "429", "message": "Rate limit exceeded."}}).encode() mock_response = _make_mock_response(429, error_body) - + async def response_coro(): return mock_response - + chunks = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -84,15 +81,13 @@ async def test_async_streaming_429_raises(): async def test_async_streaming_500_raises(): """500 from upstream should also raise, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "500", "message": "Internal server error"}} - ).encode() + + error_body = json.dumps({"error": {"code": "500", "message": "Internal server error"}}).encode() mock_response = _make_mock_response(500, error_body) - + async def response_coro(): return mock_response - + with pytest.raises(httpx.HTTPStatusError) as exc_info: async for _ in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -100,7 +95,7 @@ async def test_async_streaming_500_raises(): provider_config=MagicMock(), ): pass - + assert exc_info.value.response.status_code == 500 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py similarity index 94% rename from tests/test_litellm/passthrough/test_passthrough_main.py rename to tests/unit/passthrough/test_passthrough_main.py index 3f2c434cc00..82825ec2802 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -3,14 +3,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi.testclient import TestClient - -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - - - import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.passthrough.main import allm_passthrough_route, llm_passthrough_route @@ -37,10 +32,7 @@ def test_llm_passthrough_route(): client=client, ) - assert ( - mock_post.call_args.kwargs["request"].url - == "http://localhost:8090/v1/chat/completions" - ) + assert mock_post.call_args.kwargs["request"].url == "http://localhost:8090/v1/chat/completions" assert response.status_code == 200 assert response.json == {"message": "Hello, world!"} @@ -74,12 +66,9 @@ def test_bedrock_application_inference_profile_url_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -132,12 +121,9 @@ def test_bedrock_non_application_inference_profile_no_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -202,7 +188,6 @@ def test_update_stream_param_based_on_request_body(): @pytest.fixture def mock_request(): """Create a mock request with headers""" - from typing import Optional class QueryParams: def __init__(self): @@ -215,9 +200,7 @@ def mock_request(): return self._dict.items() class MockRequest: - def __init__( - self, headers=None, method="POST", request_body: Optional[dict] = None - ): + def __init__(self, headers=None, method="POST", request_body: dict | None = None): self.headers = headers or {} self.query_params = QueryParams() self.method = method @@ -245,9 +228,7 @@ def mock_user_api_key_dict(): @pytest.mark.asyncio -async def test_pass_through_request_stream_param_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_override(mock_request, mock_user_api_key_dict): """ Test that when stream=None is passed as parameter but stream=True is in request body, the request body value takes precedence and @@ -346,9 +327,7 @@ async def test_pass_through_request_stream_param_override( @pytest.mark.asyncio -async def test_pass_through_request_stream_param_no_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_no_override(mock_request, mock_user_api_key_dict): """ Test that when stream=False is passed as parameter and no stream is in request body, the function parameter is used and @@ -448,15 +427,11 @@ def test_azure_with_custom_api_base_and_key(): # Mock the provider config and its methods mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01" - ), + httpx.URL("https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01"), "https://my-custom-base", ) mock_provider_config.get_api_key.return_value = "my-custom-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "my-custom-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "my-custom-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "my-custom-key"}, None, @@ -484,13 +459,10 @@ def test_azure_with_custom_api_base_and_key(): patch.object( client.client, "send", - return_value=MagicMock( - status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []} - ), - ) as mock_send, + return_value=MagicMock(status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []}), + ), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -541,9 +513,7 @@ def test_content_param_forwarded_to_build_request(): mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "test-key" @@ -575,7 +545,6 @@ def test_content_param_forwarded_to_build_request(): patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -656,15 +625,11 @@ async def test_allm_passthrough_route_429_streaming_raises(): """ mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "fake-azure-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "fake-azure-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "fake-azure-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "fake-azure-key"}, None, @@ -752,9 +717,7 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): headers={"content-type": "application/json"}, ) - sync_client = HTTPHandler( - client=httpx.Client(transport=httpx.MockTransport(_handler)) - ) + sync_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_handler))) mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( @@ -762,18 +725,14 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): "https://gigachat.devices.sberbank.ru/api/v1", ) mock_provider_config.get_api_key.return_value = "fake-key" - mock_provider_config.validate_environment.return_value = { - "Authorization": "Bearer fake-key" - } + mock_provider_config.validate_environment.return_value = {"Authorization": "Bearer fake-key"} mock_provider_config.sign_request.return_value = ( {"Authorization": "Bearer fake-key"}, None, ) mock_provider_config.is_streaming_request.return_value = True - mock_provider_config.get_error_class.side_effect = ( - lambda error_message, status_code, headers: BaseLLMException( - status_code=status_code, message=error_message, headers=headers - ) + mock_provider_config.get_error_class.side_effect = lambda error_message, status_code, headers: BaseLLMException( + status_code=status_code, message=error_message, headers=headers ) mock_logging_obj = MagicMock() diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py similarity index 91% rename from tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py rename to tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py index 5e13db9439b..922643f9834 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py @@ -68,9 +68,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -88,7 +86,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): received.append(chunk) assert received == chunks - + assert received_response.headers["content-type"] == "application/octet-stream" assert received_response.headers["x-request-id"] == "req-123" @@ -107,9 +105,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): b'{"chunk": 3, "outputTokens": 8}', ] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -138,17 +134,13 @@ async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx(): err_response = MagicMock(spec=httpx.Response) err_response.status_code = 429 - err_response.headers = httpx.Headers( - {"content-type": "application/octet-stream"} - ) + err_response.headers = httpx.Headers({"content-type": "application/octet-stream"}) def _raise(): raise httpx.HTTPStatusError( "429", request=httpx.Request("POST", "https://example.com"), - response=httpx.Response( - 429, request=httpx.Request("POST", "https://example.com") - ), + response=httpx.Response(429, request=httpx.Request("POST", "https://example.com")), ) err_response.raise_for_status = _raise @@ -180,9 +172,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_response.status_code = 200 mock_response.raise_for_status = MagicMock(return_value=None) mock_response.aclose = AsyncMock() - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def _aiter_bytes_then_raise(): for c in partial_chunks: @@ -197,6 +187,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_logging_obj = _make_logging_obj() received = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -222,9 +213,7 @@ def test_passthroughstreamingresponse_flushes_on_normal_completion(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks @@ -258,9 +247,7 @@ def test_passthroughstreamingresponse_flushes_on_early_close(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks diff --git a/tests/unit/rag/ingestion/__init__.py b/tests/unit/rag/ingestion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py similarity index 94% rename from tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py rename to tests/unit/rag/ingestion/test_s3_vectors_ingestion.py index 07fd2b765f3..1e5b62456b6 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py @@ -21,10 +21,14 @@ class _RecordingRouter: def _ingestion(embedding=REQUEST_EMBEDDING, router=None, **vector_store): vector_store_options = {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store} - ingest_options = {"vector_store": vector_store_options} if embedding is None else { - "embedding": embedding, - "vector_store": vector_store_options, - } + ingest_options = ( + {"vector_store": vector_store_options} + if embedding is None + else { + "embedding": embedding, + "vector_store": vector_store_options, + } + ) return S3VectorsRAGIngestion(ingest_options=ingest_options, router=router) diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py similarity index 98% rename from tests/test_litellm/realtime_api/test_main.py rename to tests/unit/realtime_api/test_main.py index 86b25b2f9c8..5d3276dfae1 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -12,6 +12,15 @@ from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class FakeLogging: def update_from_kwargs(self, **kwargs): pass @@ -502,8 +511,8 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None): - from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig captured: dict[str, object] = {} diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py similarity index 98% rename from tests/test_litellm/repositories/test_repositories.py rename to tests/unit/repositories/test_repositories.py index 63fde9b2b8f..87cf2fc4268 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -78,17 +78,11 @@ class MockTable: record_data = dict(data) if self._pk_field and self._pk_field not in record_data: record_data[self._pk_field] = f"{self._pk_field}-{len(self._records)}" - key = ( - record_data.get(self._pk_field) - if self._pk_field - else record_data.get("id", str(len(self._records))) - ) + key = record_data.get(self._pk_field) if self._pk_field else record_data.get("id", str(len(self._records))) self._records[key] = record_data return MockRecord(record_data) - async def update( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Optional[MockRecord]: + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Optional[MockRecord]: key_field = list(where.keys())[0] key_value = where[key_field] if key_value in self._records: @@ -140,9 +134,7 @@ class MockPrismaClient: self.db.litellm_config = MockTable() self.db.litellm_organizationtable = MockTable() self.db.litellm_projecttable = MockTable(pk_field="project_id") - self.db.litellm_objectpermissiontable = MockTable( - pk_field="object_permission_id" - ) + self.db.litellm_objectpermissiontable = MockTable(pk_field="object_permission_id") self.db.litellm_credentialstable = MockTable() @@ -200,9 +192,7 @@ class TestBaseRepository: prisma_client.db.litellm_budgettable._records = { "b1": {"budget_id": "b1", "max_budget": 100.0}, } - budgets = await repo.find_many( - where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"} - ) + budgets = await repo.find_many(where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"}) assert len(budgets) == 1 def test_record_to_dict_branches(self): @@ -1518,9 +1508,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1563,9 +1551,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1578,9 +1564,7 @@ class TestVerificationTokenRepositoryExtended: await repo.delete_token("sk-arch", deleted_by="admin") - archived = list( - repo._prisma_client.db.litellm_deletedverificationtoken._records.values() - )[0] + archived = list(repo._prisma_client.db.litellm_deletedverificationtoken._records.values())[0] assert isinstance(archived["aliases"], str) assert json.loads(archived["aliases"]) == {"a": "b"} @@ -1599,9 +1583,7 @@ class TestVerificationTokenRepositoryExtended: ): assert relation_field not in archived - assert ( - "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records - ) + assert "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records @pytest.mark.asyncio async def test_find_by_id_maps_org_and_budget_columns(self, repo): @@ -1977,9 +1959,7 @@ class TestDomainModelExtended: DomainModel.from_db_record(None) def test_from_db_record_dict(self): - model = _SampleDomainModel.from_db_record( - {"budget_id": "b1", "max_budget": 100.0} - ) + model = _SampleDomainModel.from_db_record({"budget_id": "b1", "max_budget": 100.0}) assert model.budget_id == "b1" def test_from_db_record_model_dump(self): @@ -2174,9 +2154,7 @@ class TestPrismaTableRepository: assert self.CONFIG_SYNCED_TABLE_NAMES <= seen -def _json_path_equals( - metadata: Optional[Dict[str, Any]], path: List[str], expected: Any -) -> bool: +def _json_path_equals(metadata: Optional[Dict[str, Any]], path: List[str], expected: Any) -> bool: """Reproduce Postgres jsonb path-equals semantics: a missing path yields SQL NULL, which never matches `equals`.""" value: Any = metadata @@ -2201,11 +2179,7 @@ class _ScimAwareUserTable: json_filter = where["metadata"] path = json_filter["path"] expected = getattr(json_filter["equals"], "data", json_filter["equals"]) - return sum( - 1 - for metadata in self._metadatas - if _json_path_equals(metadata, path, expected) - ) + return sum(1 for metadata in self._metadatas if _json_path_equals(metadata, path, expected)) class TestCountBillableUsers: diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py similarity index 100% rename from tests/test_litellm/repositories/test_unit_of_work.py rename to tests/unit/repositories/test_unit_of_work.py diff --git a/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py similarity index 100% rename from tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py rename to tests/unit/router_strategy/complexity_router/test_jev_classifier.py diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index b5651062098..d8bc4c45ab8 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -14,6 +14,30 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data @@ -43,9 +67,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -348,9 +370,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): model_group=model_group, user_key=user_api_key_hash, ) - await router.cache.async_set_cache( - affinity_cache_key, {"model_id": other_model_id}, ttl=3600 - ) + await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600) # Even though user-key affinity points elsewhere, previous_response_id should pin # to the deployment that created the original response. @@ -519,9 +539,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -542,9 +560,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s model="some-router-model-group", healthy_deployments=healthy_deployments, messages=None, - request_kwargs={ - "metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"} - }, + request_kwargs={"metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"}}, parent_otel_span=None, ) @@ -580,9 +596,7 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -618,9 +632,7 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -660,9 +672,7 @@ async def test_async_user_key_affinity_ttl_expiry_allows_reroute(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -706,9 +716,7 @@ def test_cache_key_does_not_double_hash_user_api_key_hash(): The affinity cache key should not hash it again. """ - user_api_key_hash = ( - "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" - ) + user_api_key_hash = "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" key = DeploymentAffinityCheck.get_affinity_cache_key( model_group="any-model-group", user_key=user_api_key_hash, @@ -746,9 +754,7 @@ def test_get_effective_flags_returns_per_group_config(): assert session_id is True # unconfigured-model: falls back to global flags - user_key, responses_api, session_id = callback._get_effective_flags( - "unconfigured-model" - ) + user_key, responses_api, session_id = callback._get_effective_flags("unconfigured-model") assert user_key is True assert responses_api is True assert session_id is False @@ -980,12 +986,8 @@ async def test_model_group_affinity_config_overrides_global(): ] # Set up user-key affinity cache for claude-3 - cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group=stable_model_map_key, user_key=user_key - ) - await callback.cache.async_set_cache( - cache_key, {"model_id": "deployment-1"}, ttl=60 - ) + cache_key = DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) + await callback.cache.async_set_cache(cache_key, {"model_id": "deployment-1"}, ttl=60) # claude-3 has per-group config (session_affinity only), so user-key affinity # should NOT apply even though it's globally enabled @@ -1050,7 +1052,7 @@ async def test_async_jwt_user_affinity_routes_to_same_deployment(): return seq[0] return seq[1] if len(seq) > 1 else seq[0] - with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + with patch( "litellm.router_strategy.simple_shuffle.random.choice", side_effect=deterministic_choice, ): diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py similarity index 96% rename from tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index b93b8c1cdfc..aa34fbd6bf7 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -25,6 +25,31 @@ from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse + +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -1088,21 +1113,19 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) @@ -1114,21 +1137,19 @@ def test_boundary_key_matches_named_credential_precedence(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1146,21 +1167,19 @@ def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1178,37 +1197,35 @@ def test_boundary_fallback_matches_deployments_with_same_named_credential_values EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-a-peer", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-b", - credential_values={ - "api_base": "https://account-b.example.com", - "api_key": "credential-key-b", - }, - credential_info={}, - ), - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-a-peer", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-b", + credential_values={ + "api_base": "https://account-b.example.com", + "api_key": "credential-key-b", + }, + credential_info={}, + ), + ], ): router = litellm.Router( model_list=[ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py similarity index 94% rename from tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff31..849edc8c537 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,10 +1,9 @@ import asyncio import copy -from typing import List, cast +from typing import cast import pytest - import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT @@ -22,6 +21,15 @@ MODEL_GROUP_ALIAS = "my-claude-group" OPUS_4_6_MIN_TOKENS = 4096 +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + @pytest.fixture(autouse=True) def _local_model_cost_map_autouse(local_model_cost_map): """Every test here reads `prompt_cache_min_tokens`, which only the in-repo map @@ -30,8 +38,7 @@ def _local_model_cost_map_autouse(local_model_cost_map): yield - -def _deployments(*models: str) -> List[dict]: +def _deployments(*models: str) -> list[dict]: return [ { "model_name": MODEL_GROUP_ALIAS, @@ -42,9 +49,9 @@ def _deployments(*models: str) -> List[dict]: ] -def _messages(word_count: int) -> List[AllMessageValues]: +def _messages(word_count: int) -> list[AllMessageValues]: return cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "user", @@ -84,7 +91,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum(): """ messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True + ) assert 1024 < token_count < 4096 assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False @@ -110,7 +119,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -136,7 +147,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=5000) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert token_count > OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -197,10 +210,10 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" -def _auto_caching_messages() -> List[AllMessageValues]: +def _auto_caching_messages() -> list[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast( - List[AllMessageValues], + list[AllMessageValues], [ {"role": "system", "content": "word " * 3000}, {"role": "user", "content": "hello"}, @@ -208,7 +221,7 @@ def _auto_caching_messages() -> List[AllMessageValues]: ) -def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValues]: +def _affinity_messages(messages: list[AllMessageValues]) -> list[AllMessageValues]: """The messages the check keys deployment affinity on, for a group of `AUTO_CACHING_MODEL`.""" return AnthropicCacheControlHook.messages_with_default_injections( messages=messages, @@ -218,7 +231,7 @@ def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValue class _SentMessagesCapture(CustomLogger): def __init__(self): - self.messages: List[AllMessageValues] | None = None + self.messages: list[AllMessageValues] | None = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_logging_object = kwargs.get("standard_logging_object") @@ -338,7 +351,7 @@ async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_aff cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) - messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) request_kwargs = { "system": [ { @@ -441,7 +454,7 @@ def test_client_supplied_cache_control_keeps_its_own_prefix_boundary(monkeypatch """ monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "system", @@ -491,7 +504,7 @@ async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): warm_tokenizer("anthropic/claude-fable-5") check = PromptCachingDeploymentCheck(cache=DualCache()) deployments = _deployments("anthropic/claude-fable-5") - messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": text * 100}]) result, took, lags = await timed_with_loop_lags( lambda: check.async_filter_deployments( @@ -516,7 +529,7 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], ) standard_logging_object = { diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py index ee7fab7d19f..78cafbec70a 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py @@ -1,21 +1,10 @@ import asyncio -from typing import Optional +import json from unittest.mock import AsyncMock, patch import pytest -import json - import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.llms.openai import ( - IncompleteDetails, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, -) -from litellm.types.utils import StandardLoggingPayload @pytest.mark.asyncio @@ -119,14 +108,11 @@ async def test_async_responses_api_routing_with_previous_response_id(): input="Hello, how are you?", truncation="auto", ) - print("RESPONSE", response) # Store the model_id from the response expected_model_id = response._hidden_params["model_id"] response_id = response.id - print("Response ID=", response_id, "came from model_id=", expected_model_id) - # Make 10 other requests with previous_response_id, assert that they are sent to the same model_id for i in range(10): # Reset the mock for the next call @@ -137,7 +123,7 @@ async def test_async_responses_api_routing_with_previous_response_id(): response = await router.aresponses( model=MODEL, - input=f"Follow-up question {i+1}", + input=f"Follow-up question {i + 1}", truncation="auto", previous_response_id=response_id, ) @@ -163,9 +149,7 @@ async def test_async_routing_without_previous_response_id(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -266,9 +250,7 @@ async def test_async_routing_without_previous_response_id(): used_model_ids.add(response._hidden_params["model_id"]) # We should have used more than one model_id if load balancing is working - assert ( - len(used_model_ids) > 1 - ), "Load balancing isn't working, only one deployment was used" + assert len(used_model_ids) > 1, "Load balancing isn't working, only one deployment was used" @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py rename to tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 780300bf9e1..9bbaed0ae1b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,12 +1,10 @@ import asyncio +import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest - -import json - import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache @@ -46,9 +44,7 @@ async def test_async_session_id_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -164,9 +160,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): ) await callback.cache.async_set_cache( - DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1", user_key="user1" - ), + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "session1", user_key="user1"), {"model_id": "deployment-2"}, ) @@ -175,9 +169,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): model="model_group", healthy_deployments=healthy_deployments, messages=[], - request_kwargs={ - "metadata": {"user_api_key_hash": "user1", "session_id": "session1"} - }, + request_kwargs={"metadata": {"user_api_key_hash": "user1", "session_id": "session1"}}, ) assert len(filtered) == 1 @@ -575,16 +567,17 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): (None, {"model": "second"}), ], ) -async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( - stored: object, expected: object -) -> None: +async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) clock.return_value = 105.0 winner: Final = await claim_affinity_pin( - cache, "tier-pin", {"model": "second"}, 30, + cache, + "tier-pin", + {"model": "second"}, + 30, eligible_values=({"model": "first"}, {"model": "second"}), ) @@ -600,13 +593,18 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( async def test_concurrent_eligible_claims_return_one_winner() -> None: cache: Final = DualCache() candidates: Final = ({"model": "first"}, {"model": "second"}) - winners: Final = await asyncio.gather(*( - claim_affinity_pin( - cache, "tier-pin", candidates[index % 2], 30, - eligible_values=candidates, + winners: Final = await asyncio.gather( + *( + claim_affinity_pin( + cache, + "tier-pin", + candidates[index % 2], + 30, + eligible_values=candidates, + ) + for index in range(20) ) - for index in range(20) - )) + ) assert winners == [{"model": "first"}] * 20 assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} @@ -628,23 +626,19 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) clock.return_value = 105.0 - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "7"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30) assert winner == expected - assert cache.in_memory_cache.ttl_dict["deployment-pin"] == ( - 135.0 if refresh else 110.0 - ) - assert cache.in_memory_cache.get_cache("deployment-pin") == ( - {"model_id": "7"} if refresh else stored - ) + assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0) + assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored) @pytest.mark.asyncio @@ -668,13 +662,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( redis.async_register_script.return_value = AsyncMock(return_value=raw) cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "candidate"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) assert winner == expected assert cache.in_memory_cache.get_cache("deployment-pin") == stored diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/chat_completions/__init__.py b/tests/unit/rust_bridge/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/chat_completions/test_route_host.py b/tests/unit/rust_bridge/chat_completions/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/chat_completions/test_route_host.py rename to tests/unit/rust_bridge/chat_completions/test_route_host.py diff --git a/tests/unit/rust_bridge/messages/__init__.py b/tests/unit/rust_bridge/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/messages/test_route_host.py rename to tests/unit/rust_bridge/messages/test_route_host.py From 73a35abeb303dff68f232719bac0893cddf69366 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:57:53 +0000 Subject: [PATCH 081/127] refactor(types): keep object permission parsing as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/object_permission_utils.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 61e432daa16..daab38d3662 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,13 +156,12 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) - if raw_object_permission is None: + new_object_permission: dict | str | None = data_json.pop("object_permission", None) + if new_object_permission is None: return None - new_object_permission: Final[object] = ( - json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission - ) + if isinstance(new_object_permission, str): + new_object_permission = json.loads(new_object_permission) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -231,7 +230,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), From eea00175655063cd398f7200e9663cc8be2aa6c5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:13 +0000 Subject: [PATCH 082/127] test: add __init__.py to new tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/models/__init__.py | 0 tests/unit/realtime_api/__init__.py | 0 tests/unit/repositories/__init__.py | 0 tests/unit/router_strategy/complexity_router/__init__.py | 0 tests/unit/router_utils/pre_call_checks/__init__.py | 0 5 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/models/__init__.py create mode 100644 tests/unit/realtime_api/__init__.py create mode 100644 tests/unit/repositories/__init__.py create mode 100644 tests/unit/router_strategy/complexity_router/__init__.py create mode 100644 tests/unit/router_utils/pre_call_checks/__init__.py diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/realtime_api/__init__.py b/tests/unit/realtime_api/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/repositories/__init__.py b/tests/unit/repositories/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/complexity_router/__init__.py b/tests/unit/router_strategy/complexity_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/pre_call_checks/__init__.py b/tests/unit/router_utils/pre_call_checks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 42dd6a130016083166339c3ede9539a2c94a162c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:15 +0000 Subject: [PATCH 083/127] test: add __init__.py to phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/responses/__init__.py | 0 tests/unit/llms/parallel_ai/__init__.py | 0 tests/unit/llms/parasail/__init__.py | 0 tests/unit/llms/perplexity/chat/__init__.py | 0 tests/unit/llms/perplexity/embedding/__init__.py | 0 tests/unit/llms/perplexity/responses/__init__.py | 0 tests/unit/llms/publicai/__init__.py | 0 tests/unit/llms/ragflow/chat/__init__.py | 0 tests/unit/llms/recraft/image_edit/__init__.py | 0 tests/unit/llms/recraft/image_generation/__init__.py | 0 tests/unit/llms/runwayml/__init__.py | 0 tests/unit/llms/runwayml/videos/__init__.py | 0 tests/unit/llms/s3_vectors/vector_stores/__init__.py | 0 tests/unit/llms/sap/__init__.py | 0 tests/unit/llms/scaleway/__init__.py | 0 tests/unit/llms/snowflake/__init__.py | 0 tests/unit/llms/soniox/__init__.py | 0 tests/unit/llms/stability/image_generation/__init__.py | 0 tests/unit/llms/tencent/chat/__init__.py | 0 19 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/responses/__init__.py create mode 100644 tests/unit/llms/parallel_ai/__init__.py create mode 100644 tests/unit/llms/parasail/__init__.py create mode 100644 tests/unit/llms/perplexity/chat/__init__.py create mode 100644 tests/unit/llms/perplexity/embedding/__init__.py create mode 100644 tests/unit/llms/perplexity/responses/__init__.py create mode 100644 tests/unit/llms/publicai/__init__.py create mode 100644 tests/unit/llms/ragflow/chat/__init__.py create mode 100644 tests/unit/llms/recraft/image_edit/__init__.py create mode 100644 tests/unit/llms/recraft/image_generation/__init__.py create mode 100644 tests/unit/llms/runwayml/__init__.py create mode 100644 tests/unit/llms/runwayml/videos/__init__.py create mode 100644 tests/unit/llms/s3_vectors/vector_stores/__init__.py create mode 100644 tests/unit/llms/sap/__init__.py create mode 100644 tests/unit/llms/scaleway/__init__.py create mode 100644 tests/unit/llms/snowflake/__init__.py create mode 100644 tests/unit/llms/soniox/__init__.py create mode 100644 tests/unit/llms/stability/image_generation/__init__.py create mode 100644 tests/unit/llms/tencent/chat/__init__.py diff --git a/tests/unit/llms/openrouter/responses/__init__.py b/tests/unit/llms/openrouter/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parallel_ai/__init__.py b/tests/unit/llms/parallel_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parasail/__init__.py b/tests/unit/llms/parasail/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/chat/__init__.py b/tests/unit/llms/perplexity/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/embedding/__init__.py b/tests/unit/llms/perplexity/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/responses/__init__.py b/tests/unit/llms/perplexity/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/publicai/__init__.py b/tests/unit/llms/publicai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/chat/__init__.py b/tests/unit/llms/ragflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_edit/__init__.py b/tests/unit/llms/recraft/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_generation/__init__.py b/tests/unit/llms/recraft/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/__init__.py b/tests/unit/llms/runwayml/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/videos/__init__.py b/tests/unit/llms/runwayml/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/vector_stores/__init__.py b/tests/unit/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/sap/__init__.py b/tests/unit/llms/sap/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/scaleway/__init__.py b/tests/unit/llms/scaleway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/snowflake/__init__.py b/tests/unit/llms/snowflake/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/soniox/__init__.py b/tests/unit/llms/soniox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/image_generation/__init__.py b/tests/unit/llms/stability/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/chat/__init__.py b/tests/unit/llms/tencent/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From cd07acbdea3bc1d579e88e612c870b1b92b8ea28 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:37:21 +0000 Subject: [PATCH 084/127] test: add __init__.py to intermediate phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/perplexity/__init__.py | 0 tests/unit/llms/ragflow/__init__.py | 0 tests/unit/llms/recraft/__init__.py | 0 tests/unit/llms/s3_vectors/__init__.py | 0 tests/unit/llms/stability/__init__.py | 0 tests/unit/llms/tencent/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/perplexity/__init__.py create mode 100644 tests/unit/llms/ragflow/__init__.py create mode 100644 tests/unit/llms/recraft/__init__.py create mode 100644 tests/unit/llms/s3_vectors/__init__.py create mode 100644 tests/unit/llms/stability/__init__.py create mode 100644 tests/unit/llms/tencent/__init__.py diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/__init__.py b/tests/unit/llms/perplexity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/__init__.py b/tests/unit/llms/ragflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/__init__.py b/tests/unit/llms/recraft/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/__init__.py b/tests/unit/llms/s3_vectors/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/__init__.py b/tests/unit/llms/stability/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/__init__.py b/tests/unit/llms/tencent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 4fb1a3e61043a8308de3159417f1fca3a329d52e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:55:27 +0000 Subject: [PATCH 085/127] test: add __init__.py to intermediate tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rag/__init__.py | 0 tests/unit/router_strategy/__init__.py | 0 tests/unit/router_utils/__init__.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/rag/__init__.py create mode 100644 tests/unit/router_strategy/__init__.py create mode 100644 tests/unit/router_utils/__init__.py diff --git a/tests/unit/rag/__init__.py b/tests/unit/rag/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/__init__.py b/tests/unit/router_strategy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/__init__.py b/tests/unit/router_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30f3075010964a024480df032cc2bf818ac43607 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:44:30 +0000 Subject: [PATCH 086/127] test: use main's local_model_cost_map fixture in tests/unit/conftest.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 73a555e4455..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,7 +2,6 @@ import os from collections.abc import Iterator from typing import Final -import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -23,19 +22,6 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( ) -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) From d086f8574dcaac12c494fa9d37d60cab3b5a5ab7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:46:52 +0000 Subject: [PATCH 087/127] test: add sync and async custom_llm_provider bridge propagation tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py index 09ef1889818..f2e36137a19 100644 --- a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -40,6 +40,58 @@ def _bedrock_mantle_kwargs() -> dict: } +def _openai_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="completion", + model="gpt-5.5", + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": "gpt-5.5", + "custom_llm_provider": "openai", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +def test_completion_forwards_custom_llm_provider_to_responses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + with patch("litellm.responses", return_value=cached) as fake_responses: + result = bridge.completion(**_openai_kwargs()) + + assert result is cached + assert fake_responses.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_custom_llm_provider_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_openai_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_acompletion_forwards_aws_region_name_to_aresponses(): bridge = ResponsesToCompletionBridgeHandler() From e1556ce32bc87fa17ccac09fd1abf27cc3fb3979 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:18:49 +0000 Subject: [PATCH 088/127] test: migrate legacy provider tests to tests/unit (wave 2, phase 13) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../files/test_vertex_ai_files_integration.py | 214 ------------------ tests/unit/llms/tencent/__init__.py | 0 tests/unit/llms/tencent/messages/__init__.py | 0 ...ncent_anthropic_messages_transformation.py | 0 tests/unit/llms/together_ai/__init__.py | 0 tests/unit/llms/together_ai/chat/__init__.py | 0 .../test_together_ai_chat_transformation.py | 0 tests/unit/llms/valkey/__init__.py | 0 .../llms/valkey/vector_stores/__init__.py | 0 .../test_valkey_transformation.py | 10 +- tests/unit/llms/vercel_ai_gateway/__init__.py | 0 .../llms/vercel_ai_gateway/chat/__init__.py | 0 .../test_vercel_ai_gateway_transformation.py | 0 .../vercel_ai_gateway/embedding/__init__.py | 0 .../test_vercel_ai_gateway_embedding.py | 0 tests/unit/llms/vertex_ai/__init__.py | 0 .../llms/vertex_ai/agent_engine/__init__.py | 0 .../agent_engine/test_transformation.py | 0 .../vertex_ai/context_caching/__init__.py | 0 .../test_context_caching_ttl.py | 0 .../test_vertex_ai_context_caching.py | 94 -------- tests/unit/llms/vertex_ai/files/__init__.py | 0 .../test_file_retrieve_provider_routing.py | 0 .../test_vertex_ai_binary_file_upload.py | 68 ------ .../files/test_vertex_ai_files_handler.py | 104 --------- .../files/test_vertex_ai_files_integration.py | 93 ++++++++ .../files/test_vertex_ai_files_streaming.py | 0 .../test_vertex_ai_files_transformation.py | 88 ------- .../vertex_ai/gemini_embeddings/__init__.py | 0 ...test_batch_embed_content_transformation.py | 0 .../llms/vertex_ai/image_edit/__init__.py | 0 ...est_vertex_ai_image_edit_transformation.py | 0 .../llms/vertex_ai/interactions/__init__.py | 0 ...t_vertex_ai_interactions_transformation.py | 0 .../multimodal_embeddings/__init__.py | 0 ..._ai_multimodal_embedding_transformation.py | 0 .../unit/llms/vertex_ai/realtime/__init__.py | 0 .../test_vertex_ai_realtime_transformation.py | 44 ---- .../llms/vertex_ai/text_to_speech/__init__.py | 0 .../text_to_speech/test_transformation.py | 0 40 files changed, 94 insertions(+), 621 deletions(-) delete mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py create mode 100644 tests/unit/llms/tencent/__init__.py create mode 100644 tests/unit/llms/tencent/messages/__init__.py rename tests/{test_litellm => unit}/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py (100%) create mode 100644 tests/unit/llms/together_ai/__init__.py create mode 100644 tests/unit/llms/together_ai/chat/__init__.py rename tests/{test_litellm => unit}/llms/together_ai/chat/test_together_ai_chat_transformation.py (100%) create mode 100644 tests/unit/llms/valkey/__init__.py create mode 100644 tests/unit/llms/valkey/vector_stores/__init__.py rename tests/{test_litellm => unit}/llms/valkey/vector_stores/test_valkey_transformation.py (97%) create mode 100644 tests/unit/llms/vercel_ai_gateway/__init__.py create mode 100644 tests/unit/llms/vercel_ai_gateway/chat/__init__.py rename tests/{test_litellm => unit}/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py (100%) create mode 100644 tests/unit/llms/vercel_ai_gateway/embedding/__init__.py rename tests/{test_litellm => unit}/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py (100%) create mode 100644 tests/unit/llms/vertex_ai/__init__.py create mode 100644 tests/unit/llms/vertex_ai/agent_engine/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/agent_engine/test_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/context_caching/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/context_caching/test_context_caching_ttl.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py (95%) create mode 100644 tests/unit/llms/vertex_ai/files/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_file_retrieve_provider_routing.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py (71%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_handler.py (75%) create mode 100644 tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_streaming.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_vertex_ai_files_transformation.py (95%) create mode 100644 tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/image_edit/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/interactions/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/realtime/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py (92%) create mode 100644 tests/unit/llms/vertex_ai/text_to_speech/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/text_to_speech/test_transformation.py (100%) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py deleted file mode 100644 index 8f9acafa49d..00000000000 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ /dev/null @@ -1,214 +0,0 @@ -""" -Test Vertex AI files integration with main files API -""" - -import pytest -from unittest.mock import AsyncMock, MagicMock, patch - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent - - -class TestVertexAIFilesIntegration: - """Test integration of Vertex AI files with main litellm API""" - - @pytest.mark.asyncio - async def test_litellm_afile_content_vertex_ai_provider(self): - """Test litellm.afile_content with vertex_ai provider""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content since the code - # now routes through ProviderConfigManager -> base_llm_http_handler - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - # Make it return a coroutine for async path - mock_retrieve.return_value = mock_result - - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - - def test_litellm_file_content_vertex_ai_provider(self): - """Test litellm.file_content with vertex_ai provider (sync)""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - return_value=mock_result, - ) as mock_retrieve: - result = litellm.file_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 - - # Verify the mock was called - mock_retrieve.assert_called_once() - - def test_litellm_file_content_vertex_ai_with_model_provider_detection(self): - """Test litellm.file_content with model parameter for provider detection""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - return_value=mock_result, - ): - # Mock get_llm_provider to return vertex_ai - with patch("litellm.files.main.get_llm_provider") as mock_get_provider: - mock_get_provider.return_value = ( - "vertex_ai/gemini-pro", - "vertex_ai", - None, - None, - ) - - # Call litellm.file_content with model to trigger provider detection - result = litellm.file_content( - file_id=file_id, - model="vertex_ai/gemini-pro", - vertex_project="test-project", - vertex_location="us-central1", - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - - # Verify provider detection was called - mock_get_provider.assert_called_once() - - def test_litellm_file_content_vertex_ai_error_cases(self): - """Test error handling in vertex_ai file_content""" - # Test missing file_id - the VertexAI provider config's - # transform_file_content_request should handle empty file_id. - # Since the code now goes through base_llm_http_handler, we mock - # ProviderConfigManager to return None so it falls through to the - # old vertex_ai code path that validates file_id. - with patch( - "litellm.files.main.ProviderConfigManager.get_provider_files_config", - return_value=None, - ): - with pytest.raises(ValueError, match="file_id is required"): - litellm.file_content( - file_id="", # Empty file_id should cause error - custom_llm_provider="vertex_ai", - vertex_project="test-project", - ) - - def test_vertex_ai_provider_in_supported_providers_list(self): - """Test that vertex_ai is included in supported providers for file_content""" - # This test ensures the type annotations and error messages include vertex_ai - - # Test that calling with unsupported provider raises appropriate error - with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: - litellm.file_content( - file_id="test-file-id", - custom_llm_provider="unsupported_provider", # This should fail - ) - - # The error message should mention supported providers including vertex_ai - error_message = str(exc_info.value) - assert "vertex_ai" in error_message or "supported" in error_message.lower() - - @pytest.mark.asyncio - async def test_vertex_ai_file_content_with_timeout_and_retries(self): - """Test vertex_ai file_content with timeout and retry configuration""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock the base_llm_http_handler.retrieve_file_content - with patch( - "litellm.files.main.base_llm_http_handler.retrieve_file_content", - new_callable=MagicMock, - ) as mock_retrieve: - mock_retrieve.return_value = mock_result - - # Call with custom timeout and max_retries - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="vertex_ai", - vertex_project="test-project", - vertex_location="us-central1", - timeout=120, - max_retries=5, - ) - - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - - # Verify the mock was called - mock_retrieve.assert_called_once() - # Verify the timeout was passed through - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["timeout"] == 120 diff --git a/tests/unit/llms/tencent/__init__.py b/tests/unit/llms/tencent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/messages/__init__.py b/tests/unit/llms/tencent/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py b/tests/unit/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py rename to tests/unit/llms/tencent/messages/test_tencent_anthropic_messages_transformation.py diff --git a/tests/unit/llms/together_ai/__init__.py b/tests/unit/llms/together_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/together_ai/chat/__init__.py b/tests/unit/llms/together_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py rename to tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py diff --git a/tests/unit/llms/valkey/__init__.py b/tests/unit/llms/valkey/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/valkey/vector_stores/__init__.py b/tests/unit/llms/valkey/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py similarity index 97% rename from tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py rename to tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py index aa114f128c5..64e15008536 100644 --- a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py @@ -1,8 +1,7 @@ import struct -import sys from types import SimpleNamespace from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from urllib.parse import unquote, urlsplit import httpx @@ -355,13 +354,6 @@ def test_search_treats_an_explicit_null_max_num_results_as_the_default(): assert client.index.searched_query.query_string() == "*=>[KNN 10 @embedding $vec AS vector_distance]" -def test_missing_redis_dependency_raises_actionable_error(): - config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0])) - blocked = {name: None for name in list(sys.modules) if name == "redis" or name.startswith("redis.")} - - with patch.dict(sys.modules, blocked): - with pytest.raises(ValueError, match="pip install redis"): - _search(config) @pytest.mark.asyncio diff --git a/tests/unit/llms/vercel_ai_gateway/__init__.py b/tests/unit/llms/vercel_ai_gateway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vercel_ai_gateway/chat/__init__.py b/tests/unit/llms/vercel_ai_gateway/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py b/tests/unit/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py rename to tests/unit/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py diff --git a/tests/unit/llms/vercel_ai_gateway/embedding/__init__.py b/tests/unit/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py b/tests/unit/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py rename to tests/unit/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py diff --git a/tests/unit/llms/vertex_ai/__init__.py b/tests/unit/llms/vertex_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/agent_engine/__init__.py b/tests/unit/llms/vertex_ai/agent_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/agent_engine/test_transformation.py rename to tests/unit/llms/vertex_ai/agent_engine/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/context_caching/__init__.py b/tests/unit/llms/vertex_ai/context_caching/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/unit/llms/vertex_ai/context_caching/test_context_caching_ttl.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py rename to tests/unit/llms/vertex_ai/context_caching/test_context_caching_ttl.py diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py similarity index 95% rename from tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py rename to tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 34c00e84d2e..4ba2d69a38b 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -9,7 +9,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( - MAX_PAGINATION_PAGES, ContextCachingEndpoints, ) @@ -1892,100 +1891,7 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): - """Test that pagination stops after MAX_PAGINATION_PAGES iterations""" - # Setup - mock_get_token_url.return_value = ("token", "https://test-url.com") - cache_key_to_find = "nonexistent_cache_key" - # Create mock response that always has nextPageToken (infinite pagination scenario) - def create_page_response(page_num): - response = MagicMock() - response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], - "nextPageToken": f"token_page_{page_num + 1}", - } - return response - - # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken - self.mock_client.get.side_effect = [ - create_page_response(i) for i in range(MAX_PAGINATION_PAGES) - ] - - # Execute - result = self.context_caching.check_cache( - cache_key=cache_key_to_find, - client=self.mock_client, - headers={"Authorization": "Bearer token"}, - api_key="test_key", - api_base=None, - logging_obj=self.mock_logging, - custom_llm_provider=custom_llm_provider, - vertex_project="test_project", - vertex_location="us-central1", - vertex_auth_header="Bearer test-token", - ) - - # Assert - should return None after exhausting all pages without finding match - assert result is None - # Verify exactly MAX_PAGINATION_PAGES API calls were made (not more) - assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): - """Test that async pagination stops after MAX_PAGINATION_PAGES iterations""" - # Setup - mock_get_token_url.return_value = ("token", "https://test-url.com") - cache_key_to_find = "nonexistent_cache_key" - - # Create mock response that always has nextPageToken (infinite pagination scenario) - def create_page_response(page_num): - response = MagicMock() - response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], - "nextPageToken": f"token_page_{page_num + 1}", - } - return response - - # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken - self.mock_async_client.get = AsyncMock( - side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)] - ) - - # Execute - result = await self.context_caching.async_check_cache( - cache_key=cache_key_to_find, - client=self.mock_async_client, - headers={"Authorization": "Bearer token"}, - api_key="test_key", - api_base=None, - logging_obj=self.mock_logging, - custom_llm_provider=custom_llm_provider, - vertex_project="test_project", - vertex_location="us-central1", - vertex_auth_header="Bearer test-token", - ) - - # Assert - should return None after exhausting all pages without finding match - assert result is None - # Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more) - assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES class TestVertexAIGlobalLocation: diff --git a/tests/unit/llms/vertex_ai/files/__init__.py b/tests/unit/llms/vertex_ai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py b/tests/unit/llms/vertex_ai/files/test_file_retrieve_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py rename to tests/unit/llms/vertex_ai/files/test_file_retrieve_provider_routing.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py similarity index 71% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index d2ee9d7d659..f4aee11c140 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -11,8 +11,6 @@ import io import json import pytest -import httpx - from litellm.llms.custom_httpx.llm_http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig from litellm.types.llms.openai import CreateFileRequest @@ -96,39 +94,6 @@ class TestVertexAIBinaryFileUpload: assert isinstance(transformed_request, bytes) assert transformed_request == mock_png_content - @pytest.mark.asyncio - async def test_http_handler_accepts_bytes_without_decoding(self): - """ - Test that httpx correctly accepts binary data without decoding. - - This test verifies that bytes can be passed to httpx's post/put methods - without needing UTF-8 decoding, which is the core of our fix. - """ - # Create mock binary data with non-UTF-8 bytes - mock_binary_data = b"\x00\x01\x02\x03\xff\xfe\xfd\xc4\xe5\xf2" - - # Test that httpx accepts bytes in the data parameter - # We're testing the behavior, not making an actual request - - # Verify that attempting to decode would fail (proving it's binary) - with pytest.raises(UnicodeDecodeError): - mock_binary_data.decode("utf-8") - - # Verify that httpx Request accepts bytes - try: - request = httpx.Request( - method="POST", - url="https://example.com/upload", - data=mock_binary_data, - headers={"Content-Type": "application/octet-stream"}, - ) - # If we get here, httpx accepts bytes - which is what we need - assert request.content == mock_binary_data - except Exception as e: - pytest.fail(f"httpx should accept bytes in data parameter: {e}") - - # Document the expected behavior - assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes" @pytest.mark.asyncio async def test_jsonl_file_upload_returns_streaming_body(self): @@ -224,36 +189,3 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) assert isinstance(result3, bytes) - - def test_bytes_type_preservation_documentation(self): - """ - Documentation test: Verify that bytes are the correct type for binary uploads. - - This test documents the expected behavior: - - Binary files (PDF, images, etc.) should remain as bytes - - Text files (JSONL) should be strings - - httpx accepts both bytes and strings in the 'data' parameter - - bytes should NEVER be decoded to UTF-8 for binary files - """ - # This is a documentation test - it always passes - # but serves as a reference for the expected behavior - - expected_behavior = { - "binary_files": { - "input_type": "bytes", - "output_type": "bytes", - "examples": ["PDF", "PNG", "JPEG", "binary data"], - "http_method": "POST or PUT", - "encoding": "none - preserve raw bytes", - }, - "text_files": { - "input_type": "str or bytes", - "output_type": "bytes", - "examples": ["JSONL", "CSV", "TXT"], - "http_method": "POST", - "encoding": "UTF-8", - }, - } - - assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" - assert expected_behavior["text_files"]["encoding"] == "UTF-8" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py similarity index 75% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 0a44f0a9a74..e0f0b7e5c0b 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -2,14 +2,11 @@ Test Vertex AI files handler functionality """ -import asyncio import re from types import MappingProxyType import pytest from unittest.mock import AsyncMock, patch -import httpx - from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent @@ -312,104 +309,3 @@ class TestVertexAIFilesHandler: assert isinstance(result, HttpxBinaryResponseContent) dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"] assert dynamic_params["gcs_bucket_name"] == "my-model-bucket" - - def test_file_content_sync_success(self): - """Test successful sync file content retrieval""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - # Create expected response - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - expected_result = HttpxBinaryResponseContent(response=mock_response) - - # Mock asyncio.run to return our expected result - with patch("asyncio.run") as mock_run: - mock_run.return_value = expected_result - - result = self.handler.file_content( - _is_async=False, - file_content_request=file_content_request, - api_base="", - vertex_credentials=None, - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=3, - ) - - # Verify the result - assert result == expected_result - - # Verify asyncio.run was called (indicating sync execution) - mock_run.assert_called_once() - - @pytest.mark.asyncio - async def test_file_content_async_mode(self): - """Test async file content retrieval when _is_async=True""" - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - expected_content = b"test file content" - - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - # Mock the afile_content method - with patch.object(self.handler, "afile_content", new_callable=AsyncMock) as mock_afile_content: - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), - ) - mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response) - - # Call the method with _is_async=True - result = self.handler.file_content( - _is_async=True, - file_content_request=file_content_request, - api_base="", - vertex_credentials=None, - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=3, - ) - - # Should return a coroutine since _is_async=True - assert asyncio.iscoroutine(result) - - # Await the result - final_result = await result - assert isinstance(final_result, HttpxBinaryResponseContent) - assert final_result.response.content == expected_content - - def test_httpx_response_compatibility(self): - """Test that the created HttpxBinaryResponseContent is compatible with expected interface""" - # Test the mock response creation logic - expected_content = b"test file content" - decoded_path = "gs://test-bucket/test-file.txt" - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=decoded_path), - ) - - result = HttpxBinaryResponseContent(response=mock_response) - - # Verify the response properties - assert result.response.status_code == 200 - assert result.response.content == expected_content - assert result.response.headers["content-type"] == "application/octet-stream" - - # Verify it has the expected interface (matching OpenAI file content response) - assert hasattr(result, "response") - assert hasattr(result.response, "content") - assert hasattr(result.response, "status_code") - assert hasattr(result.response, "headers") diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py new file mode 100644 index 00000000000..402c463d134 --- /dev/null +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -0,0 +1,93 @@ +""" +Test Vertex AI files integration with main files API +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class TestVertexAIFilesIntegration: + """Test integration of Vertex AI files with main litellm API""" + + + + + def test_litellm_file_content_vertex_ai_error_cases(self): + """Test error handling in vertex_ai file_content""" + # Test missing file_id - the VertexAI provider config's + # transform_file_content_request should handle empty file_id. + # Since the code now goes through base_llm_http_handler, we mock + # ProviderConfigManager to return None so it falls through to the + # old vertex_ai code path that validates file_id. + with patch( + "litellm.files.main.ProviderConfigManager.get_provider_files_config", + return_value=None, + ): + with pytest.raises(ValueError, match="file_id is required"): + litellm.file_content( + file_id="", # Empty file_id should cause error + custom_llm_provider="vertex_ai", + vertex_project="test-project", + ) + + def test_vertex_ai_provider_in_supported_providers_list(self): + """Test that vertex_ai is included in supported providers for file_content""" + # This test ensures the type annotations and error messages include vertex_ai + + # Test that calling with unsupported provider raises appropriate error + with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: + litellm.file_content( + file_id="test-file-id", + custom_llm_provider="unsupported_provider", # This should fail + ) + + # The error message should mention supported providers including vertex_ai + error_message = str(exc_info.value) + assert "vertex_ai" in error_message or "supported" in error_message.lower() + + @pytest.mark.asyncio + async def test_vertex_ai_file_content_with_timeout_and_retries(self): + """Test vertex_ai file_content with timeout and retry configuration""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Create a mock HttpxBinaryResponseContent response + import httpx + + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"), + ) + mock_result = HttpxBinaryResponseContent(response=mock_response) + + # Mock the base_llm_http_handler.retrieve_file_content + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file_content", + new_callable=MagicMock, + ) as mock_retrieve: + mock_retrieve.return_value = mock_result + + # Call with custom timeout and max_retries + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + timeout=120, + max_retries=5, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + + # Verify the mock was called + mock_retrieve.assert_called_once() + # Verify the timeout was passed through + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs["timeout"] == 120 diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py similarity index 95% rename from tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py rename to tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8a249820cbd..48464e79876 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -860,94 +860,6 @@ class TestVertexBatchOutputTransformation: binary = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\n" + b"\x00\x01\x02\xff\xfe" * 64 assert config._try_transform_vertex_batch_output_to_openai(binary) == binary - def test_streaming_transform_peaks_below_list_pipeline(self, config): - """The output transform must stream row-by-row, not build a list of every - parsed row and a second list of transformed rows. This guards against a - regression to the list pipeline, which peaks at several full copies and - OOMs on large result files. The relative comparison cancels shared noise - (per-row transform cost, GC timing) and only the list overhead differs. - """ - import gc - import tracemalloc - - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - def vertex_row(index: int) -> dict: - return { - "status": "", - "processed_time": "2024-11-01T18:13:16.826+00:00", - "request": { - "contents": [{"role": "user", "parts": [{"text": "hi"}]}], - "labels": {"litellm_custom_id": f"r-{index}"}, - }, - "response": { - "candidates": [ - { - "content": { - "parts": [{"text": "hello " * 20}], - "role": "model", - }, - "finishReason": "STOP", - } - ], - "modelVersion": "gemini-2.0-flash-001", - "usageMetadata": { - "promptTokenCount": 10, - "candidatesTokenCount": 20, - "totalTokenCount": 30, - }, - }, - } - - content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode("utf-8") - - def list_pipeline() -> bytes: - gemini_config = VertexGeminiConfig() - logging_obj = Logging( - model="", - messages=[], - stream=False, - call_type="batch_transform", - start_time=0.1, - litellm_call_id="", - function_id="", - ) - logging_obj.optional_params = {} - mock_response = httpx.Response( - status_code=200, - headers={"content-type": "application/json"}, - request=httpx.Request("POST", "https://example.com"), - ) - rows = content.decode("utf-8").strip().split("\n") - transformed = [ - json.dumps( - config._transform_single_vertex_batch_output_to_openai( - json.loads(row), gemini_config, logging_obj, mock_response - ) - ) - for row in rows - ] - return "\n".join(transformed).encode("utf-8") - - def peak_of(fn) -> int: - gc.collect() - tracemalloc.start() - try: - fn() - return tracemalloc.get_traced_memory()[1] - finally: - tracemalloc.stop() - - streaming_peak = peak_of(lambda: config._try_transform_vertex_batch_output_to_openai(content)) - list_peak = peak_of(list_pipeline) - - assert streaming_peak < list_peak * 0.75, ( - f"streaming peak {streaming_peak} is not a clear win over the list " - f"pipeline {list_peak} (ratio {streaming_peak / list_peak:.2f})" - ) class TestTryTransformDoesNotMutateCallerLoggingObj: diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py rename to tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py diff --git a/tests/unit/llms/vertex_ai/image_edit/__init__.py b/tests/unit/llms/vertex_ai/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py rename to tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py diff --git a/tests/unit/llms/vertex_ai/interactions/__init__.py b/tests/unit/llms/vertex_ai/interactions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py rename to tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py rename to tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py diff --git a/tests/unit/llms/vertex_ai/realtime/__init__.py b/tests/unit/llms/vertex_ai/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py similarity index 92% rename from tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py rename to tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index d4cf58bc0b4..a71ef21d997 100644 --- a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -387,50 +387,6 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) ) -@pytest.mark.asyncio -async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( - monkeypatch, -): - """Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors. - - Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params`` - (the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here. - """ - import websockets - - from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler - - cfg = VertexAIRealtimeConfig( - access_token="tok", project="my-proj", location="us-central1" - ) - - captured: dict = {} - - def fake_connect(url, *args, **kwargs): - captured["url"] = url - raise RuntimeError("stop before establishing the backend connection") - - monkeypatch.setattr(websockets, "connect", fake_connect) - - try: - await BaseLLMHTTPHandler().async_realtime( - model="gemini-live-2.5-flash-preview-native-audio-09-2025", - websocket=AsyncMock(), - logging_obj=MagicMock(), - provider_config=cfg, - headers={}, - query_params={ - "model": "gemini-live-2.5-flash-preview-native-audio-09-2025", - "intent": "chat", - }, - ) - except (RuntimeError, Exception): - pass - - assert "url" in captured, "websockets.connect was never called" - assert "?" not in captured["url"] - assert "model=" not in captured["url"] - assert "intent=" not in captured["url"] def test_vertex_function_call_output_omits_id(): diff --git a/tests/unit/llms/vertex_ai/text_to_speech/__init__.py b/tests/unit/llms/vertex_ai/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py rename to tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py From a5033fff6e243dcd58dfa108bb83f3a638287283 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:20:04 +0000 Subject: [PATCH 089/127] test: define local_model_cost_map fixture for migrated context caching tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_context_caching.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 4ba2d69a38b..a438b62c94c 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -13,6 +14,25 @@ from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ) +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled in-repo cost map so capability and pricing assertions do not + depend on the network-fetched ``main`` copy, which lags this branch until merge. + + ``get_model_info`` is lru_cached, so swapping ``model_cost`` is not enough on its + own; clear on the way in and out so entries warmed against either map never leak + across tests.""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + class TestContextCachingEndpoints: """Test class for ContextCachingEndpoints methods""" From 5ddf6ff3963d0700bebbb08213f303bcc0adbfdc Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 14:53:26 +0000 Subject: [PATCH 090/127] test: restore pagination limit tests and realtime query param regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_context_caching.py | 81 +++++++++++++++++++ .../test_vertex_ai_realtime_transformation.py | 37 +++++++++ 2 files changed, 118 insertions(+) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index a438b62c94c..7913700c8a7 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1911,6 +1911,87 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(100) + ] + + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_client.get.call_count == 100 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(100)] + ) + + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + assert result is None + assert self.mock_async_client.get.call_count == 100 + diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a71ef21d997..14b3bdb48a1 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -387,6 +387,43 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) ) +@pytest.mark.asyncio +async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( + monkeypatch, +): + import websockets + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + captured = {} + + def fake_connect(url, *args, **kwargs): + captured["url"] = url + raise RuntimeError("stop before establishing the backend connection") + + monkeypatch.setattr(websockets, "connect", fake_connect) + + await BaseLLMHTTPHandler().async_realtime( + model="gemini-live-2.5-flash-preview-native-audio-09-2025", + websocket=AsyncMock(), + logging_obj=MagicMock(), + provider_config=cfg, + headers={}, + query_params={ + "model": "gemini-live-2.5-flash-preview-native-audio-09-2025", + "intent": "chat", + }, + ) + + assert "?" not in captured["url"] + assert "model=" not in captured["url"] + assert "intent=" not in captured["url"] + + def test_vertex_function_call_output_omits_id(): From d93cc8defda03a9c2cb7fb86903190cce1dcdc07 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sun, 20 Sep 2026 08:39:51 -0700 Subject: [PATCH 091/127] test(mcp): run public client regressions in the MCP shard --- .github/workflows/test-unit.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index aa82a0bf3ee..49e6d7040d4 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -51,7 +51,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/mcp_tests" + test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" workers: 2 reruns: 0 timeout-minutes: 20 @@ -113,7 +113,6 @@ jobs: tests/test_litellm/compression tests/test_litellm/containers tests/test_litellm/endpoints - tests/test_litellm/experimental_mcp_client tests/test_litellm/models tests/test_litellm/repositories tests/test_litellm/images From e3911c71f74a46eef73d789c433cfa72ed72c3fd Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 16:24:36 +0000 Subject: [PATCH 092/127] test: make path-sourced streaming peak test differential Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../files/test_vertex_ai_files_streaming.py | 53 +++++++++---------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index b94ea1ea269..a824a1b75d9 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -99,6 +99,17 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st ) +def _measure_peak(fn) -> int: + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + class TestStreamingOutputParity: def test_transform_create_file_request_returns_streaming_body_parity(self): cfg = VertexAIFilesConfig() @@ -258,16 +269,6 @@ class TestStreamingPeakMemory: measurement removes any garbage the previous run left behind. """ - def _measure(self, fn): - gc.collect() - tracemalloc.start() - try: - fn() - _, peak = tracemalloc.get_traced_memory() - finally: - tracemalloc.stop() - return peak - def test_streaming_peak_well_below_list_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(8000) @@ -279,8 +280,8 @@ class TestStreamingPeakMemory: for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes(): pass - streaming_peak = self._measure(drain_stream) - list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str)) + streaming_peak = _measure_peak(drain_stream) + list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) # Core guard: the lazily consumed streaming body peaks well under a list # pipeline that materializes every transformed row. Building full @@ -298,7 +299,7 @@ class TestStreamingPeakMemory: # The payload bytes already exist before measurement starts, so a lazy # first-row parse should allocate only a small fraction of the payload; # parsing every row would blow past this bound. - peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) + peak = _measure_peak(lambda: cfg.get_object_name(file_data, purpose="batch")) assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload" @@ -344,12 +345,12 @@ class TestPathSourcedStreaming: first_labels = json.loads(lines[0])["request"]["labels"] assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" - def test_path_source_peak_stays_below_payload(self, tmp_path): + def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): cfg = VertexAIFilesConfig() - path, raw = self._write_jsonl(tmp_path, 8000) + path, _ = self._write_jsonl(tmp_path, 8000) data = self._batch_request(path) - def run(): + def drain_stream(): cfg.get_complete_file_url( api_base=None, api_key=None, @@ -362,19 +363,17 @@ class TestPathSourcedStreaming: model="", create_file_data=data, optional_params={}, litellm_params={} ) for _ in _upload_stream(out).iter_bytes(): - pass # drain without accumulating + pass - gc.collect() - tracemalloc.start() - try: - run() - _, peak = tracemalloc.get_traced_memory() - finally: - tracemalloc.stop() + streaming_peak = _measure_peak(drain_stream) + list_peak = _measure_peak( + lambda: _reference_vertex_jsonl_string(cfg, path.read_bytes().decode("utf-8")) + ) - # Streaming from disk must not materialize the payload. Reading the whole - # file into bytes (the pre-fix path) would push peak past the file size. - assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})" + assert streaming_peak < list_peak * 0.3, ( + f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline " + f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) def test_path_source_stream_is_reiterable(self, tmp_path): cfg = VertexAIFilesConfig() From 8554d1f8275d9b6e3eebaa0c7def5d4b24354c02 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 17:08:16 +0000 Subject: [PATCH 093/127] test: prepare reference input before tracing path-sourced peak Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/test_vertex_ai_files_streaming.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index a824a1b75d9..29eaf38b427 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -347,8 +347,9 @@ class TestPathSourcedStreaming: def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): cfg = VertexAIFilesConfig() - path, _ = self._write_jsonl(tmp_path, 8000) + path, raw = self._write_jsonl(tmp_path, 8000) data = self._batch_request(path) + content_str = raw.decode("utf-8") def drain_stream(): cfg.get_complete_file_url( @@ -366,9 +367,7 @@ class TestPathSourcedStreaming: pass streaming_peak = _measure_peak(drain_stream) - list_peak = _measure_peak( - lambda: _reference_vertex_jsonl_string(cfg, path.read_bytes().decode("utf-8")) - ) + list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) assert streaming_peak < list_peak * 0.3, ( f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline " From a6361dc08cc1d9c50b072ecf38d92bfa27df8049 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 17:30:49 +0000 Subject: [PATCH 094/127] chore(prices): sync OpenRouter prices: 11 models openrouter/~deepseek/deepseek-flash-latest: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-flash-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/meta/muse-glimmer-30b: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token openrouter/mistralai/mistral-small-3.1-24b-instruct: supports_tool_choice, supports_function_calling openrouter/qwen/qwen-2.5-coder-32b-instruct: supports_response_schema openrouter/qwen/qwen3.8-27b: input_cost_per_token, output_cost_per_token openrouter/z-ai/glm-5.3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 65 ++++++++++--------- model_prices_and_context_window.json | 65 ++++++++++--------- 2 files changed, 66 insertions(+), 64 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bfb4729ef32..e92712fd6d9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,21 +41579,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.3724e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost": 1.7908e-08, "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, @@ -42097,12 +42097,12 @@ "max_tokens": 102400, "mode": "chat", "output_cost_per_token": 5.55e-07, - "supports_tool_choice": false, + "supports_tool_choice": true, "max_input_tokens": 128000, "max_output_tokens": 102400, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, @@ -42775,7 +42775,7 @@ "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_response_schema": false, + "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": false, "supports_web_search": false @@ -66570,9 +66570,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.96e-07, - "output_cost_per_token": 2.816e-06, - "cache_read_input_token_cost": 1.664e-07, + "input_cost_per_token": 9.1e-07, + "output_cost_per_token": 2.86e-06, + "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -66591,8 +66591,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2.55e-06, "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.668e-08, - "output_cost_per_token": 7.336e-08, - "cache_read_input_token_cost": 7.336e-09, + "input_cost_per_token": 3.556e-08, + "output_cost_per_token": 7.112e-08, + "cache_read_input_token_cost": 7.112e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71296,14 +71296,15 @@ "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 2.6e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 5.2e-07, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71316,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9228e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.7908e-08, + "input_cost_per_token": 5.3724e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71568,14 +71569,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 1.8e-08, + "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71588,14 +71589,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.5678e-07, - "input_cost_per_token": 8.442e-07, + "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.6532e-06, + "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -73022,13 +73023,13 @@ }, "openrouter/meta/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 3.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bfb4729ef32..e92712fd6d9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,21 +41579,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.3724e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 4.4e-08, + "cache_read_input_token_cost": 1.7908e-08, "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, @@ -42097,12 +42097,12 @@ "max_tokens": 102400, "mode": "chat", "output_cost_per_token": 5.55e-07, - "supports_tool_choice": false, + "supports_tool_choice": true, "max_input_tokens": 128000, "max_output_tokens": 102400, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, @@ -42775,7 +42775,7 @@ "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_response_schema": false, + "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": false, "supports_web_search": false @@ -66570,9 +66570,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.96e-07, - "output_cost_per_token": 2.816e-06, - "cache_read_input_token_cost": 1.664e-07, + "input_cost_per_token": 9.1e-07, + "output_cost_per_token": 2.86e-06, + "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -66591,8 +66591,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2.55e-06, "cache_read_input_token_cost": 8.5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.668e-08, - "output_cost_per_token": 7.336e-08, - "cache_read_input_token_cost": 7.336e-09, + "input_cost_per_token": 3.556e-08, + "output_cost_per_token": 7.112e-08, + "cache_read_input_token_cost": 7.112e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71296,14 +71296,15 @@ "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 2.6e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 5.2e-07, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71316,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.9228e-08, - "input_cost_per_token": 5.7684e-07, + "cache_read_input_token_cost": 1.7908e-08, + "input_cost_per_token": 5.3724e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.73052e-06, + "output_cost_per_token": 1.61172e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71568,14 +71569,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 1.8e-08, + "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71588,14 +71589,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.5678e-07, - "input_cost_per_token": 8.442e-07, + "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.6532e-06, + "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -73022,13 +73023,13 @@ }, "openrouter/meta/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 3.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 1.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 229d5cbfcbeac05a4e30d5dc24e22d0713016900 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 18:00:52 +0000 Subject: [PATCH 095/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e92712fd6d9..771657cd4d9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3724e-07, + "input_cost_per_token": 5.346e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7908e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.782e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7908e-08, - "input_cost_per_token": 5.3724e-07, + "cache_read_input_token_cost": 1.782e-08, + "input_cost_per_token": 5.346e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e92712fd6d9..771657cd4d9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3724e-07, + "input_cost_per_token": 5.346e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7908e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.782e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7908e-08, - "input_cost_per_token": 5.3724e-07, + "cache_read_input_token_cost": 1.782e-08, + "input_cost_per_token": 5.346e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.61172e-06, + "output_cost_per_token": 1.6038e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 010a5a5b1c71dd5b8e548051236ed6212dffbd80 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 18:30:53 +0000 Subject: [PATCH 096/127] chore(prices): sync OpenRouter prices: 5 models openrouter/~anthropic/claude-fable-latest: supports_tool_choice openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/anthropic/claude-fable-5.1: supports_tool_choice openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing openrouter/qwen/qwen3.6-35b-a3b: input_cost_per_token, output_cost_per_token --- ...odel_prices_and_context_window_backup.json | 30 +++++++++---------- model_prices_and_context_window.json | 30 +++++++++---------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 771657cd4d9..bf65675f3c5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.346e-07, + "input_cost_per_token": 5.3328e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.782e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, + "cache_read_input_token_cost": 1.6968e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65764,7 +65764,7 @@ "thinking_always_on": true, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, @@ -67155,8 +67155,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-35b-a3b": { - "input_cost_per_token": 1e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1e-06, "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -71225,7 +71225,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true }, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.782e-08, - "input_cost_per_token": 5.346e-07, + "cache_read_input_token_cost": 1.6968e-08, + "input_cost_per_token": 5.3328e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 771657cd4d9..bf65675f3c5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.346e-07, + "input_cost_per_token": 5.3328e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.782e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.346e-7,"output_cost_per_token":0.0000016038,"cache_read_input_token_cost":1.782e-8}, + "cache_read_input_token_cost": 1.6968e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65764,7 +65764,7 @@ "thinking_always_on": true, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, @@ -67155,8 +67155,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-35b-a3b": { - "input_cost_per_token": 1e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1e-06, "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -71225,7 +71225,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": false, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true }, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.782e-08, - "input_cost_per_token": 5.346e-07, + "cache_read_input_token_cost": 1.6968e-08, + "input_cost_per_token": 5.3328e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.6038e-06, + "output_cost_per_token": 1.59984e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 6a8b131fbdc0b041ffe8112924919ea4bb2dd2df Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:00:55 +0000 Subject: [PATCH 097/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bf65675f3c5..d325f46017a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3328e-07, + "input_cost_per_token": 5.2932e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6968e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7644e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6968e-08, - "input_cost_per_token": 5.3328e-07, + "cache_read_input_token_cost": 1.7644e-08, + "input_cost_per_token": 5.2932e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bf65675f3c5..d325f46017a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.3328e-07, + "input_cost_per_token": 5.2932e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6968e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.3328e-7,"output_cost_per_token":0.00000159984,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7644e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6968e-08, - "input_cost_per_token": 5.3328e-07, + "cache_read_input_token_cost": 1.7644e-08, + "input_cost_per_token": 5.2932e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.59984e-06, + "output_cost_per_token": 1.58796e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 9e84b8b048d2cec984bd4094c7243a48f6b95b3d Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:30:56 +0000 Subject: [PATCH 098/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d325f46017a..a5419899773 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2932e-07, + "input_cost_per_token": 5.2668e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7644e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7556e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7644e-08, - "input_cost_per_token": 5.2932e-07, + "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 5.2668e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d325f46017a..a5419899773 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2932e-07, + "input_cost_per_token": 5.2668e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7644e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2932e-7,"output_cost_per_token":0.00000158796,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.7556e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7644e-08, - "input_cost_per_token": 5.2932e-07, + "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 5.2668e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.58796e-06, + "output_cost_per_token": 1.58004e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 83e1f0834ebd3ea17738af52733cc01adaa19194 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:30:54 +0000 Subject: [PATCH 099/127] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-v4-flash-latest: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-0731: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-vision-exp: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a5419899773..a59ba1082da 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66550,13 +66550,13 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { - "input_cost_per_token": 2.156e-07, - "output_cost_per_token": 6.468e-07, - "cache_read_input_token_cost": 6.86e-09, + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 7e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a5419899773..a59ba1082da 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66550,13 +66550,13 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { - "input_cost_per_token": 2.156e-07, - "output_cost_per_token": 6.468e-07, - "cache_read_input_token_cost": 6.86e-09, + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 7e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8e-08, + "output_cost_per_token": 1.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 8a8a15af932ffb5ef3c5cf25a875d0c41b05fe5a Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:00:50 +0000 Subject: [PATCH 100/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-v4-flash-latest: output_cost_per_token openrouter/deepseek/deepseek-v4-flash-0731: output_cost_per_token --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a59ba1082da..d25cf0c78df 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a59ba1082da..d25cf0c78df 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66690,7 +66690,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, @@ -71344,7 +71344,7 @@ "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 84c26978a3dc144e44c5e5e55b41e6a3d129093e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 14:10:31 -0700 Subject: [PATCH 101/127] refactor(rust): split token counter backends --- litellm-rust/Cargo.lock | 83 ++++- litellm-rust/Cargo.toml | 4 + litellm-rust/crates/python-bridge/Cargo.toml | 7 +- .../crates/python-bridge/src/token_counter.rs | 72 ++++- .../crates/token-counter-fast/Cargo.toml | 19 ++ .../src/byte_level.rs | 20 +- .../src/cl100k.rs | 0 .../crates/token-counter-fast/src/error.rs | 13 + .../crates/token-counter-fast/src/lib.rs | 70 +++++ .../src/o200k.rs | 0 .../src/scanner.rs | 0 .../crates/token-counter-fast/src/tiktoken.rs | 222 ++++++++++++++ .../src/unicode_classes.rs | 0 .../tests/fixtures/cl100k/requests.jsonl | 0 .../tests/fixtures/cl100k/texts.jsonl | 0 .../tests/fixtures/generate.py | 0 .../tests/fixtures/o200k/requests.jsonl | 0 .../tests/fixtures/o200k/texts.jsonl | 0 .../token-counter-huggingface/Cargo.toml | 10 + .../token-counter-huggingface/src/lib.rs | 29 ++ .../crates/token-counter-tiktoken/Cargo.toml | 10 + .../crates/token-counter-tiktoken/src/lib.rs | 54 ++++ litellm-rust/crates/token-counter/Cargo.toml | 15 +- litellm-rust/crates/token-counter/README.md | 23 ++ .../token-counter/benches/allocations.rs | 2 +- .../token-counter/benches/token_counter.rs | 2 +- .../crates/token-counter/src/counter.rs | 65 +--- .../crates/token-counter/src/error.rs | 6 +- litellm-rust/crates/token-counter/src/fast.rs | 41 +++ .../crates/token-counter/src/huggingface.rs | 27 ++ litellm-rust/crates/token-counter/src/lib.rs | 15 +- .../crates/token-counter/src/tiktoken.rs | 226 +------------- .../crates/token-counter/src/tokenizer.rs | 31 ++ .../token-counter/tests/token_counter.rs | 288 +++++++++--------- litellm/rust_bridge/_native.pyi | 4 +- litellm/rust_bridge/token_counter.py | 40 ++- .../spend_tracking/test_budget_reservation.py | 14 +- .../rust_bridge/test_token_counter.py | 33 +- 38 files changed, 944 insertions(+), 501 deletions(-) create mode 100644 litellm-rust/crates/token-counter-fast/Cargo.toml rename litellm-rust/crates/{token-counter => token-counter-fast}/src/byte_level.rs (97%) rename litellm-rust/crates/{token-counter => token-counter-fast}/src/cl100k.rs (100%) create mode 100644 litellm-rust/crates/token-counter-fast/src/error.rs create mode 100644 litellm-rust/crates/token-counter-fast/src/lib.rs rename litellm-rust/crates/{token-counter => token-counter-fast}/src/o200k.rs (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/src/scanner.rs (100%) create mode 100644 litellm-rust/crates/token-counter-fast/src/tiktoken.rs rename litellm-rust/crates/{token-counter => token-counter-fast}/src/unicode_classes.rs (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/cl100k/requests.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/cl100k/texts.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/generate.py (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/o200k/requests.jsonl (100%) rename litellm-rust/crates/{token-counter => token-counter-fast}/tests/fixtures/o200k/texts.jsonl (100%) create mode 100644 litellm-rust/crates/token-counter-huggingface/Cargo.toml create mode 100644 litellm-rust/crates/token-counter-huggingface/src/lib.rs create mode 100644 litellm-rust/crates/token-counter-tiktoken/Cargo.toml create mode 100644 litellm-rust/crates/token-counter-tiktoken/src/lib.rs create mode 100644 litellm-rust/crates/token-counter/README.md create mode 100644 litellm-rust/crates/token-counter/src/fast.rs create mode 100644 litellm-rust/crates/token-counter/src/huggingface.rs create mode 100644 litellm-rust/crates/token-counter/src/tokenizer.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2720cf01f2e..62bde5806a4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -61,6 +61,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + [[package]] name = "arc-swap" version = "1.9.2" @@ -598,6 +604,17 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -1181,6 +1198,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fancy-regex" version = "0.19.2" @@ -1926,6 +1954,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -2083,7 +2117,7 @@ dependencies = [ name = "litellm-core-utils" version = "0.1.0" dependencies = [ - "fancy-regex", + "fancy-regex 0.19.2", "litellm-types", "rstest", "serde", @@ -2212,10 +2246,24 @@ dependencies = [ name = "litellm-token-counter" version = "0.1.0" dependencies = [ - "base64 0.22.1", "criterion", "indexmap 2.14.0", "itoa", + "litellm-token-counter-fast", + "litellm-token-counter-huggingface", + "litellm-token-counter-tiktoken", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "litellm-token-counter-fast" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", "rand 0.8.7", "rstest", "rustc-hash", @@ -2226,6 +2274,22 @@ dependencies = [ "unicode-normalization-alignments", ] +[[package]] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tokenizers", +] + +[[package]] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tiktoken-rs", +] + [[package]] name = "litellm-types" version = "0.1.0" @@ -3794,6 +3858,21 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.17.0", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.53" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a6185632871..109a2edaf4d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -24,6 +24,9 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-token-counter = { path = "crates/token-counter" } +litellm-token-counter-fast = { path = "crates/token-counter-fast" } +litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } +litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } bytes = "1" @@ -45,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 3a4a579efa3..2c868ebc769 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,10 +10,13 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3"] +default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] +token-counter-fast = ["litellm-token-counter/fast"] +token-counter-huggingface = ["litellm-token-counter/huggingface"] +token-counter-tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true @@ -26,7 +29,7 @@ litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true -litellm-token-counter.workspace = true +litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 7dc86b78ad6..a7c6bce4d0b 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -1,6 +1,19 @@ -use std::{num::NonZero, sync::Arc, thread::available_parallelism}; +use std::sync::Arc; -use litellm_host_python::{release_gil, run_async}; +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use std::{num::NonZero, thread::available_parallelism}; + +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use litellm_host_python::release_gil; +use litellm_host_python::run_async; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, }; @@ -28,17 +41,47 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + #[cfg(feature = "token-counter-huggingface")] + { + Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-huggingface"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-huggingface feature", + )) + } } #[staticmethod] - fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) + fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { + #[cfg(feature = "token-counter-fast")] + { + Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-fast"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-fast feature", + )) + } } #[staticmethod] - fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) + fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + #[cfg(feature = "token-counter-tiktoken")] + { + Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) + } + #[cfg(not(feature = "token-counter-tiktoken"))] + { + let _ = (py, encoding); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-tiktoken feature", + )) + } } fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { @@ -62,6 +105,11 @@ impl TokenCounter { } impl TokenCounter { + #[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" + ))] fn load( py: Python<'_>, load: impl FnOnce() -> Result + Send, @@ -74,6 +122,11 @@ impl TokenCounter { } } +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] fn encode_parallelism() -> usize { available_parallelism().map_or(1, NonZero::get) } @@ -86,7 +139,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { let message = error.to_string(); match error { - Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message), + Error::Load(_) + | Error::Ranks(_) + | Error::UnicodeClasses + | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), Error::RequestParse(_) | Error::MissingInput | Error::FloatText diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml new file mode 100644 index 00000000000..4950277f5b0 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-token-counter-fast" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +base64.workspace = true +rustc-hash = "2.1.3" +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +unicode-normalization-alignments = "0.1.12" + +[dev-dependencies] +rand.workspace = true +rstest.workspace = true +serde.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/token-counter/src/byte_level.rs b/litellm-rust/crates/token-counter-fast/src/byte_level.rs similarity index 97% rename from litellm-rust/crates/token-counter/src/byte_level.rs rename to litellm-rust/crates/token-counter-fast/src/byte_level.rs index ec6134a252e..6fc7d9146b9 100644 --- a/litellm-rust/crates/token-counter/src/byte_level.rs +++ b/litellm-rust/crates/token-counter-fast/src/byte_level.rs @@ -473,13 +473,13 @@ mod tests { _ => unreachable!(), } assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none()); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", " stop"] { assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -545,7 +545,7 @@ mod tests { .rstrip(rstrip)]) .expect("add token"); let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -557,7 +557,7 @@ mod tests { ] { assert_eq!(fast.count(&anthropic_tokenizer, text), None); assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -571,17 +571,17 @@ mod tests { assert_eq!(fast.count(&tokenizer, "hello"), None); assert!(tokenizer.encode_fast("hello", true).is_err()); let counter = - crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize")) + crate::FastTokenizer::from_json(&tokenizer.to_string(false).expect("serialize")) .expect("load"); assert!(matches!( - counter.count_text("hello"), + counter.count_tokens("hello"), Err(crate::Error::Encode(_)) )); } #[rstest] fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) { - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -598,7 +598,7 @@ mod tests { scope.spawn(move || { for _ in 0..100 { for (text, count) in inputs.iter().zip(expected) { - assert_eq!(counter.count_text(text).expect("count"), count); + assert_eq!(counter.count_tokens(text).expect("count"), count); } } }); @@ -614,10 +614,10 @@ mod tests { let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1); assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); - assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1); + assert_eq!(counter.count_tokens("ABCD EFGH").expect("count"), 1); } } diff --git a/litellm-rust/crates/token-counter/src/cl100k.rs b/litellm-rust/crates/token-counter-fast/src/cl100k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/cl100k.rs rename to litellm-rust/crates/token-counter-fast/src/cl100k.rs diff --git a/litellm-rust/crates/token-counter-fast/src/error.rs b/litellm-rust/crates/token-counter-fast/src/error.rs new file mode 100644 index 00000000000..e63ccf3ad39 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/error.rs @@ -0,0 +1,13 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("failed to load tokenizer: tiktoken rank file: {0}")] + Ranks(String), + #[error("failed to load tokenizer: Unicode character classes are unavailable")] + UnicodeClasses, + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} diff --git a/litellm-rust/crates/token-counter-fast/src/lib.rs b/litellm-rust/crates/token-counter-fast/src/lib.rs new file mode 100644 index 00000000000..ce91af642ea --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/lib.rs @@ -0,0 +1,70 @@ +#![forbid(unsafe_code)] + +mod byte_level; +mod cl100k; +mod error; +mod o200k; +mod scanner; +mod tiktoken; +mod unicode_classes; + +use byte_level::ByteLevelCounter; +use scanner::{SplitPattern, TiktokenCounter}; + +pub use error::Error; + +enum Encoder { + HuggingFace { + tokenizer: Box, + byte_level: Option, + }, + Tiktoken(TiktokenCounter), +} + +pub struct FastTokenizer(Encoder); + +impl FastTokenizer { + pub fn from_json(json: &str) -> Result { + let tokenizer = json.parse::().map_err(Error::Load)?; + let byte_level = ByteLevelCounter::detect(&tokenizer); + Ok(Self(Encoder::HuggingFace { + tokenizer: Box::new(tokenizer), + byte_level, + })) + } + + pub fn from_cl100k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::Cl100k, ranks) + } + + pub fn from_o200k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::O200k, ranks) + } + + fn from_ranks(split: SplitPattern, ranks: &str) -> Result { + TiktokenCounter::from_ranks(split, ranks) + .map(Encoder::Tiktoken) + .map(Self) + } + + pub fn count_tokens(&self, text: &str) -> Result { + match &self.0 { + Encoder::Tiktoken(counter) => Ok(counter.count(text)), + Encoder::HuggingFace { + tokenizer, + byte_level, + } => { + if let Some(count) = byte_level + .as_ref() + .and_then(|counter| counter.count(tokenizer, text)) + { + return Ok(count); + } + tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } + } + } +} diff --git a/litellm-rust/crates/token-counter/src/o200k.rs b/litellm-rust/crates/token-counter-fast/src/o200k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/o200k.rs rename to litellm-rust/crates/token-counter-fast/src/o200k.rs diff --git a/litellm-rust/crates/token-counter/src/scanner.rs b/litellm-rust/crates/token-counter-fast/src/scanner.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/scanner.rs rename to litellm-rust/crates/token-counter-fast/src/scanner.rs diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs new file mode 100644 index 00000000000..7a9e71ed587 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -0,0 +1,222 @@ +//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and +//! the merge loop that turns one regex piece into tokens. The merge order is +//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is +//! identical, but pairs are tracked in a heap so a long piece costs +//! `O(n log n)` instead of tiktoken's `O(n^2)`. + +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +use base64::Engine; +use base64::engine::general_purpose::STANDARD; +use rustc_hash::FxHashMap; + +use crate::Error; + +type Rank = u32; + +const NO_RANK: Rank = Rank::MAX; +const END: usize = usize::MAX; + +pub(super) struct MergeRanks(FxHashMap, Rank>); + +impl MergeRanks { + pub(super) fn parse(text: &str) -> Result { + let ranks = text + .lines() + .filter(|line| !line.is_empty()) + .map(parse_line) + .collect::, _>>()?; + if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { + return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); + } + Ok(Self(ranks)) + } + + fn rank(&self, bytes: &[u8]) -> Rank { + self.0.get(bytes).copied().unwrap_or(NO_RANK) + } + + /// Token count of one regex piece, as `encode_ordinary` would produce. + pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { + if piece.len() < 2 || self.0.contains_key(piece) { + return 1; + } + scratch.reset(piece.len()); + for start in 0..piece.len() - 1 { + scratch.set_rank(start, self.rank(&piece[start..start + 2])); + } + let mut parts = piece.len(); + while let Some(Reverse((rank, start))) = scratch.heap.pop() { + if scratch.next[start] == END || scratch.rank[start] != rank { + continue; + } + let merged = scratch.next[start]; + let after = scratch.next[merged]; + scratch.next[merged] = END; + scratch.next[start] = after; + parts -= 1; + if after < piece.len() { + scratch.prev[after] = start; + scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); + } else { + scratch.rank[start] = NO_RANK; + } + let before = scratch.prev[start]; + if before != END { + scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); + } + } + parts + } +} + +fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { + let (token, rank) = line + .split_once(' ') + .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; + let bytes = STANDARD + .decode(token) + .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; + let rank = rank + .parse() + .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; + Ok((bytes.into_boxed_slice(), rank)) +} + +/// Buffers reused across the pieces of one text. Parts are addressed by the +/// byte offset they start at, which also gives the leftmost-pair tie break. +#[derive(Default)] +pub(super) struct MergeScratch { + next: Vec, + prev: Vec, + rank: Vec, + heap: BinaryHeap>, +} + +impl MergeScratch { + fn reset(&mut self, len: usize) { + self.next.clear(); + self.next.extend(1..=len); + self.prev.clear(); + self.prev.push(END); + self.prev.extend(0..len - 1); + self.rank.clear(); + self.rank.resize(len, NO_RANK); + self.heap.clear(); + } + + fn end(&self, start: usize) -> usize { + self.next[start] + } + + fn set_rank(&mut self, start: usize, rank: Rank) { + self.rank[start] = rank; + if rank != NO_RANK { + self.heap.push(Reverse((rank, start))); + } + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + + fn ranks() -> MergeRanks { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" + ); + MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) + .expect("rank file parses") + } + + /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. + fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { + if piece.len() < 2 || ranks.0.contains_key(piece) { + return 1; + } + let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) + .map(|index| (index, ranks.rank(&piece[index..index + 2]))) + .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) + .collect(); + let get_rank = |parts: &[(usize, Rank)], index: usize| { + if index + 3 < parts.len() { + ranks.rank(&piece[parts[index].0..parts[index + 3].0]) + } else { + NO_RANK + } + }; + loop { + let Some(index) = parts[..parts.len() - 1] + .iter() + .enumerate() + .filter(|(_, (_, rank))| *rank != NO_RANK) + .min_by_key(|(index, (_, rank))| (*rank, *index)) + .map(|(index, _)| index) + else { + return parts.len() - 1; + }; + if index > 0 { + parts[index - 1].1 = get_rank(&parts, index - 1); + } + parts[index].1 = get_rank(&parts, index); + parts.remove(index + 1); + } + } + + #[test] + fn every_byte_is_a_token() { + let ranks = ranks(); + assert_eq!(ranks.0.len(), 100_256); + assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); + } + + #[test] + fn heap_merge_matches_tiktokens_merge_loop() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut rng = StdRng::seed_from_u64(99); + let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; + for _ in 0..20_000 { + let piece: Vec = (0..rng.gen_range(1..24)) + .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) + .collect(); + assert_eq!( + ranks.count_piece(&piece, &mut scratch), + reference_count(&ranks, &piece), + "piece {:?}", + String::from_utf8_lossy(&piece) + ); + } + } + + #[test] + fn long_repeated_runs_cost_close_to_linear() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut time = |len: usize| { + let piece = vec![b' '; len]; + let started = std::time::Instant::now(); + assert!(ranks.count_piece(&piece, &mut scratch) > 0); + started.elapsed() + }; + let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); + let large = time(1 << 18); + assert!( + large < small * 64, + "{small:?} for 2^14 bytes, {large:?} for 2^18" + ); + } + + #[test] + fn malformed_rank_files_are_rejected() { + assert!(MergeRanks::parse("IQ==").is_err()); + assert!(MergeRanks::parse("IQ== x").is_err()); + assert!(MergeRanks::parse("!!! 1").is_err()); + assert!(MergeRanks::parse("IQ== 1").is_err()); + } +} diff --git a/litellm-rust/crates/token-counter/src/unicode_classes.rs b/litellm-rust/crates/token-counter-fast/src/unicode_classes.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/unicode_classes.rs rename to litellm-rust/crates/token-counter-fast/src/unicode_classes.rs diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/generate.py rename to litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml new file mode 100644 index 00000000000..ad15d9f19d0 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs new file mode 100644 index 00000000000..184956b5f39 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -0,0 +1,29 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} + +pub struct HuggingFaceTokenizer(Box); + +impl HuggingFaceTokenizer { + pub fn from_json(json: &str) -> Result { + json.parse::() + .map(Box::new) + .map(Self) + .map_err(Error::Load) + } + + pub fn count_tokens(&self, text: &str) -> Result { + self.0 + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } +} diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml new file mode 100644 index 00000000000..494a9233e69 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tiktoken-rs.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs new file mode 100644 index 00000000000..b35c4876d00 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -0,0 +1,54 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +#[error("unsupported tokenizer: {0}")] +pub struct UnsupportedTokenizer(pub String); + +pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); + +impl TiktokenTokenizer { + pub fn from_name(name: &str) -> Result { + let tokenizer = match name { + "cl100k_base" => tiktoken_rs::cl100k_base_singleton(), + "o200k_base" => tiktoken_rs::o200k_base_singleton(), + "o200k_harmony" => tiktoken_rs::o200k_harmony_singleton(), + "p50k_base" => tiktoken_rs::p50k_base_singleton(), + "p50k_edit" => tiktoken_rs::p50k_edit_singleton(), + "r50k_base" | "gpt2" => tiktoken_rs::r50k_base_singleton(), + _ => return Err(UnsupportedTokenizer(name.to_owned())), + }; + Ok(Self(tokenizer)) + } + + pub fn count_tokens(&self, text: &str) -> usize { + self.0.count_ordinary(text) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn special_tokens_are_counted_as_ordinary_text() { + let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap(); + assert!(counter.count_tokens("<|endoftext|>") > 1); + } + + #[test] + fn all_python_tiktoken_encodings_are_available() { + for name in [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ] { + assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}"); + } + } +} diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index d0369631682..59e4f9a6a1d 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -5,16 +5,21 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +default = ["huggingface", "tiktoken"] +fast = ["dep:litellm-token-counter-fast"] +huggingface = ["dep:litellm-token-counter-huggingface"] +tiktoken = ["dep:litellm-token-counter-tiktoken"] + [dependencies] -base64.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } itoa = "1.0" -rustc-hash = "2.1.3" +litellm-token-counter-fast = { workspace = true, optional = true } +litellm-token-counter-huggingface = { workspace = true, optional = true } +litellm-token-counter-tiktoken = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } -unicode-normalization-alignments = "0.1.12" [dev-dependencies] criterion.workspace = true @@ -24,7 +29,9 @@ rstest.workspace = true [[bench]] name = "token_counter" harness = false +required-features = ["fast"] [[bench]] name = "allocations" harness = false +required-features = ["fast"] diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md new file mode 100644 index 00000000000..c8ed0737a7e --- /dev/null +++ b/litellm-rust/crates/token-counter/README.md @@ -0,0 +1,23 @@ +# Token counting + +`Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface + +The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation + +The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation + +The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` + +The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend + +Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits + +Run the feature matrix with: + +```sh +cargo test -p litellm-token-counter +cargo test -p litellm-token-counter --no-default-features +cargo test -p litellm-token-counter --no-default-features --features fast +cargo test -p litellm-token-counter --no-default-features --features huggingface +cargo test -p litellm-token-counter --no-default-features --features tiktoken +``` diff --git a/litellm-rust/crates/token-counter/benches/allocations.rs b/litellm-rust/crates/token-counter/benches/allocations.rs index 343f815749d..7aebceb6e9c 100644 --- a/litellm-rust/crates/token-counter/benches/allocations.rs +++ b/litellm-rust/crates/token-counter/benches/allocations.rs @@ -87,7 +87,7 @@ fn main() { }, ); - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("tokenizer loads"); let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses"); counter .count_request(&object) diff --git a/litellm-rust/crates/token-counter/benches/token_counter.rs b/litellm-rust/crates/token-counter/benches/token_counter.rs index a7c3177b88a..5e30cf6f77e 100644 --- a/litellm-rust/crates/token-counter/benches/token_counter.rs +++ b/litellm-rust/crates/token-counter/benches/token_counter.rs @@ -42,7 +42,7 @@ fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> { } fn token_counter(c: &mut Criterion) { - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("token counter should load"); let tokenizer = TOKENIZER_JSON .parse::() .expect("reference tokenizer should load"); diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 7eedc449dd1..ce08e225be4 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -1,9 +1,7 @@ use serde::Serialize; use crate::Error; -use crate::byte_level::ByteLevelCounter; use crate::python_json; -use crate::scanner::{SplitPattern, TiktokenCounter}; use crate::tools::format_function_definitions; use crate::types::{ ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, @@ -24,73 +22,22 @@ pub struct InputTokenCount { pub input_tokens: usize, } -enum Encoder { - HuggingFace { - tokenizer: Box, - byte_level: Option, - }, - Tiktoken(TiktokenCounter), -} - /// A loaded tokenizer plus the message accounting Python applies on top of /// it. Encoding is CPU-bound and synchronous; hosts run it off their event /// loop. pub struct TokenCounter { - encoder: Encoder, + encoder: Box, } impl TokenCounter { - /// Load a HuggingFace `tokenizer.json` document. The host reads the file. - pub fn from_json(tokenizer_json: &str) -> Result { - let tokenizer = tokenizer_json - .parse::() - .map_err(Error::Load)?; - let byte_level = ByteLevelCounter::detect(&tokenizer); - Ok(Self { - encoder: Encoder::HuggingFace { - tokenizer: Box::new(tokenizer), - byte_level, - }, - }) - } - - /// Load tiktoken's `cl100k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_cl100k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::Cl100k, rank_file) - } - - /// Load tiktoken's `o200k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_o200k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::O200k, rank_file) - } - - fn from_tiktoken_ranks(split: SplitPattern, rank_file: &str) -> Result { - Ok(Self { - encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?), - }) + pub fn new(tokenizer: impl crate::Tokenizer + 'static) -> Self { + Self { + encoder: Box::new(tokenizer), + } } pub fn count_text(&self, text: &str) -> Result { - match &self.encoder { - Encoder::Tiktoken(counter) => Ok(counter.count(text)), - Encoder::HuggingFace { - tokenizer, - byte_level, - } => { - if let Some(count) = byte_level - .as_ref() - .and_then(|counter| counter.count(tokenizer, text)) - { - return Ok(count); - } - tokenizer - .encode_fast(text, true) - .map(|encoding| encoding.len()) - .map_err(Error::Encode) - } - } + self.encoder.count_tokens(text) } /// Mirrors the host's key precedence: `messages`, then `prompt`, then diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index 6b8668fe182..b05ce007e46 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -4,8 +4,10 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { + #[error("unsupported tokenizer: {0}")] + UnsupportedTokenizer(String), #[error("failed to load tokenizer: {0}")] - Load(#[source] tokenizers::Error), + Load(#[source] Box), #[error("failed to load tokenizer: tiktoken rank file: {0}")] Ranks(String), #[error("failed to load tokenizer: Unicode character classes are unavailable")] @@ -29,7 +31,7 @@ pub enum Error { #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] JsonUtf8(#[source] FromUtf8Error), #[error("tokenization failed: {0}")] - Encode(#[source] tokenizers::Error), + Encode(#[source] Box), #[error("token counting task failed: {0}")] Task(String), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs new file mode 100644 index 00000000000..de3f86abd68 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -0,0 +1,41 @@ +use litellm_token_counter_fast::Error as BackendError; +pub use litellm_token_counter_fast::FastTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json_fast(tokenizer_json: &str) -> Result { + FastTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_cl100k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_cl100k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_o200k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_o200k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for FastTokenizer { + fn count_tokens(&self, text: &str) -> Result { + FastTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Ranks(message) => Self::Ranks(message), + BackendError::UnicodeClasses => Self::UnicodeClasses, + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs new file mode 100644 index 00000000000..fb7683b373e --- /dev/null +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -0,0 +1,27 @@ +use litellm_token_counter_huggingface::Error as BackendError; +pub use litellm_token_counter_huggingface::HuggingFaceTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json(tokenizer_json: &str) -> Result { + HuggingFaceTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for HuggingFaceTokenizer { + fn count_tokens(&self, text: &str) -> Result { + HuggingFaceTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index fa0014e2bad..446c91049de 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -4,18 +4,21 @@ #![forbid(unsafe_code)] -mod byte_level; -mod cl100k; mod counter; mod error; -mod o200k; mod python_json; -mod scanner; -mod tiktoken; +mod tokenizer; mod tools; mod types; -mod unicode_classes; + +#[cfg(feature = "fast")] +pub mod fast; +#[cfg(feature = "huggingface")] +pub mod huggingface; +#[cfg(feature = "tiktoken")] +pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; pub use error::Error; +pub use tokenizer::Tokenizer; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index 7a9e71ed587..07c1c9f5b73 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,222 +1,24 @@ -//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and -//! the merge loop that turns one regex piece into tokens. The merge order is -//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is -//! identical, but pairs are tracked in a heap so a long piece costs -//! `O(n log n)` instead of tiktoken's `O(n^2)`. +pub use litellm_token_counter_tiktoken::TiktokenTokenizer; +use litellm_token_counter_tiktoken::UnsupportedTokenizer; -use std::cmp::Reverse; -use std::collections::BinaryHeap; +use crate::{Error, TokenCounter, Tokenizer}; -use base64::Engine; -use base64::engine::general_purpose::STANDARD; -use rustc_hash::FxHashMap; - -use crate::Error; - -type Rank = u32; - -const NO_RANK: Rank = Rank::MAX; -const END: usize = usize::MAX; - -pub(super) struct MergeRanks(FxHashMap, Rank>); - -impl MergeRanks { - pub(super) fn parse(text: &str) -> Result { - let ranks = text - .lines() - .filter(|line| !line.is_empty()) - .map(parse_line) - .collect::, _>>()?; - if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { - return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); - } - Ok(Self(ranks)) - } - - fn rank(&self, bytes: &[u8]) -> Rank { - self.0.get(bytes).copied().unwrap_or(NO_RANK) - } - - /// Token count of one regex piece, as `encode_ordinary` would produce. - pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { - if piece.len() < 2 || self.0.contains_key(piece) { - return 1; - } - scratch.reset(piece.len()); - for start in 0..piece.len() - 1 { - scratch.set_rank(start, self.rank(&piece[start..start + 2])); - } - let mut parts = piece.len(); - while let Some(Reverse((rank, start))) = scratch.heap.pop() { - if scratch.next[start] == END || scratch.rank[start] != rank { - continue; - } - let merged = scratch.next[start]; - let after = scratch.next[merged]; - scratch.next[merged] = END; - scratch.next[start] = after; - parts -= 1; - if after < piece.len() { - scratch.prev[after] = start; - scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); - } else { - scratch.rank[start] = NO_RANK; - } - let before = scratch.prev[start]; - if before != END { - scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); - } - } - parts +impl TokenCounter { + pub fn from_tiktoken(encoding: &str) -> Result { + TiktokenTokenizer::from_name(encoding) + .map(Self::new) + .map_err(Error::from) } } -fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { - let (token, rank) = line - .split_once(' ') - .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; - let bytes = STANDARD - .decode(token) - .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; - let rank = rank - .parse() - .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; - Ok((bytes.into_boxed_slice(), rank)) -} - -/// Buffers reused across the pieces of one text. Parts are addressed by the -/// byte offset they start at, which also gives the leftmost-pair tie break. -#[derive(Default)] -pub(super) struct MergeScratch { - next: Vec, - prev: Vec, - rank: Vec, - heap: BinaryHeap>, -} - -impl MergeScratch { - fn reset(&mut self, len: usize) { - self.next.clear(); - self.next.extend(1..=len); - self.prev.clear(); - self.prev.push(END); - self.prev.extend(0..len - 1); - self.rank.clear(); - self.rank.resize(len, NO_RANK); - self.heap.clear(); - } - - fn end(&self, start: usize) -> usize { - self.next[start] - } - - fn set_rank(&mut self, start: usize, rank: Rank) { - self.rank[start] = rank; - if rank != NO_RANK { - self.heap.push(Reverse((rank, start))); - } +impl Tokenizer for TiktokenTokenizer { + fn count_tokens(&self, text: &str) -> Result { + Ok(TiktokenTokenizer::count_tokens(self, text)) } } -#[cfg(test)] -mod tests { - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; - - use super::*; - - fn ranks() -> MergeRanks { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" - ); - MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) - .expect("rank file parses") - } - - /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. - fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { - if piece.len() < 2 || ranks.0.contains_key(piece) { - return 1; - } - let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) - .map(|index| (index, ranks.rank(&piece[index..index + 2]))) - .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) - .collect(); - let get_rank = |parts: &[(usize, Rank)], index: usize| { - if index + 3 < parts.len() { - ranks.rank(&piece[parts[index].0..parts[index + 3].0]) - } else { - NO_RANK - } - }; - loop { - let Some(index) = parts[..parts.len() - 1] - .iter() - .enumerate() - .filter(|(_, (_, rank))| *rank != NO_RANK) - .min_by_key(|(index, (_, rank))| (*rank, *index)) - .map(|(index, _)| index) - else { - return parts.len() - 1; - }; - if index > 0 { - parts[index - 1].1 = get_rank(&parts, index - 1); - } - parts[index].1 = get_rank(&parts, index); - parts.remove(index + 1); - } - } - - #[test] - fn every_byte_is_a_token() { - let ranks = ranks(); - assert_eq!(ranks.0.len(), 100_256); - assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); - } - - #[test] - fn heap_merge_matches_tiktokens_merge_loop() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut rng = StdRng::seed_from_u64(99); - let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; - for _ in 0..20_000 { - let piece: Vec = (0..rng.gen_range(1..24)) - .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) - .collect(); - assert_eq!( - ranks.count_piece(&piece, &mut scratch), - reference_count(&ranks, &piece), - "piece {:?}", - String::from_utf8_lossy(&piece) - ); - } - } - - #[test] - fn long_repeated_runs_cost_close_to_linear() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut time = |len: usize| { - let piece = vec![b' '; len]; - let started = std::time::Instant::now(); - assert!(ranks.count_piece(&piece, &mut scratch) > 0); - started.elapsed() - }; - let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); - let large = time(1 << 18); - assert!( - large < small * 64, - "{small:?} for 2^14 bytes, {large:?} for 2^18" - ); - } - - #[test] - fn malformed_rank_files_are_rejected() { - assert!(MergeRanks::parse("IQ==").is_err()); - assert!(MergeRanks::parse("IQ== x").is_err()); - assert!(MergeRanks::parse("!!! 1").is_err()); - assert!(MergeRanks::parse("IQ== 1").is_err()); +impl From for Error { + fn from(error: UnsupportedTokenizer) -> Self { + Self::UnsupportedTokenizer(error.0) } } diff --git a/litellm-rust/crates/token-counter/src/tokenizer.rs b/litellm-rust/crates/token-counter/src/tokenizer.rs new file mode 100644 index 00000000000..88c29c672a7 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/tokenizer.rs @@ -0,0 +1,31 @@ +use crate::Error; + +pub trait Tokenizer: Send + Sync { + fn count_tokens(&self, text: &str) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{CountableRequest, TokenCounter}; + + struct Characters; + + impl Tokenizer for Characters { + fn count_tokens(&self, text: &str) -> Result { + Ok(text.chars().count()) + } + } + + #[test] + fn request_accounting_works_with_an_injected_backend() { + let counter = TokenCounter::new(Characters); + let request = + CountableRequest::parse(br#"{"messages":[{"role":"user","content":"hello"}]}"#) + .unwrap(); + assert_eq!( + counter.count_request(&request).unwrap().input_tokens, + 3 + 4 + 5 + 3 + ); + } +} diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 12c59768952..378a16d5d72 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,5 +1,6 @@ +#![cfg(feature = "huggingface")] + use rstest::rstest; -use serde::Deserialize; use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; @@ -176,146 +177,155 @@ fn loading_a_bad_tokenizer_is_a_load_error() { assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); } -/// A tiktoken encoding: its fixture directory, the vendored rank file Python -/// loads, the constructor, and the model `generate.py` counted the requests for. -#[derive(Clone, Copy)] -struct TiktokenEncoding { - fixtures: &'static str, - rank_file: &'static str, - load: fn(&str) -> Result, - model: &'static str, -} +#[cfg(feature = "fast")] +mod fast { + use super::*; + use serde::Deserialize; -const CL100K: TiktokenEncoding = TiktokenEncoding { - fixtures: "cl100k", - rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", - load: TokenCounter::from_cl100k_ranks, - model: "gpt-4", -}; + /// A tiktoken encoding: its fixture directory, the vendored rank file Python + /// loads, the constructor, and the model `generate.py` counted the requests for. + #[derive(Clone, Copy)] + struct TiktokenEncoding { + fixtures: &'static str, + rank_file: &'static str, + load: fn(&str) -> Result, + model: &'static str, + } -const O200K: TiktokenEncoding = TiktokenEncoding { - fixtures: "o200k", - rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", - load: TokenCounter::from_o200k_ranks, - model: "gpt-4o", -}; - -fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { - let path = format!( - "{}/../../../litellm/litellm_core_utils/tokenizers/{}", - env!("CARGO_MANIFEST_DIR"), - encoding.rank_file - ); - let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); - (encoding.load)(&ranks).expect("ranks load") -} - -fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { - let path = format!( - "{}/tests/fixtures/{}/{name}", - env!("CARGO_MANIFEST_DIR"), - encoding.fixtures - ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") -} - -#[derive(Deserialize)] -struct TextFixture { - text: String, - tokens: usize, -} - -#[derive(Deserialize)] -struct RequestFixture { - body: String, - input_tokens: usize, -} - -/// Reference counts come from `tiktoken.get_encoding(name)`; see -/// `tests/fixtures/generate.py`. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - assert!(fixtures.len() > 3000); - let mismatches: Vec<_> = fixtures - .iter() - .filter_map(|fixture| { - let count = counter.count_text(&fixture.text).expect("text counts"); - (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) - }) - .collect(); - assert!( - mismatches.is_empty(), - "(text, tiktoken, rust): {mismatches:?}" - ); -} - -/// Reference counts come from the proxy's admission counter -/// (`_count_input_tokens(body, model)`), so this pins the shared message, -/// tool and reply-priming accounting on the tiktoken paths as well. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - let counts: Vec = fixtures - .iter() - .map(|fixture| { - let request = CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); - let count = counter.count_request(&request).expect("fixture counts"); - assert_eq!(count.model.as_deref(), Some(encoding.model)); - assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); - count.input_tokens - }) - .collect(); - assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); -} - -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( - #[case] encoding: TiktokenEncoding, -) { - let counter = tiktoken_counter(encoding); - let count = |body: &str| { - counter - .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) - .expect("counts") - .input_tokens + const CL100K: TiktokenEncoding = TiktokenEncoding { + fixtures: "cl100k", + rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", + load: TokenCounter::from_cl100k_ranks, + model: "gpt-4", }; - let text = |text: &str| counter.count_text(text).expect("counts"); - let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); - assert_eq!(base, 3 + text("user") + text("hi") + 3); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), - base + text("al") + 1 - ); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), - base + 1 - ); -} -#[rstest] -#[case::empty("")] -#[case::not_base64("!!!! 0")] -#[case::missing_rank("YQ==")] -#[case::rank_not_a_number("YQ== x")] -#[case::single_byte_tokens_missing("YWI= 0")] -fn loading_a_bad_rank_file_is_a_load_error( - #[case] rank_file: &str, - #[values(CL100K, O200K)] encoding: TiktokenEncoding, -) { - assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + const O200K: TiktokenEncoding = TiktokenEncoding { + fixtures: "o200k", + rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", + load: TokenCounter::from_o200k_ranks, + model: "gpt-4o", + }; + + fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { + let path = format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{}", + env!("CARGO_MANIFEST_DIR"), + encoding.rank_file + ); + let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); + (encoding.load)(&ranks).expect("ranks load") + } + + fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { + let path = format!( + "{}/../token-counter-fast/tests/fixtures/{}/{name}", + env!("CARGO_MANIFEST_DIR"), + encoding.fixtures + ); + std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + } + + #[derive(Deserialize)] + struct TextFixture { + text: String, + tokens: usize, + } + + #[derive(Deserialize)] + struct RequestFixture { + body: String, + input_tokens: usize, + } + + /// Reference counts come from `tiktoken.get_encoding(name)`; see + /// `tests/fixtures/generate.py`. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + assert!(fixtures.len() > 3000); + let mismatches: Vec<_> = fixtures + .iter() + .filter_map(|fixture| { + let count = counter.count_text(&fixture.text).expect("text counts"); + (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) + }) + .collect(); + assert!( + mismatches.is_empty(), + "(text, tiktoken, rust): {mismatches:?}" + ); + } + + /// Reference counts come from the proxy's admission counter + /// (`_count_input_tokens(body, model)`), so this pins the shared message, + /// tool and reply-priming accounting on the tiktoken paths as well. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + let counts: Vec = fixtures + .iter() + .map(|fixture| { + let request = + CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); + let count = counter.count_request(&request).expect("fixture counts"); + assert_eq!(count.model.as_deref(), Some(encoding.model)); + assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); + count.input_tokens + }) + .collect(); + assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); + } + + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( + #[case] encoding: TiktokenEncoding, + ) { + let counter = tiktoken_counter(encoding); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let text = |text: &str| counter.count_text(text).expect("counts"); + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!(base, 3 + text("user") + text("hi") + 3); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), + base + text("al") + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# + ), + base + 1 + ); + } + + #[rstest] + #[case::empty("")] + #[case::not_base64("!!!! 0")] + #[case::missing_rank("YQ==")] + #[case::rank_not_a_number("YQ== x")] + #[case::single_byte_tokens_missing("YWI= 0")] + fn loading_a_bad_rank_file_is_a_load_error( + #[case] rank_file: &str, + #[values(CL100K, O200K)] encoding: TiktokenEncoding, + ) { + assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c0a06364261..7d20fc5a02f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -97,9 +97,9 @@ class ResponsesWebSocketConnection: class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... + def from_json_fast(tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_o200k_ranks(rank_file: str) -> TokenCounter: ... + def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... def gil_stats() -> dict[str, int]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d36234f56c1..e816f1ab388 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -11,14 +11,21 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt from litellm.utils import claude_json_str, huggingface_tokenizer_kind -RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] +RustTokenizer = Literal[ + "anthropic", + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", +] class RustTokenCounter(Protocol): @@ -30,10 +37,7 @@ class RustTokenCounterFactory(Protocol): def __call__(self, tokenizer_json: str) -> RustTokenCounter: raise NotImplementedError - def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: - raise NotImplementedError - - def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: + def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter: raise NotImplementedError @@ -63,9 +67,8 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count. Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace - downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust - prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in - Python.""" + downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every + message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python.""" if litellm.disable_token_counter is True: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -73,24 +76,19 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - match openai_tokenizer_encoding(model).name: - case "cl100k_base": - return "cl100k_base" - case "o200k_base": - return "o200k_base" - case _: - return None + encoding: Final = openai_tokenizer_encoding(model).name + if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}: + return cast(RustTokenizer, encoding) + return None -@lru_cache(maxsize=4) +@lru_cache(maxsize=8) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: match tokenizer: case "anthropic": return factory(claude_json_str) - case "cl100k_base": - return factory.from_cl100k_ranks(cl100k_base_rank_file()) - case "o200k_base": - return factory.from_o200k_ranks(o200k_base_rank_file()) + case _: + return factory.from_tiktoken(tokenizer) async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..c785a7b657a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -256,7 +256,7 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] @@ -264,11 +264,8 @@ class _RecordingFactory: def __call__(self, tokenizer_json: str) -> _RecordingCounter: return _RecordingCounter(self, "anthropic") - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "o200k_base") + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter: + return _RecordingCounter(self, encoding) class _DecliningCounter: @@ -280,10 +277,7 @@ class _DecliningFactory: def __call__(self, tokenizer_json: str) -> _DecliningCounter: return _DecliningCounter() - def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: - return _DecliningCounter() - - def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter: return _DecliningCounter() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index 71aa79cc4bb..b95e6a62e9c 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -8,7 +8,6 @@ cases need the extension and are skipped when it is not built. from __future__ import annotations import json -from types import MappingProxyType from typing import Final import pytest @@ -27,7 +26,6 @@ MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") -RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() @@ -55,24 +53,20 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.rank_files: list[str] = [] + self.encodings: list[str] = [] def __call__(self, tokenizer_json: str) -> _RecordingCounter: counter = _RecordingCounter(tokenizer_json) self.counters.append(counter) return counter - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("o200k_base") + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter: + self.encodings.append(encoding) + return self(encoding) class _RaisingCounter: @@ -92,10 +86,7 @@ class _RaisingFactory: def __call__(self, tokenizer_json: str) -> _RaisingCounter: return _RaisingCounter(self.error) - def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: - return _RaisingCounter(self.error) - - def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -140,7 +131,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) @@ -149,9 +140,7 @@ async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokeni second: Final = await bridge.count_input_tokens(BODY, tokenizer) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert len(factory.rank_files) == 1 - assert factory.rank_files[0].startswith("IQ== 0\n") - assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] + assert factory.encodings == [tokenizer] assert factory.counters[0].tokenizer_json == tokenizer assert factory.counters[0].bodies == [BODY, BODY] @@ -235,13 +224,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: ("model", "python_encoding"), (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")), ) -def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have( - monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str +def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust( + monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer ) -> None: monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model}) assert openai_tokenizer_encoding(model).name == python_encoding - assert bridge.rust_tokenizer(model) is None + assert bridge.rust_tokenizer(model) == python_encoding def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: From 9b2b3d0b90346c0da91a8abf4b52be7561e1dcfe Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:21:52 +0000 Subject: [PATCH 102/127] refactor(rust): align tokenizer docs and lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/token-counter/README.md | 2 +- litellm/rust_bridge/token_counter.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index c8ed0737a7e..effa5324822 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -2,7 +2,7 @@ `Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface -The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation +The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast` uses this implementation The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index e816f1ab388..d232db19a4a 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -77,7 +77,14 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: if kind is not None or uses_legacy_message_accounting(model): return None encoding: Final = openai_tokenizer_encoding(model).name - if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}: + if encoding in ( + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + ): return cast(RustTokenizer, encoding) return None From a619b765fc06b9ba9c7124debcf75f0c79033f58 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:25:19 +0000 Subject: [PATCH 103/127] refactor(rust): simplify tokenizer bridge features Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/Cargo.toml | 8 ++-- .../crates/python-bridge/src/token_counter.rs | 42 ++++++------------- 2 files changed, 17 insertions(+), 33 deletions(-) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 2c868ebc769..67f662408c4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,13 +10,13 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"] +default = ["abi3", "huggingface", "tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] -token-counter-fast = ["litellm-token-counter/fast"] -token-counter-huggingface = ["litellm-token-counter/huggingface"] -token-counter-tiktoken = ["litellm-token-counter/tiktoken"] +fast = ["litellm-token-counter/fast"] +huggingface = ["litellm-token-counter/huggingface"] +tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index a7c6bce4d0b..12749486679 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -1,17 +1,9 @@ use std::sync::Arc; -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use std::{num::NonZero, thread::available_parallelism}; -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use litellm_host_python::release_gil; use litellm_host_python::run_async; use litellm_token_counter::{ @@ -41,45 +33,45 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "token-counter-huggingface")] + #[cfg(feature = "huggingface")] { Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) } - #[cfg(not(feature = "token-counter-huggingface"))] + #[cfg(not(feature = "huggingface"))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-huggingface feature", + "tokenizer backend requires the huggingface feature", )) } } #[staticmethod] fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "token-counter-fast")] + #[cfg(feature = "fast")] { Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) } - #[cfg(not(feature = "token-counter-fast"))] + #[cfg(not(feature = "fast"))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-fast feature", + "tokenizer backend requires the fast feature", )) } } #[staticmethod] fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { - #[cfg(feature = "token-counter-tiktoken")] + #[cfg(feature = "tiktoken")] { Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) } - #[cfg(not(feature = "token-counter-tiktoken"))] + #[cfg(not(feature = "tiktoken"))] { let _ = (py, encoding); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the token-counter-tiktoken feature", + "tokenizer backend requires the tiktoken feature", )) } } @@ -105,11 +97,7 @@ impl TokenCounter { } impl TokenCounter { - #[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" - ))] + #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] fn load( py: Python<'_>, load: impl FnOnce() -> Result + Send, @@ -122,11 +110,7 @@ impl TokenCounter { } } -#[cfg(any( - feature = "token-counter-fast", - feature = "token-counter-huggingface", - feature = "token-counter-tiktoken" -))] +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] fn encode_parallelism() -> usize { available_parallelism().map_or(1, NonZero::get) } From be7b7d2ade72d0b158b13c8d1fed426d6a24cd92 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:30:55 +0000 Subject: [PATCH 104/127] chore(prices): sync OpenRouter prices: 1 model openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d25cf0c78df..850e0e5ecb4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d25cf0c78df..850e0e5ecb4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 831810248ae95a407d8cd6c3f1c94be695d894e0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:49:39 +0000 Subject: [PATCH 105/127] refactor(rust): keep the Python token counter on the fast backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- .../crates/python-bridge/src/token_counter.rs | 31 +++++++++--- litellm-rust/crates/token-counter/README.md | 2 +- litellm/rust_bridge/_native.pyi | 4 +- litellm/rust_bridge/token_counter.py | 47 +++++++++---------- .../spend_tracking/test_budget_reservation.py | 14 ++++-- .../rust_bridge/test_token_counter.py | 33 ++++++++----- 7 files changed, 83 insertions(+), 50 deletions(-) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 67f662408c4..a76b069935f 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "huggingface", "tiktoken"] +default = ["abi3", "fast"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 12749486679..244401e6696 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -33,28 +33,47 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "huggingface")] + #[cfg(feature = "fast")] + { + Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + } + #[cfg(all(not(feature = "fast"), feature = "huggingface"))] { Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) } - #[cfg(not(feature = "huggingface"))] + #[cfg(not(any(feature = "fast", feature = "huggingface")))] { let _ = (py, tokenizer_json); Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", + "tokenizer backend requires the fast or huggingface feature", )) } } #[staticmethod] - fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { + fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { #[cfg(feature = "fast")] { - Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) } #[cfg(not(feature = "fast"))] { - let _ = (py, tokenizer_json); + let _ = (py, rank_file); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the fast feature", + )) + } + } + + #[staticmethod] + fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { + #[cfg(feature = "fast")] + { + Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) + } + #[cfg(not(feature = "fast"))] + { + let _ = (py, rank_file); Err(RustBridgeDeclined::new_err( "tokenizer backend requires the fast feature", )) diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index effa5324822..3ad21a00697 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` -The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7d20fc5a02f..05a6df6d5af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -97,7 +97,9 @@ class ResponsesWebSocketConnection: class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_json_fast(tokenizer_json: str) -> TokenCounter: ... + def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... + @staticmethod + def from_o200k_ranks(rank_file: str) -> TokenCounter: ... @staticmethod def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d232db19a4a..d36234f56c1 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -11,21 +11,14 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt from litellm.utils import claude_json_str, huggingface_tokenizer_kind -RustTokenizer = Literal[ - "anthropic", - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", -] +RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] class RustTokenCounter(Protocol): @@ -37,7 +30,10 @@ class RustTokenCounterFactory(Protocol): def __call__(self, tokenizer_json: str) -> RustTokenCounter: raise NotImplementedError - def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter: + def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: + raise NotImplementedError + + def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: raise NotImplementedError @@ -67,8 +63,9 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count. Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace - downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every - message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python.""" + downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust + prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in + Python.""" if litellm.disable_token_counter is True: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -76,26 +73,24 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - encoding: Final = openai_tokenizer_encoding(model).name - if encoding in ( - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", - ): - return cast(RustTokenizer, encoding) - return None + match openai_tokenizer_encoding(model).name: + case "cl100k_base": + return "cl100k_base" + case "o200k_base": + return "o200k_base" + case _: + return None -@lru_cache(maxsize=8) +@lru_cache(maxsize=4) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: match tokenizer: case "anthropic": return factory(claude_json_str) - case _: - return factory.from_tiktoken(tokenizer) + case "cl100k_base": + return factory.from_cl100k_ranks(cl100k_base_rank_file()) + case "o200k_base": + return factory.from_o200k_ranks(o200k_base_rank_file()) async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index c785a7b657a..0e0025f5194 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -256,7 +256,7 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class.""" + """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] @@ -264,8 +264,11 @@ class _RecordingFactory: def __call__(self, tokenizer_json: str) -> _RecordingCounter: return _RecordingCounter(self, "anthropic") - def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter: - return _RecordingCounter(self, encoding) + def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: + return _RecordingCounter(self, "cl100k_base") + + def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: + return _RecordingCounter(self, "o200k_base") class _DecliningCounter: @@ -277,7 +280,10 @@ class _DecliningFactory: def __call__(self, tokenizer_json: str) -> _DecliningCounter: return _DecliningCounter() - def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter: + def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: + return _DecliningCounter() + + def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: return _DecliningCounter() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index b95e6a62e9c..71aa79cc4bb 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -8,6 +8,7 @@ cases need the extension and are skipped when it is not built. from __future__ import annotations import json +from types import MappingProxyType from typing import Final import pytest @@ -26,6 +27,7 @@ MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") +RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() @@ -53,20 +55,24 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class.""" + """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.encodings: list[str] = [] + self.rank_files: list[str] = [] def __call__(self, tokenizer_json: str) -> _RecordingCounter: counter = _RecordingCounter(tokenizer_json) self.counters.append(counter) return counter - def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter: - self.encodings.append(encoding) - return self(encoding) + def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: + self.rank_files.append(rank_file) + return self("cl100k_base") + + def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: + self.rank_files.append(rank_file) + return self("o200k_base") class _RaisingCounter: @@ -86,7 +92,10 @@ class _RaisingFactory: def __call__(self, tokenizer_json: str) -> _RaisingCounter: return _RaisingCounter(self.error) - def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter: + def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: + return _RaisingCounter(self.error) + + def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -131,7 +140,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) @@ -140,7 +149,9 @@ async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridg second: Final = await bridge.count_input_tokens(BODY, tokenizer) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert factory.encodings == [tokenizer] + assert len(factory.rank_files) == 1 + assert factory.rank_files[0].startswith("IQ== 0\n") + assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] assert factory.counters[0].tokenizer_json == tokenizer assert factory.counters[0].bodies == [BODY, BODY] @@ -224,13 +235,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: ("model", "python_encoding"), (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")), ) -def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust( - monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer +def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have( + monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str ) -> None: monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model}) assert openai_tokenizer_encoding(model).name == python_encoding - assert bridge.rust_tokenizer(model) == python_encoding + assert bridge.rust_tokenizer(model) is None def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: From 56c5d31e73d7b7ecc5f6a91d2b49ed355fb8772e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:52:48 +0000 Subject: [PATCH 106/127] test(rust): run fast token counter parity tests by default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + .../tests/fixtures/generate.py | 4 +- litellm-rust/crates/token-counter/Cargo.toml | 3 +- litellm-rust/crates/token-counter/README.md | 2 +- .../token-counter/tests/token_counter.rs | 191 +++++++++++------- 6 files changed, 129 insertions(+), 73 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 62bde5806a4..af3a31ddbfa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2257,6 +2257,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.19", + "tokenizers", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 109a2edaf4d..4250fa81d26 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -48,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } diff --git a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py index 1bfbdf00218..2ca853cbd62 100644 --- a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py +++ b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py @@ -2,8 +2,8 @@ Run from the repository root with the project environment, once per encoding: - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py cl100k_base - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py o200k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py cl100k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py o200k_base `/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens` counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`, diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index 59e4f9a6a1d..67e5bd60537 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -default = ["huggingface", "tiktoken"] +default = ["fast", "huggingface", "tiktoken"] fast = ["dep:litellm-token-counter-fast"] huggingface = ["dep:litellm-token-counter-huggingface"] tiktoken = ["dep:litellm-token-counter-tiktoken"] @@ -25,6 +25,7 @@ thiserror.workspace = true criterion.workspace = true rand.workspace = true rstest.workspace = true +tokenizers.workspace = true [[bench]] name = "token_counter" diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index 3ad21a00697..a6b2b50aac0 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` -The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +All three backends are enabled by default. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 378a16d5d72..ebbe2f00ffc 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,4 +1,4 @@ -#![cfg(feature = "huggingface")] +#![cfg(any(feature = "fast", feature = "huggingface"))] use rstest::rstest; @@ -6,13 +6,15 @@ use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCount /// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` /// so this test also guards Python parity. -fn counter() -> TokenCounter { +type JsonLoader = fn(&str) -> Result; + +fn counter(load: JsonLoader) -> TokenCounter { let path = concat!( env!("CARGO_MANIFEST_DIR"), "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" ); let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); - TokenCounter::from_json(&json).expect("anthropic tokenizer loads") + load(&json).expect("anthropic tokenizer loads") } const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; @@ -63,21 +65,15 @@ const EMBEDDINGS_TOKEN_IDS: &str = const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; -/// Expected counts are pinned from -/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. -#[rstest] -#[case::text_only(SIMPLE, 14)] -#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] -#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] -#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] -#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] -#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] -#[case::responses_input_items(RESPONSES_INPUT, 62)] -#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] -#[case::rerank_query_and_documents(RERANK, 41)] -fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { +fn assert_count_request_matches_python_token_counter( + load: JsonLoader, + body: &str, + expected: usize, +) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!( count, InputTokenCount { @@ -87,63 +83,24 @@ fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expect ); } -#[rstest] -#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] -#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] -#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] -#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] -fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { +fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!(count.input_tokens, expected); } -#[rstest] -#[case::not_json(b"not json" as &[u8])] -#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] -#[case::message_with_tool_calls( - br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# -)] -#[case::dict_content( - br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# -)] -#[case::float_enum( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# -)] -#[case::anthropic_tool_choice_without_function( - br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# -)] -fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { - assert!(matches!( - CountableRequest::parse(body), - Err(Error::RequestParse(_)) - )); -} - -#[rstest] -#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] -#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] -#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] -#[case::image_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# -)] -#[case::tool_result_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# -)] -#[case::array_without_items( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# -)] -fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { +fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { let request = CountableRequest::parse(body).expect("shape parses"); assert!(matches!( - counter().count_request(&request), + counter(load).count_request(&request), Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) )); } -#[test] -fn tool_choice_and_system_discount_change_the_count() { - let counter = counter(); +fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { + let counter = counter(load); let count = |body: &str| { counter .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) @@ -172,9 +129,104 @@ fn tool_choice_and_system_discount_change_the_count() { ); } -#[test] -fn loading_a_bad_tokenizer_is_a_load_error() { - assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); +fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { + assert!(matches!(load("{}"), Err(Error::Load(_)))); +} + +macro_rules! json_backend_tests { + ($loader:path) => { + #[rstest] + #[case::text_only(super::SIMPLE, 14)] + #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] + #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] + #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] + #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] + #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] + #[case::responses_input_items(super::RESPONSES_INPUT, 62)] + #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] + #[case::rerank_query_and_documents(super::RERANK, 41)] + fn count_request_matches_python_token_counter( + #[case] body: &str, + #[case] expected: usize, + ) { + super::assert_count_request_matches_python_token_counter($loader, body, expected); + } + + #[rstest] + #[case::null_messages_win_over_prompt( + r#"{"model":"m","messages":null,"prompt":"ignored"}"#, + 3 + )] + #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] + #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] + #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] + fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + super::assert_key_presence_follows_python($loader, body, expected); + } + + #[rstest] + #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] + #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] + #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] + #[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# + )] + #[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# + )] + #[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# + )] + fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); + } + + #[test] + fn tool_choice_and_system_discount_change_the_count() { + super::assert_tool_choice_and_system_discount_change_the_count($loader); + } + + #[test] + fn loading_a_bad_tokenizer_is_a_load_error() { + super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); + } + }; +} + +#[rstest] +#[case::not_json(b"not json" as &[u8])] +#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] +#[case::message_with_tool_calls( + br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# +)] +#[case::dict_content( + br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# +)] +#[case::float_enum( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# +)] +#[case::anthropic_tool_choice_without_function( + br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# +)] +fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { + assert!(matches!( + CountableRequest::parse(body), + Err(Error::RequestParse(_)) + )); +} + +#[cfg(feature = "fast")] +mod fast_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json_fast); +} + +#[cfg(feature = "huggingface")] +mod huggingface_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json); } #[cfg(feature = "fast")] @@ -222,7 +274,8 @@ mod fast { env!("CARGO_MANIFEST_DIR"), encoding.fixtures ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + std::fs::read_to_string(&path) + .expect("fixture generated by token-counter-fast/tests/fixtures/generate.py") } #[derive(Deserialize)] @@ -238,7 +291,7 @@ mod fast { } /// Reference counts come from `tiktoken.get_encoding(name)`; see - /// `tests/fixtures/generate.py`. + /// `token-counter-fast/tests/fixtures/generate.py`. #[rstest] #[case::cl100k(CL100K)] #[case::o200k(O200K)] From 10715b1b4607ff2260f5ad22d7bc94fbd43d48f3 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:00:46 +0000 Subject: [PATCH 107/127] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 26 +++++++++---------- model_prices_and_context_window.json | 26 +++++++++---------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 850e0e5ecb4..dc570f14a2d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2668e-07, + "input_cost_per_token": 5.2536e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7556e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.6716e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7556e-08, - "input_cost_per_token": 5.2668e-07, + "cache_read_input_token_cost": 1.6716e-08, + "input_cost_per_token": 5.2536e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 850e0e5ecb4..dc570f14a2d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2668e-07, + "input_cost_per_token": 5.2536e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7556e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2668e-7,"output_cost_per_token":0.00000158004,"cache_read_input_token_cost":1.6968e-8}, + "cache_read_input_token_cost": 1.6716e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7556e-08, - "input_cost_per_token": 5.2668e-07, + "cache_read_input_token_cost": 1.6716e-08, + "input_cost_per_token": 5.2536e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 384000, - "max_tokens": 384000, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.58004e-06, + "output_cost_per_token": 1.57608e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71593,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 661da87c910deebd487953d4f2c38c0211d88398 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 15:13:59 -0700 Subject: [PATCH 108/127] fix(rust): validate tokenizer ranks and cover backend features --- .github/workflows/test-rust.yml | 7 + .../crates/token-counter-fast/Cargo.toml | 2 +- .../crates/token-counter-fast/src/tiktoken.rs | 18 + .../token-counter-huggingface/Cargo.toml | 2 +- .../token-counter-huggingface/src/error.rs | 9 + .../token-counter-huggingface/src/lib.rs | 10 +- .../token-counter-tiktoken/src/error.rs | 5 + .../crates/token-counter-tiktoken/src/lib.rs | 54 ++- .../token-counter/tests/token_counter.rs | 423 ++++++++++-------- 9 files changed, 326 insertions(+), 204 deletions(-) create mode 100644 litellm-rust/crates/token-counter-huggingface/src/error.rs create mode 100644 litellm-rust/crates/token-counter-tiktoken/src/error.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 551f783d4f9..70d828b5c8b 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -120,6 +120,13 @@ jobs: - run: cargo test --workspace --doc --locked + - name: Test token counter feature combinations + run: | + for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do + cargo test -p litellm-token-counter --locked --no-default-features --features "$features" + cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}" + done + rust-wheel: runs-on: ubuntu-latest timeout-minutes: 30 diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml index 4950277f5b0..5127dc17f22 100644 --- a/litellm-rust/crates/token-counter-fast/Cargo.toml +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -9,7 +9,7 @@ repository.workspace = true base64.workspace = true rustc-hash = "2.1.3" thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +tokenizers.workspace = true unicode-normalization-alignments = "0.1.12" [dev-dependencies] diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs index 7a9e71ed587..16172b7a688 100644 --- a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -81,6 +81,9 @@ fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { let rank = rank .parse() .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; + if rank == NO_RANK { + return Err(Error::Ranks(format!("rank {rank} is reserved"))); + } Ok((bytes.into_boxed_slice(), rank)) } @@ -212,6 +215,21 @@ mod tests { ); } + #[test] + fn reserved_merge_rank_is_rejected() { + let bytes = (0..=u8::MAX) + .map(|byte| format!("{} {byte}\n", STANDARD.encode([byte]))) + .collect::(); + let rank_file = format!("{bytes}{} {NO_RANK}\n", STANDARD.encode(b"ab")); + assert!(matches!( + MergeRanks::parse(&rank_file), + Err(Error::Ranks(_)) + )); + let valid_rank_file = format!("{bytes}{} {}\n", STANDARD.encode(b"ab"), NO_RANK - 1); + let ranks = MergeRanks::parse(&valid_rank_file).unwrap(); + assert_eq!(ranks.count_piece(b"aab", &mut MergeScratch::default()), 2); + } + #[test] fn malformed_rank_files_are_rejected() { assert!(MergeRanks::parse("IQ==").is_err()); diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml index ad15d9f19d0..6d8cb85e524 100644 --- a/litellm-rust/crates/token-counter-huggingface/Cargo.toml +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -7,4 +7,4 @@ repository.workspace = true [dependencies] thiserror.workspace = true -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +tokenizers.workspace = true diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs new file mode 100644 index 00000000000..adc4551886f --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -0,0 +1,9 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 184956b5f39..8e05c2cca46 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -1,14 +1,8 @@ #![forbid(unsafe_code)] -use thiserror::Error as ThisError; +mod error; -#[derive(Debug, ThisError)] -pub enum Error { - #[error("failed to load tokenizer: {0}")] - Load(#[source] tokenizers::Error), - #[error("tokenization failed: {0}")] - Encode(#[source] tokenizers::Error), -} +pub use error::Error; pub struct HuggingFaceTokenizer(Box); diff --git a/litellm-rust/crates/token-counter-tiktoken/src/error.rs b/litellm-rust/crates/token-counter-tiktoken/src/error.rs new file mode 100644 index 00000000000..e28cbbfb620 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/error.rs @@ -0,0 +1,5 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +#[error("unsupported tokenizer: {0}")] +pub struct UnsupportedTokenizer(pub String); diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index b35c4876d00..ecdb3946eee 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -1,10 +1,8 @@ #![forbid(unsafe_code)] -use thiserror::Error as ThisError; +mod error; -#[derive(Debug, ThisError)] -#[error("unsupported tokenizer: {0}")] -pub struct UnsupportedTokenizer(pub String); +pub use error::UnsupportedTokenizer; pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); @@ -32,23 +30,41 @@ mod tests { use super::*; #[test] - fn special_tokens_are_counted_as_ordinary_text() { - let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap(); - assert!(counter.count_tokens("<|endoftext|>") > 1); + fn named_encodings_match_their_reference_counts() { + let encodings = [ + ("cl100k_base", tiktoken_rs::cl100k_base_singleton()), + ("o200k_base", tiktoken_rs::o200k_base_singleton()), + ("o200k_harmony", tiktoken_rs::o200k_harmony_singleton()), + ("p50k_base", tiktoken_rs::p50k_base_singleton()), + ("p50k_edit", tiktoken_rs::p50k_edit_singleton()), + ("r50k_base", tiktoken_rs::r50k_base_singleton()), + ("gpt2", tiktoken_rs::r50k_base_singleton()), + ]; + let texts = [ + "", + "Hello, how are you today?", + "é e\u{301} 漢字 ع ३ 🙂 ＡfiⅣ", + " def function():\n return 123456789\r\n", + "<|endoftext|><|fim_prefix|><|start|>assistant<|message|>", + ]; + for (name, reference) in encodings { + let counter = TiktokenTokenizer::from_name(name).unwrap(); + for text in texts { + assert_eq!( + counter.count_tokens(text), + reference.encode_ordinary(text).len(), + "{name}: {text:?}", + ); + } + } } #[test] - fn all_python_tiktoken_encodings_are_available() { - for name in [ - "cl100k_base", - "o200k_base", - "o200k_harmony", - "p50k_base", - "p50k_edit", - "r50k_base", - "gpt2", - ] { - assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}"); - } + fn unsupported_encoding_preserves_its_name() { + let Err(UnsupportedTokenizer(name)) = TiktokenTokenizer::from_name("unknown-encoding") + else { + panic!("unknown encoding must be rejected"); + }; + assert_eq!(name, "unknown-encoding"); } } diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index ebbe2f00ffc..542bd4a1fc4 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,25 +1,30 @@ -#![cfg(any(feature = "fast", feature = "huggingface"))] - use rstest::rstest; -use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; +#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] +use litellm_token_counter::TokenCounter; +use litellm_token_counter::{CountableRequest, Error}; -/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` -/// so this test also guards Python parity. -type JsonLoader = fn(&str) -> Result; +#[cfg(any(feature = "fast", feature = "huggingface"))] +mod json { + use super::*; + use litellm_token_counter::InputTokenCount; -fn counter(load: JsonLoader) -> TokenCounter { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" - ); - let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); - load(&json).expect("anthropic tokenizer loads") -} + /// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` + /// so this test also guards Python parity. + type JsonLoader = fn(&str) -> Result; -const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + fn counter(load: JsonLoader) -> TokenCounter { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); + load(&json).expect("anthropic tokenizer loads") + } -const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ + const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + + const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ {"role":"system","content":"You are a terse assistant."}, {"role":"user","name":"alice","content":[ {"type":"text","text":"Summarise this paragraph about ships and harbours."}, @@ -28,7 +33,7 @@ const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ {"type":"tool_reference","tool_name":"get_weather"}]}, {"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#; -const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], + const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], "tools":[ {"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{ "type":"object", @@ -43,154 +48,188 @@ const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":" {"type":"function","function":{"name":"noop"}}], "tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#; -const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", + const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", "messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}], "tools":[{"name":"get_weather","description":"Get weather","input_schema":{ "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}], "tool_choice":"none"}"#; -const COMPLETIONS_PROMPT: &str = - r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; + const COMPLETIONS_PROMPT: &str = + r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; -const COMPLETIONS_PROMPT_LIST: &str = - r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; + const COMPLETIONS_PROMPT_LIST: &str = + r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; -const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ + const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]}, {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#; -const EMBEDDINGS_TOKEN_IDS: &str = - r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; + const EMBEDDINGS_TOKEN_IDS: &str = + r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; -const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", + const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; -fn assert_count_request_matches_python_token_counter( - load: JsonLoader, - body: &str, - expected: usize, -) { - let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter(load) - .count_request(&request) - .expect("fixture counts"); - assert_eq!( - count, - InputTokenCount { - model: Some("claude-sonnet-4-5".to_string()), - input_tokens: expected, - } - ); -} + fn assert_count_request_matches_python_token_counter( + load: JsonLoader, + body: &str, + expected: usize, + ) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); + assert_eq!( + count, + InputTokenCount { + model: Some("claude-sonnet-4-5".to_string()), + input_tokens: expected, + } + ); + } -fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { - let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter(load) - .count_request(&request) - .expect("fixture counts"); - assert_eq!(count.input_tokens, expected); -} + fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); + assert_eq!(count.input_tokens, expected); + } -fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { - let request = CountableRequest::parse(body).expect("shape parses"); - assert!(matches!( - counter(load).count_request(&request), - Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) - )); -} + fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { + let request = CountableRequest::parse(body).expect("shape parses"); + assert!(matches!( + counter(load).count_request(&request), + Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + )); + } -fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { - let counter = counter(load); - let count = |body: &str| { - counter - .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) - .expect("counts") - .input_tokens - }; - let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), - base + 1 - ); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#), - base - ); - let with_tools = count( - r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, - ); - let with_tools_and_system = count( - r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, - ); - assert_eq!(with_tools - with_tools_and_system, 4); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), - base - ); -} + fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { + let counter = counter(load); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# + ), + base + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"# + ), + base + ); + let with_tools = count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + let with_tools_and_system = count( + r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + assert_eq!(with_tools - with_tools_and_system, 4); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), + base + ); + } -fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { - assert!(matches!(load("{}"), Err(Error::Load(_)))); -} + fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { + assert!(matches!(load("{}"), Err(Error::Load(_)))); + } -macro_rules! json_backend_tests { - ($loader:path) => { - #[rstest] - #[case::text_only(super::SIMPLE, 14)] - #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] - #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] - #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] - #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] - #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] - #[case::responses_input_items(super::RESPONSES_INPUT, 62)] - #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] - #[case::rerank_query_and_documents(super::RERANK, 41)] - fn count_request_matches_python_token_counter( - #[case] body: &str, - #[case] expected: usize, - ) { - super::assert_count_request_matches_python_token_counter($loader, body, expected); - } + macro_rules! json_backend_tests { + ($loader:path) => { + #[rstest] + #[case::text_only(super::SIMPLE, 14)] + #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] + #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] + #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] + #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] + #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] + #[case::responses_input_items(super::RESPONSES_INPUT, 62)] + #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] + #[case::rerank_query_and_documents(super::RERANK, 41)] + fn count_request_matches_python_token_counter( + #[case] body: &str, + #[case] expected: usize, + ) { + super::assert_count_request_matches_python_token_counter($loader, body, expected); + } - #[rstest] - #[case::null_messages_win_over_prompt( - r#"{"model":"m","messages":null,"prompt":"ignored"}"#, - 3 - )] - #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] - #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] - #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] - fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { - super::assert_key_presence_follows_python($loader, body, expected); - } + #[rstest] + #[case::null_messages_win_over_prompt( + r#"{"model":"m","messages":null,"prompt":"ignored"}"#, + 3 + )] + #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] + #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] + #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] + fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + super::assert_key_presence_follows_python($loader, body, expected); + } - #[rstest] - #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] - #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] - #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] - #[case::image_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# - )] - #[case::tool_result_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# - )] - #[case::array_without_items( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# - )] - fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { - super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); - } + #[rstest] + #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] + #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] + #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] + #[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# + )] + #[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# + )] + #[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# + )] + fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); + } - #[test] - fn tool_choice_and_system_discount_change_the_count() { - super::assert_tool_choice_and_system_discount_change_the_count($loader); - } + #[test] + fn tool_choice_and_system_discount_change_the_count() { + super::assert_tool_choice_and_system_discount_change_the_count($loader); + } - #[test] - fn loading_a_bad_tokenizer_is_a_load_error() { - super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); - } - }; + #[test] + fn encoding_errors_preserve_the_backend_source() { + use std::error::Error as _; + + let tokenizer = tokenizers::Tokenizer::new( + tokenizers::models::wordpiece::WordPiece::default(), + ); + let expected = tokenizer.encode_fast("hello", true).unwrap_err(); + let counter = $loader(&tokenizer.to_string(false).unwrap()).unwrap(); + let request = CountableRequest::parse(br#"{"prompt":"hello"}"#).unwrap(); + let error = counter.count_request(&request).unwrap_err(); + assert!(matches!(error, Error::Encode(_))); + assert_eq!(error.source().unwrap().to_string(), expected.to_string()); + } + + #[test] + fn loading_a_bad_tokenizer_is_a_load_error() { + super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); + } + }; + } + + #[cfg(feature = "fast")] + mod fast_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json_fast); + } + + #[cfg(feature = "huggingface")] + mod huggingface_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json); + } } #[rstest] @@ -215,22 +254,8 @@ fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { )); } -#[cfg(feature = "fast")] -mod fast_json { - use super::*; - - json_backend_tests!(TokenCounter::from_json_fast); -} - -#[cfg(feature = "huggingface")] -mod huggingface_json { - use super::*; - - json_backend_tests!(TokenCounter::from_json); -} - -#[cfg(feature = "fast")] -mod fast { +#[cfg(any(feature = "fast", feature = "tiktoken"))] +mod tiktoken { use super::*; use serde::Deserialize; @@ -239,33 +264,65 @@ mod fast { #[derive(Clone, Copy)] struct TiktokenEncoding { fixtures: &'static str, - rank_file: &'static str, + source: TokenizerSource, load: fn(&str) -> Result, model: &'static str, } + #[cfg(feature = "fast")] const CL100K: TiktokenEncoding = TiktokenEncoding { fixtures: "cl100k", - rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", + source: TokenizerSource::RankFile("9b5ad71b2ce5302211f9c61530b329a4922fc6a4"), load: TokenCounter::from_cl100k_ranks, model: "gpt-4", }; + #[cfg(feature = "fast")] const O200K: TiktokenEncoding = TiktokenEncoding { fixtures: "o200k", - rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", + source: TokenizerSource::RankFile("fb374d419588a4632f3f557e76b4b70aebbca790"), load: TokenCounter::from_o200k_ranks, model: "gpt-4o", }; + #[cfg(feature = "tiktoken")] + const TIKTOKEN_CL100K: TiktokenEncoding = TiktokenEncoding { + fixtures: "cl100k", + source: TokenizerSource::Name("cl100k_base"), + load: TokenCounter::from_tiktoken, + model: "gpt-4", + }; + + #[cfg(feature = "tiktoken")] + const TIKTOKEN_O200K: TiktokenEncoding = TiktokenEncoding { + fixtures: "o200k", + source: TokenizerSource::Name("o200k_base"), + load: TokenCounter::from_tiktoken, + model: "gpt-4o", + }; + + #[derive(Clone, Copy)] + enum TokenizerSource { + #[cfg(feature = "fast")] + RankFile(&'static str), + #[cfg(feature = "tiktoken")] + Name(&'static str), + } + fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { - let path = format!( - "{}/../../../litellm/litellm_core_utils/tokenizers/{}", - env!("CARGO_MANIFEST_DIR"), - encoding.rank_file - ); - let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); - (encoding.load)(&ranks).expect("ranks load") + match encoding.source { + #[cfg(feature = "tiktoken")] + TokenizerSource::Name(name) => (encoding.load)(name).expect("encoding loads"), + #[cfg(feature = "fast")] + TokenizerSource::RankFile(file) => { + let path = format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{file}", + env!("CARGO_MANIFEST_DIR"), + ); + let ranks = std::fs::read_to_string(path).expect("rank file is in the repo"); + (encoding.load)(&ranks).expect("ranks load") + } + } } fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { @@ -293,8 +350,10 @@ mod fast { /// Reference counts come from `tiktoken.get_encoding(name)`; see /// `token-counter-fast/tests/fixtures/generate.py`. #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { let counter = tiktoken_counter(encoding); let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") @@ -319,8 +378,10 @@ mod fast { /// (`_count_input_tokens(body, model)`), so this pins the shared message, /// tool and reply-priming accounting on the tiktoken paths as well. #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { let counter = tiktoken_counter(encoding); let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") @@ -342,8 +403,10 @@ mod fast { } #[rstest] - #[case::cl100k(CL100K)] - #[case::o200k(O200K)] + #[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))] + #[cfg_attr(feature = "fast", case::fast_o200k(O200K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))] + #[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_O200K))] fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( #[case] encoding: TiktokenEncoding, ) { @@ -369,6 +432,7 @@ mod fast { ); } + #[cfg(feature = "fast")] #[rstest] #[case::empty("")] #[case::not_base64("!!!! 0")] @@ -382,3 +446,12 @@ mod fast { assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); } } + +#[cfg(feature = "tiktoken")] +#[test] +fn unsupported_encoding_reaches_the_counter_caller() { + assert!(matches!( + TokenCounter::from_tiktoken("unknown-encoding"), + Err(Error::UnsupportedTokenizer(name)) if name == "unknown-encoding" + )); +} From 668a89fbfd6273c9ba92da38fd08c34f287d2109 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:30:41 +0000 Subject: [PATCH 109/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~deepseek/deepseek-pro-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro-0813: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost, off_peak_pricing --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dc570f14a2d..39c418de524 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2536e-07, + "input_cost_per_token": 5.2404e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6716e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 1.7468e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6716e-08, - "input_cost_per_token": 5.2536e-07, + "cache_read_input_token_cost": 1.7468e-08, + "input_cost_per_token": 5.2404e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dc570f14a2d..39c418de524 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2536e-07, + "input_cost_per_token": 5.2404e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.6716e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2536e-7,"output_cost_per_token":0.00000157608,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 1.7468e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.6716e-08, - "input_cost_per_token": 5.2536e-07, + "cache_read_input_token_cost": 1.7468e-08, + "input_cost_per_token": 5.2404e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57608e-06, + "output_cost_per_token": 1.57212e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 3fd2dd635b119ddd3afd1771c1cda217d71df33c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 15:45:38 -0700 Subject: [PATCH 110/127] refactor(rust): split auth facade from shared types --- litellm-rust/Cargo.lock | 26 +++++--- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-aws/Cargo.toml | 2 +- litellm-rust/crates/auth-aws/src/error.rs | 6 +- litellm-rust/crates/auth-azure/Cargo.toml | 2 +- .../src/credential_provider_cache.rs | 2 +- litellm-rust/crates/auth-azure/src/native.rs | 13 ++-- litellm-rust/crates/auth-azure/src/resolve.rs | 16 ++--- litellm-rust/crates/auth-azure/src/types.rs | 4 +- litellm-rust/crates/auth-gcp/Cargo.toml | 2 +- litellm-rust/crates/auth-gcp/src/lib.rs | 2 +- litellm-rust/crates/auth-types/Cargo.toml | 15 +++++ .../{auth => auth-types}/src/credential.rs | 4 +- .../crates/{auth => auth-types}/src/error.rs | 0 .../crates/{auth => auth-types}/src/http.rs | 3 - litellm-rust/crates/auth-types/src/lib.rs | 57 +++++++++++++++++ .../crates/{auth => auth-types}/src/policy.rs | 6 +- .../crates/{auth => auth-types}/src/secret.rs | 0 .../crates/{auth => auth-types}/src/token.rs | 4 +- litellm-rust/crates/auth/Cargo.toml | 17 +++--- litellm-rust/crates/auth/src/lib.rs | 61 +++---------------- litellm-rust/crates/auth/tests/facade.rs | 33 ++++++++++ 22 files changed, 172 insertions(+), 104 deletions(-) create mode 100644 litellm-rust/crates/auth-types/Cargo.toml rename litellm-rust/crates/{auth => auth-types}/src/credential.rs (98%) rename litellm-rust/crates/{auth => auth-types}/src/error.rs (100%) rename litellm-rust/crates/{auth => auth-types}/src/http.rs (92%) create mode 100644 litellm-rust/crates/auth-types/src/lib.rs rename litellm-rust/crates/{auth => auth-types}/src/policy.rs (96%) rename litellm-rust/crates/{auth => auth-types}/src/secret.rs (100%) rename litellm-rust/crates/{auth => auth-types}/src/token.rs (95%) create mode 100644 litellm-rust/crates/auth/tests/facade.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index af3a31ddbfa..5cf9116d752 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1976,11 +1976,10 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" name = "litellm-auth" version = "0.1.0" dependencies = [ - "serde", - "subtle", - "thiserror 2.0.19", - "tokio", - "veil", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", ] [[package]] @@ -1993,7 +1992,7 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "litellm-auth", + "litellm-auth-types", "litellm-http", "moka", "reqwest 0.12.28", @@ -2009,7 +2008,7 @@ version = "0.1.0" dependencies = [ "azure_core", "azure_identity", - "litellm-auth", + "litellm-auth-types", "moka", "rstest", "serde_json", @@ -2024,13 +2023,24 @@ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ "gcp_auth", - "litellm-auth", + "litellm-auth-types", "moka", "serde_json", "sha2 0.10.9", "tokio", ] +[[package]] +name = "litellm-auth-types" +version = "0.1.0" +dependencies = [ + "serde", + "subtle", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-cache" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 4250fa81d26..b781d564805 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -14,6 +14,7 @@ litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } +litellm-auth-types = { path = "crates/auth-types" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 1f27c7bc990..1a35af48574 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index f80fbce456e..d4c6ae8cf6c 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -22,7 +22,7 @@ pub enum Error { AwsMissingWebIdentityCredentials, } -impl From for litellm_auth::Error { +impl From for litellm_auth_types::Error { fn from(error: Error) -> Self { Self::ProviderAuthentication(error.to_string()) } @@ -34,11 +34,11 @@ mod tests { #[test] fn converts_to_shared_auth_error_without_losing_context() { - let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into())); + let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into())); assert_eq!( error, - litellm_auth::Error::ProviderAuthentication( + litellm_auth_types::Error::ProviderAuthentication( "AWS profile credentials failed: profile not found".into() ) ); diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 8099506d2e5..1fd9d39d113 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs index ab9ffc719df..cd16b27f66d 100644 --- a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs +++ b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use azure_core::credentials::TokenCredential; use moka::future::Cache; -use litellm_auth::Error; +use litellm_auth_types::Error; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub(crate) struct AzureCredentialProviderCacheKey { diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 5f913a8ad01..d635e559641 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -12,8 +12,8 @@ use azure_identity::{ }; use sha2::{Digest, Sha256}; -use litellm_auth::Error; -use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::Error; +use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -484,7 +484,7 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, SecretValue, Sourced}; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) @@ -649,7 +649,7 @@ mod tests { assert!(matches!( error, - litellm_auth::Error::MixedAzureCredentialSources + litellm_auth_types::Error::MixedAzureCredentialSources )); } @@ -679,7 +679,10 @@ mod tests { authority, )) .unwrap_err(); - assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority)); + assert!(matches!( + error, + litellm_auth_types::Error::InvalidAzureAuthority + )); } } } diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 4e18cbb89aa..4d564b6e68a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,5 +1,5 @@ -use litellm_auth::Error; -use litellm_auth::{ +use litellm_auth_types::Error; +use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, SecretValue, Sourced, TokenProviderHandle, }; @@ -451,9 +451,9 @@ mod tests { }; use crate::native::ValidatedAzureRequest; use crate::types::AzureAuthInputs; - use litellm_auth::Error; - use litellm_auth::ResolvedCredential; - use litellm_auth::{ + use litellm_auth_types::Error; + use litellm_auth_types::ResolvedCredential; + use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, }; @@ -661,8 +661,8 @@ mod tests { #[derive(Debug)] struct CallerToken(&'static str); - impl litellm_auth::TokenProvider for CallerToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + impl litellm_auth_types::TokenProvider for CallerToken { + fn acquire(&self) -> litellm_auth_types::TokenFuture<'_> { Box::pin(async move { Ok(ResolvedCredential::AccessToken { token: SecretValue::new(self.0), @@ -675,7 +675,7 @@ mod tests { fn caller_inputs(token: &'static str) -> AzureAuthInputs { let params = json!({"azure_ad_token": "static-token"}); AzureAuthInputs { - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + azure_ad_token_provider: Some(litellm_auth_types::TokenProviderHandle::new(Arc::new( CallerToken(token), ))), ..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap() diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index 87e883a6a54..d5a00f09751 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_auth::{ +use litellm_auth_types::{ CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, }; use serde_json::{Map, Value}; @@ -126,7 +126,7 @@ fn source_for(sources: &BTreeMap, name: &str) -> InputSourc mod tests { use std::collections::BTreeMap; - use litellm_auth::{InputSource, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; use serde_json::json; use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index f24582db13e..045320efc80 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index bf619fee144..1ff9487e3f4 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc}; use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth::{ +use litellm_auth_types::{ CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential, }; use moka::future::Cache; diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml new file mode 100644 index 00000000000..cd65412127d --- /dev/null +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-auth-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +tokio.workspace = true diff --git a/litellm-rust/crates/auth/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs similarity index 98% rename from litellm-rust/crates/auth/src/credential.rs rename to litellm-rust/crates/auth-types/src/credential.rs index 8ed1867622a..a5d14a43b71 100644 --- a/litellm-rust/crates/auth/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,9 +5,7 @@ use std::sync::Arc; use veil::Redact; -use crate::Error; - -use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { diff --git a/litellm-rust/crates/auth/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs similarity index 100% rename from litellm-rust/crates/auth/src/error.rs rename to litellm-rust/crates/auth-types/src/error.rs diff --git a/litellm-rust/crates/auth/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs similarity index 92% rename from litellm-rust/crates/auth/src/http.rs rename to litellm-rust/crates/auth-types/src/http.rs index dd87d00e70f..0cb5839f965 100644 --- a/litellm-rust/crates/auth/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -40,9 +40,6 @@ pub fn apply_credential( ) } -/// How the upstream call is authenticated. API-key strategies become headers -/// in `prepare`; SigV4 covers the serialized body, so it is applied where the -/// outbound request is built. #[derive(Clone, Debug, PartialEq, Eq)] pub enum RequestAuth { Header { diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs new file mode 100644 index 00000000000..9d399249c05 --- /dev/null +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -0,0 +1,57 @@ +#![forbid(unsafe_code)] + +mod credential; +mod error; +pub mod http; +mod policy; +mod secret; +mod token; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + Request, + #[default] + Deployment, + Environment, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Sourced { + value: T, + source: InputSource, +} + +impl Sourced { + pub fn new(value: T, source: InputSource) -> Self { + Self { value, source } + } + + pub fn value(&self) -> &T { + &self.value + } + + pub fn source(&self) -> InputSource { + self.source + } + + pub fn into_value(self) -> T { + self.value + } + + pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { + Sourced::new(map(self.value), self.source) + } +} + +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; +pub use error::Error; +pub use http::{CredentialPlacement, RequestAuth}; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; +pub use secret::SecretValue; +pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs similarity index 96% rename from litellm-rust/crates/auth/src/policy.rs rename to litellm-rust/crates/auth-types/src/policy.rs index 4a1f5eeecf9..4c5c0365f0b 100644 --- a/litellm-rust/crates/auth/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -1,7 +1,5 @@ -use crate::Error; - -use super::http::apply_credential; -use super::{CredentialPlacement, ResolvedCredential}; +use crate::http::apply_credential; +use crate::{CredentialPlacement, Error, ResolvedCredential}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CredentialPlanKind { diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth-types/src/secret.rs similarity index 100% rename from litellm-rust/crates/auth/src/secret.rs rename to litellm-rust/crates/auth-types/src/secret.rs diff --git a/litellm-rust/crates/auth/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs similarity index 95% rename from litellm-rust/crates/auth/src/token.rs rename to litellm-rust/crates/auth-types/src/token.rs index 94da5f259fb..4175641ce10 100644 --- a/litellm-rust/crates/auth/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,9 +5,7 @@ use std::time::SystemTime; use veil::Redact; -use crate::Error; - -use super::secret::SecretValue; +use crate::{Error, SecretValue}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index 128a05c1a25..ee4900ebcc2 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -5,11 +5,14 @@ edition.workspace = true license.workspace = true repository.workspace = true -[dependencies] -serde.workspace = true -subtle.workspace = true -thiserror.workspace = true -veil.workspace = true +[features] +default = [] +aws = ["dep:litellm-auth-aws"] +azure = ["dep:litellm-auth-azure"] +gcp = ["dep:litellm-auth-gcp"] -[dev-dependencies] -tokio.workspace = true +[dependencies] +litellm-auth-types.workspace = true +litellm-auth-aws = { workspace = true, optional = true } +litellm-auth-azure = { workspace = true, optional = true } +litellm-auth-gcp = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth/src/lib.rs b/litellm-rust/crates/auth/src/lib.rs index c8d73c239b0..622a5b2d58b 100644 --- a/litellm-rust/crates/auth/src/lib.rs +++ b/litellm-rust/crates/auth/src/lib.rs @@ -1,55 +1,10 @@ -mod credential; -mod error; -pub mod http; -mod policy; -mod secret; -mod token; +#![forbid(unsafe_code)] -use serde::{Deserialize, Serialize}; +pub use litellm_auth_types::*; -#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum InputSource { - Request, - #[default] - Deployment, - Environment, -} - -#[derive(Clone, Debug, Default, Eq, PartialEq)] -pub struct Sourced { - value: T, - source: InputSource, -} - -impl Sourced { - pub fn new(value: T, source: InputSource) -> Self { - Self { value, source } - } - - pub fn value(&self) -> &T { - &self.value - } - - pub fn source(&self) -> InputSource { - self.source - } - - pub fn into_value(self) -> T { - self.value - } - - pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { - Sourced::new(map(self.value), self.source) - } -} - -pub use credential::{ - CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, - CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, -}; -pub use error::Error; -pub use http::{CredentialPlacement, RequestAuth}; -pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; -pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +#[cfg(feature = "aws")] +pub use litellm_auth_aws as aws; +#[cfg(feature = "azure")] +pub use litellm_auth_azure as azure; +#[cfg(feature = "gcp")] +pub use litellm_auth_gcp as gcp; diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs new file mode 100644 index 00000000000..f1092b15def --- /dev/null +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -0,0 +1,33 @@ +use litellm_auth::{ + CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, + ProviderAuthPolicy, ResolvedCredential, SecretValue, +}; + +const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), +}]; + +#[test] +fn facade_applies_shared_auth_policy() { + let policy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + let headers = policy + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .expect("facade policy applies"); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); +} From 82bc67b1220389ffa7e396e285e2443e6933d48e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 16:09:20 -0700 Subject: [PATCH 111/127] feat(rust): add typed secret managers and shared auth adapters --- .github/workflows/test-rust.yml | 7 + litellm-rust/Cargo.lock | 1037 ++++++++++++++++- litellm-rust/Cargo.toml | 6 + litellm-rust/crates/auth-aws/src/constants.rs | 2 + litellm-rust/crates/auth-gcp/Cargo.toml | 5 + litellm-rust/crates/auth-gcp/src/lib.rs | 60 +- litellm-rust/crates/auth-gcp/src/sdk.rs | 106 ++ litellm-rust/crates/secrets-aws/Cargo.toml | 24 + litellm-rust/crates/secrets-aws/src/auth.rs | 79 ++ litellm-rust/crates/secrets-aws/src/error.rs | 29 + litellm-rust/crates/secrets-aws/src/kms.rs | 63 + litellm-rust/crates/secrets-aws/src/lib.rs | 10 + .../crates/secrets-aws/src/secret_manager.rs | 297 +++++ litellm-rust/crates/secrets-aws/tests/kms.rs | 59 + .../secrets-aws/tests/secret_manager.rs | 292 +++++ litellm-rust/crates/secrets-google/Cargo.toml | 28 + .../crates/secrets-google/src/auth.rs | 21 + .../crates/secrets-google/src/error.rs | 43 + litellm-rust/crates/secrets-google/src/kms.rs | 67 ++ litellm-rust/crates/secrets-google/src/lib.rs | 10 + .../secrets-google/src/secret_manager.rs | 169 +++ .../crates/secrets-google/tests/kms.rs | 49 + .../secrets-google/tests/secret_manager.rs | 176 +++ litellm-rust/crates/secrets-types/Cargo.toml | 17 + .../secrets-types/src/base_secret_manager.rs | 58 + .../crates/secrets-types/src/config.rs | 92 ++ .../crates/secrets-types/src/error.rs | 9 + litellm-rust/crates/secrets-types/src/lib.rs | 12 + .../crates/secrets-types/src/value.rs | 32 + .../crates/secrets-types/tests/config.rs | 60 + .../crates/secrets-types/tests/rotation.rs | 105 ++ litellm-rust/crates/secrets/Cargo.toml | 37 + litellm-rust/crates/secrets/src/error.rs | 39 + litellm-rust/crates/secrets/src/handler.rs | 118 ++ litellm-rust/crates/secrets/src/lib.rs | 21 + litellm-rust/crates/secrets/src/oidc.rs | 264 +++++ litellm-rust/crates/secrets/src/resolver.rs | 140 +++ litellm-rust/crates/secrets/src/state.rs | 105 ++ litellm-rust/crates/secrets/tests/handler.rs | 107 ++ litellm-rust/crates/secrets/tests/oidc.rs | 295 +++++ .../crates/secrets/tests/resolution.rs | 343 ++++++ 41 files changed, 4459 insertions(+), 34 deletions(-) create mode 100644 litellm-rust/crates/auth-gcp/src/sdk.rs create mode 100644 litellm-rust/crates/secrets-aws/Cargo.toml create mode 100644 litellm-rust/crates/secrets-aws/src/auth.rs create mode 100644 litellm-rust/crates/secrets-aws/src/error.rs create mode 100644 litellm-rust/crates/secrets-aws/src/kms.rs create mode 100644 litellm-rust/crates/secrets-aws/src/lib.rs create mode 100644 litellm-rust/crates/secrets-aws/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-aws/tests/kms.rs create mode 100644 litellm-rust/crates/secrets-aws/tests/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-google/Cargo.toml create mode 100644 litellm-rust/crates/secrets-google/src/auth.rs create mode 100644 litellm-rust/crates/secrets-google/src/error.rs create mode 100644 litellm-rust/crates/secrets-google/src/kms.rs create mode 100644 litellm-rust/crates/secrets-google/src/lib.rs create mode 100644 litellm-rust/crates/secrets-google/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-google/tests/kms.rs create mode 100644 litellm-rust/crates/secrets-google/tests/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-types/Cargo.toml create mode 100644 litellm-rust/crates/secrets-types/src/base_secret_manager.rs create mode 100644 litellm-rust/crates/secrets-types/src/config.rs create mode 100644 litellm-rust/crates/secrets-types/src/error.rs create mode 100644 litellm-rust/crates/secrets-types/src/lib.rs create mode 100644 litellm-rust/crates/secrets-types/src/value.rs create mode 100644 litellm-rust/crates/secrets-types/tests/config.rs create mode 100644 litellm-rust/crates/secrets-types/tests/rotation.rs create mode 100644 litellm-rust/crates/secrets/Cargo.toml create mode 100644 litellm-rust/crates/secrets/src/error.rs create mode 100644 litellm-rust/crates/secrets/src/handler.rs create mode 100644 litellm-rust/crates/secrets/src/lib.rs create mode 100644 litellm-rust/crates/secrets/src/oidc.rs create mode 100644 litellm-rust/crates/secrets/src/resolver.rs create mode 100644 litellm-rust/crates/secrets/src/state.rs create mode 100644 litellm-rust/crates/secrets/tests/handler.rs create mode 100644 litellm-rust/crates/secrets/tests/oidc.rs create mode 100644 litellm-rust/crates/secrets/tests/resolution.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 70d828b5c8b..278fa7c425f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -127,6 +127,13 @@ jobs: cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}" done + - name: Test secret manager feature combinations + run: | + cargo test -p litellm-auth-gcp --locked --no-default-features + for features in '' aws google aws,google; do + cargo test -p litellm-secrets --locked --no-default-features --features "$features" + done + rust-wheel: runs-on: ubuntu-latest timeout-minutes: 30 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5cf9116d752..81fcabaf122 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -82,6 +82,16 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-compression" version = "0.4.46" @@ -191,9 +201,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.8.1" +version = "1.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +checksum = "25b43ad47adc2517efe3d706559d94b97e50e80e0321b3cadbc9f77cee88adcd" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -214,6 +224,58 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-kms" +version = "1.120.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6b0fe38fee2ba5b6cd24d32d08365b314ae1adea123649e061a7eb6300b6f5b" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-secretsmanager" +version = "1.117.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d32d781b34ab083e0dc54b4c68fc5e89ddf35e97d97bdcb9386d21325c14767" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sts" version = "1.108.0" @@ -243,9 +305,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.5.1" +version = "1.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" +checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -308,9 +370,9 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.2.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" +checksum = "7bd25384a4e437aa8d8f339afad4b69e786b936a7cb10db668a7aaf66717b1a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -338,9 +400,9 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.63.0" +version = "0.63.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" +checksum = "3385d469edbe8b60cc72002784652b5efca39178192aa9cc4b44c9875c6bdc18" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", @@ -368,9 +430,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.12.0" +version = "1.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +checksum = "3296253d3a91b3f938a3f2bcce4daebadcb4aa4228153fcec90f9d23532b4484" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -394,9 +456,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.13.0" +version = "1.16.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +checksum = "6d881a7b7ad179fd6611680c9de89f716fb00ab40299a9a7b8c6913e8f7511a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -423,9 +485,9 @@ dependencies = [ [[package]] name = "aws-smithy-schema" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +checksum = "e8f395d93304280b64b7632fea798d177e74897fe7f063416ce627cd6fa24829" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -434,9 +496,9 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.6.1" +version = "1.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" +checksum = "0b791f3ac597193fe1d08b82366986eb1f5bc31f2ac6c194c0855276116c76cd" dependencies = [ "base64-simd", "bytes", @@ -472,9 +534,9 @@ dependencies = [ [[package]] name = "aws-types" -version = "1.4.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +checksum = "209f3a6d82a6e9e5f94abbed94c7a26e1c052341002bf57a5fb5481f625896fc" dependencies = [ "aws-credential-types", "aws-smithy-async", @@ -580,6 +642,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.1" @@ -632,6 +700,9 @@ name = "bytes" version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "serde", +] [[package]] name = "bytes-utils" @@ -1065,6 +1136,55 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be1e0bca6c3637f992fc1cc7cbc52a78c1ef6db076dbf1059c4323d6a2048376" +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.19", +] + [[package]] name = "deranged" version = "0.5.8" @@ -1395,6 +1515,15 @@ dependencies = [ "version_check", ] +[[package]] +name = "getopts" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" +dependencies = [ + "unicode-width", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -1440,6 +1569,224 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "google-cloud-auth" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff461519b1a948200f163574be072753bcfb462a323f0eb426629d89872dd685" +dependencies = [ + "async-trait", + "base64 0.23.1", + "bytes", + "google-cloud-gax", + "hex", + "hmac", + "http 1.4.2", + "jiff", + "reqwest 0.13.5", + "rustc_version", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.19", + "time", + "tokio", + "url", +] + +[[package]] +name = "google-cloud-gax" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5615cff28ee59cfe52fbb4c11b8b1e77f650296e2ea4f4c2b7757ac6b19e752" +dependencies = [ + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.2", + "pin-project", + "rand 0.10.2", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", +] + +[[package]] +name = "google-cloud-gax-internal" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2766757d877a7a8ac23da9884cb0e3f10ed9b75a0ce59801ce6b19bf9d5819e" +dependencies = [ + "bytes", + "futures", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-rpc", + "google-cloud-wkt", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "lazy_static", + "opentelemetry", + "opentelemetry-semantic-conventions", + "opentelemetry_sdk", + "percent-encoding", + "pin-project", + "prost", + "prost-types", + "reqwest 0.13.5", + "rustc_version", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", + "tower", + "tracing", + "tracing-opentelemetry", +] + +[[package]] +name = "google-cloud-iam-v1" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f962b40234b1531e6ef73f7558871c96e117231e962086c98804323fb8d2c82" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-type", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-kms-v1" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d0f6eab19254d9abd98035cd54e3f2522d2c49abf9d93bf5425ee738d198c15" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-iam-v1", + "google-cloud-location", + "google-cloud-longrunning", + "google-cloud-lro", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-location" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "280d5acdba8fcb1232c0719ed788d85b7e362b82cbb425b7050d3ce46f075ede" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-longrunning" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c0363c5389ffda2b55cd8a86eef4b19a3481a48467c91dc5d77f626a9572766" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-lro" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47af3deef75c14a2983c430898d960c765bddbcc9f9188ca0563108e9227cfe7" +dependencies = [ + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-longrunning", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "tokio", + "tracing", +] + +[[package]] +name = "google-cloud-rpc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2162c08a89118130979ba261080e960e44cdcb2d6e2ab8ca9b1da245285d353" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-type" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63acc3a92a85f96bab021c3a3e29b53bbacc97651e1b524d4c2991960a63eb82" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fccf98cfd5481a5f5a285181ab0c62123d7d47cd2bb7299448440649349e4e7" +dependencies = [ + "base64 0.22.1", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.19", + "time", + "url", +] + [[package]] name = "h2" version = "0.3.27" @@ -1507,6 +1854,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + [[package]] name = "hex" version = "0.4.3" @@ -1636,6 +1989,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -1675,6 +2029,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.10.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -1860,6 +2227,27 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "is-macro" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.13.0" @@ -1884,6 +2272,43 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jiff" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ab1baf72f08796de0260609515130699b890ac25f30e610ad894bc5856cafdb" +dependencies = [ + "defmt", + "jiff-core", + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", +] + +[[package]] +name = "jiff-core" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e52fe76043ccecc9005d2305ebaadf7d7fc0cc89ca6baa10a94d6bc68c7128c" +dependencies = [ + "defmt", + "log", +] + +[[package]] +name = "jiff-static" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378268a1116ad67ae6228701118ac9f491d78fda38a40a1f1a9e1348de6f7212" +dependencies = [ + "jiff-core", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "jni" version = "0.22.4" @@ -1954,6 +2379,27 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "11.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75fe14a82d81e5f5af639997db37d8b96045938a7ac6ab18cdbe1c7467e05e1" +dependencies = [ + "base64 0.22.1", + "getrandom 0.2.17", + "js-sys", + "serde", + "serde_json", + "signature", + "zeroize", +] + +[[package]] +name = "lalrpop-util" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" + [[package]] name = "lazy_static" version = "1.5.0" @@ -2023,6 +2469,8 @@ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ "gcp_auth", + "google-cloud-auth", + "http 1.4.2", "litellm-auth-types", "moka", "serde_json", @@ -2252,6 +2700,90 @@ dependencies = [ "tokio-tungstenite", ] +[[package]] +name = "litellm-secrets" +version = "0.1.0" +dependencies = [ + "aws-sdk-kms", + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-kms-v1", + "jsonwebtoken", + "litellm-core-utils", + "litellm-secrets-aws", + "litellm-secrets-google", + "litellm-secrets-types", + "moka", + "reqwest 0.12.28", + "rstest", + "rustpython-parser", + "serde", + "serde_json", + "strum", + "tempfile", + "thiserror 2.0.19", + "tokio", + "tracing", + "wiremock", +] + +[[package]] +name = "litellm-secrets-aws" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-kms", + "aws-sdk-secretsmanager", + "base64 0.22.1", + "litellm-auth-aws", + "litellm-core-utils", + "litellm-secrets-types", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-google" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-kms-v1", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-types" +version = "0.1.0" +dependencies = [ + "litellm-auth-types", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-token-counter" version = "0.1.0" @@ -2277,7 +2809,7 @@ dependencies = [ "base64 0.22.1", "rand 0.8.7", "rstest", - "rustc-hash", + "rustc-hash 2.1.3", "serde", "serde_json", "thiserror 2.0.19", @@ -2453,6 +2985,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.5.1" @@ -2487,6 +3029,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2499,7 +3051,7 @@ version = "6.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" dependencies = [ - "bitflags", + "bitflags 2.13.1", "libc", "once_cell", "onig_sys", @@ -2527,6 +3079,42 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "opentelemetry" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" +dependencies = [ + "futures-core", + "futures-sink", + "js-sys", + "pin-project-lite", + "thiserror 2.0.19", + "tracing", +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c913ac17a6c451661ee255f4625d143e51647ae78ebd969b75e41c4442f4fe47" + +[[package]] +name = "opentelemetry_sdk" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" +dependencies = [ + "futures-channel", + "futures-executor", + "futures-util", + "opentelemetry", + "percent-encoding", + "portable-atomic", + "rand 0.9.5", + "thiserror 2.0.19", +] + [[package]] name = "outref" version = "0.5.2" @@ -2590,6 +3178,44 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phf" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +dependencies = [ + "phf_shared", + "rand 0.8.7", +] + +[[package]] +name = "phf_shared" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -2662,6 +3288,15 @@ version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -2712,7 +3347,7 @@ checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" dependencies = [ "bit-set", "bit-vec", - "bitflags", + "bitflags 2.13.1", "num-traits", "rand 0.9.5", "rand_chacha 0.9.0", @@ -2723,6 +3358,38 @@ dependencies = [ "unarray", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -2821,7 +3488,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.3", "rustls 0.23.42", "socket2 0.6.5", "thiserror 2.0.19", @@ -2843,7 +3510,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.3", "rustls 0.23.42", "rustls-pki-types", "slab", @@ -3022,7 +3689,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint", + "num-bigint 0.5.1", "percent-encoding", "ryu", "sha1_smol", @@ -3049,7 +3716,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.1", ] [[package]] @@ -3180,6 +3847,9 @@ dependencies = [ "rustls 0.23.42", "rustls-pki-types", "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", @@ -3248,6 +3918,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.3" @@ -3269,7 +3945,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys", @@ -3295,6 +3971,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -3374,6 +4051,63 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustpython-ast" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" +dependencies = [ + "is-macro", + "num-bigint 0.4.8", + "rustpython-parser-core", + "static_assertions", +] + +[[package]] +name = "rustpython-parser" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" +dependencies = [ + "anyhow", + "is-macro", + "itertools 0.11.0", + "lalrpop-util", + "log", + "num-bigint 0.4.8", + "num-traits", + "phf", + "phf_codegen", + "rustc-hash 1.1.0", + "rustpython-ast", + "rustpython-parser-core", + "tiny-keccak", + "unic-emoji-char", + "unic-ucd-ident", + "unicode_names2", +] + +[[package]] +name = "rustpython-parser-core" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" +dependencies = [ + "is-macro", + "memchr", + "rustpython-parser-vendored", +] + +[[package]] +name = "rustpython-parser-vendored" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" +dependencies = [ + "memchr", + "once_cell", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -3462,7 +4196,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "core-foundation", "core-foundation-sys", "libc", @@ -3622,6 +4356,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -3638,6 +4381,15 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -3660,6 +4412,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "slab" version = "0.4.12" @@ -3869,6 +4627,15 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tiktoken-rs" version = "0.12.0" @@ -3881,7 +4648,7 @@ dependencies = [ "fancy-regex 0.17.0", "lazy_static", "regex", - "rustc-hash", + "rustc-hash 2.1.3", ] [[package]] @@ -3914,6 +4681,15 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -4029,6 +4805,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-tungstenite" version = "0.24.0" @@ -4088,6 +4876,44 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "base64 0.22.1", + "bytes", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "rustls-native-certs", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -4096,11 +4922,15 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.0", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -4110,7 +4940,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "async-compression", - "bitflags", + "bitflags 2.13.1", "bytes", "futures-core", "futures-util", @@ -4167,6 +4997,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", ] [[package]] @@ -4179,6 +5010,31 @@ dependencies = [ "tracing", ] +[[package]] +name = "tracing-opentelemetry" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" +dependencies = [ + "js-sys", + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", + "web-time", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + [[package]] name = "try-lock" version = "0.2.5" @@ -4268,6 +5124,58 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" +[[package]] +name = "unic-char-property" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" +dependencies = [ + "unic-char-range", +] + +[[package]] +name = "unic-char-range" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" + +[[package]] +name = "unic-common" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" + +[[package]] +name = "unic-emoji-char" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-ident" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-version" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" +dependencies = [ + "unic-common", +] + [[package]] name = "unicase" version = "2.9.0" @@ -4295,12 +5203,40 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode_categories" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unicode_names2" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" +dependencies = [ + "phf", + "unicode_names2_generator", +] + +[[package]] +name = "unicode_names2_generator" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" +dependencies = [ + "getopts", + "log", + "phf_codegen", + "rand 0.8.7", +] + [[package]] name = "untrusted" version = "0.9.0" @@ -4348,6 +5284,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "veil" version = "0.3.0" @@ -4724,6 +5666,29 @@ dependencies = [ "memchr", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64 0.22.1", + "deadpool", + "futures", + "http 1.4.2", + "http-body-util", + "hyper 1.10.1", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "wit-bindgen" version = "0.57.1" @@ -4817,6 +5782,20 @@ name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "zerotrie" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b781d564805..2f6f5feb4ad 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -18,6 +18,10 @@ litellm-auth-types = { path = "crates/auth-types" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } +litellm-secrets = { path = "crates/secrets" } +litellm-secrets-types = { path = "crates/secrets-types" } +litellm-secrets-aws = { path = "crates/secrets-aws" } +litellm-secrets-google = { path = "crates/secrets-google" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } @@ -32,6 +36,8 @@ litellm-host-python = { path = "crates/host-python" } bytes = "1" http = "1" +google-cloud-auth = { version = "1.16.0", default-features = false } +jsonwebtoken = { version = "11.1.0", default-features = false } hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] } proptest = "1.7.0" pyo3 = "0.29.2" diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index be215cc9016..9e7c6bfab43 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -3,6 +3,8 @@ pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; pub const AWS_REGION: &str = "AWS_REGION"; +pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION"; +pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT"; pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 045320efc80..0c6258a193c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +google-sdk = ["dep:google-cloud-auth", "dep:http"] + [dependencies] litellm-auth-types.workspace = true @@ -14,3 +17,5 @@ sha2.workspace = true tokio.workspace = true gcp_auth = "0.12.7" +google-cloud-auth = { workspace = true, optional = true } +http = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 1ff9487e3f4..8aeddae9efc 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -8,6 +8,11 @@ use moka::future::Cache; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; +#[cfg(feature = "google-sdk")] +mod sdk; +#[cfg(feature = "google-sdk")] +pub use sdk::GoogleCredentials; + const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; @@ -26,19 +31,31 @@ pub struct VertexConfig { } impl VertexConfig { + pub fn new( + credentials: Option>, + project_id: Option, + location: Option, + ) -> Self { + Self { + credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()), + project_id: project_id.filter(|value| !value.trim().is_empty()), + location: location.filter(|value| !value.trim().is_empty()), + } + } + pub fn from_sourced_optional_params( params: &Map, sources: &BTreeMap, ) -> Result { - Ok(Self { - credentials: optional_credentials( + Ok(Self::new( + optional_credentials( params, sources, &["vertex_credentials", "vertex_ai_credentials"], )?, - project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?, - location: optional_string(params, &["vertex_location", "vertex_ai_location"])?, - }) + optional_string(params, &["vertex_project", "vertex_ai_project"])?, + optional_string(params, &["vertex_location", "vertex_ai_location"])?, + )) } pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self { @@ -469,6 +486,39 @@ mod tests { assert_eq!(config.location(), Some("alias-location")); } + #[test] + fn typed_config_preserves_source_and_empty_value_fallback() { + let configured = VertexConfig::new( + Some(Sourced::new( + SecretValue::new("inline-json"), + InputSource::Request, + )), + Some("project".into()), + Some("location".into()), + ); + assert!(matches!( + credential_source(&configured, &|_| Some("environment-json".into())), + CredentialSource::Inline(value) if value.expose() == "inline-json" + )); + let empty = VertexConfig::new( + Some(Sourced::new(SecretValue::new(" "), InputSource::Request)), + Some(" ".into()), + Some(" ".into()), + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(), + Some("env-location") + ); + } + #[test] fn project_and_location_prefer_input_then_environment() { let configured = diff --git a/litellm-rust/crates/auth-gcp/src/sdk.rs b/litellm-rust/crates/auth-gcp/src/sdk.rs new file mode 100644 index 00000000000..566df291b7d --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/sdk.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag}; +use google_cloud_auth::errors::CredentialsError; +use http::{Extensions, HeaderMap, HeaderName, HeaderValue}; +use litellm_auth_types::Error; + +use crate::{VertexAuth, VertexConfig}; + +type EnvironmentLookup = dyn Fn(&str) -> Option + Send + Sync; + +pub struct GoogleCredentials { + auth: VertexAuth, + config: VertexConfig, + environment: Arc, +} + +impl GoogleCredentials { + pub fn new(config: VertexConfig, environment: Arc) -> Self { + Self { + auth: VertexAuth::default(), + config, + environment, + } + } + + pub async fn request_headers(&self) -> Result { + let response = self + .auth + .validate_environment(Vec::new(), None, &self.config, &|name| { + (self.environment)(name) + }) + .await?; + response + .headers + .into_iter() + .map(|(key, value)| { + let name = + HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidHeader)?; + let value = HeaderValue::from_str(&value).map_err(|_| Error::InvalidHeader)?; + Ok((name, value)) + }) + .collect() + } +} + +impl CredentialsProvider for GoogleCredentials { + async fn headers( + &self, + _: Extensions, + ) -> Result, CredentialsError> { + self.request_headers() + .await + .map(|data| CacheableResource::New { + entity_tag: EntityTag::new(), + data, + }) + .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed")) + } + + async fn universe_domain(&self) -> Option { + None + } +} + +impl std::fmt::Debug for GoogleCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GoogleCredentials").finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn sdk_and_http_credentials_share_token_resolution_and_redaction() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private-token".into())), + ); + let direct = credentials.request_headers().await.unwrap(); + let CacheableResource::New { data, .. } = + credentials.headers(Extensions::new()).await.unwrap() + else { + panic!("first request did not return headers"); + }; + assert_eq!(direct, data); + assert_eq!(data[http::header::AUTHORIZATION], "Bearer private-token"); + assert!(!format!("{credentials:?}").contains("private-token")); + } + + #[tokio::test] + async fn invalid_token_headers_return_a_redacted_sdk_error() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private\nvalue".into())), + ); + assert_eq!( + credentials.request_headers().await.unwrap_err(), + Error::InvalidHeader + ); + let error = credentials.headers(Extensions::new()).await.unwrap_err(); + assert!(!format!("{error:?}").contains("private")); + } +} diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml new file mode 100644 index 00000000000..b1dc5b33cda --- /dev/null +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-secrets-aws" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tracing = "0.1" +veil.workspace = true +aws-sdk-kms = "1.120.0" +aws-sdk-secretsmanager = "1.117.0" +aws-credential-types = "1.3.0" + +[dev-dependencies] +base64.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs new file mode 100644 index 00000000000..954cfa2f8fd --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; +use litellm_auth_aws::{ + AwsAuthConfig, + constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME}, + resolve_credentials, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::Error; + +#[derive(Clone)] +pub(crate) struct Credentials { + config: AwsAuthConfig, + environment: Arc, +} + +impl Credentials { + pub(crate) fn new( + settings: &KeyManagementSettings, + environment: Arc, + ) -> Self { + Self { + config: AwsAuthConfig { + region_name: region(settings, environment.as_ref()).ok(), + role_name: settings.aws_role_name.clone(), + session_name: settings.aws_session_name.clone(), + external_id: settings + .aws_external_id + .as_ref() + .map(|v| v.expose().to_owned()), + profile_name: settings.aws_profile_name.clone(), + web_identity_token: settings + .aws_web_identity_token + .as_ref() + .map(|v| v.expose().to_owned()), + sts_endpoint: settings.aws_sts_endpoint.clone(), + ..Default::default() + }, + environment, + } + } +} + +impl ProvideCredentials for Credentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) + .await + .map_err(|_| { + CredentialsError::provider_error("secret manager authentication failed") + }) + }) + } +} + +pub(crate) fn region( + settings: &KeyManagementSettings, + environment: &dyn Lookup, +) -> Result { + settings + .aws_region_name + .clone() + .or_else(|| environment.get(AWS_REGION_NAME)) + .or_else(|| environment.get(AWS_REGION)) + .or_else(|| environment.get(AWS_DEFAULT_REGION)) + .ok_or(Error::MissingRegion) +} + +impl std::fmt::Debug for Credentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Credentials").finish_non_exhaustive() + } +} diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs new file mode 100644 index 00000000000..3f2c83a6a6d --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -0,0 +1,29 @@ +use aws_sdk_secretsmanager::error::SdkError; + +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("AWS authentication failed")] + Auth(#[from] #[redact] litellm_auth_aws::Error), + #[error("AWS region is not configured")] + MissingRegion, + #[error("KMS response has no plaintext")] + MissingPlaintext, + #[error("AWS request timed out")] + Timeout, + #[error("AWS KMS decrypt failed")] + Decrypt(#[from] #[redact] Box>), + #[error("AWS Secrets Manager request preparation failed")] + Read(#[from] #[redact] Box>), + #[error("AWS Secrets Manager create failed")] + Create(#[from] #[redact] Box>), + #[error("AWS Secrets Manager update failed")] + Put(#[from] #[redact] Box>), + #[error("AWS Secrets Manager delete failed")] + Delete(#[from] #[redact] Box>), + #[error("AWS Secrets Manager replication failed")] + Replicate(#[from] #[redact] Box>), + #[error("primary secret is not a JSON object")] + PrimarySecret, + #[error(transparent)] + Operation(#[from] litellm_secrets_types::Error), +} diff --git a/litellm-rust/crates/secrets-aws/src/kms.rs b/litellm-rust/crates/secrets-aws/src/kms.rs new file mode 100644 index 00000000000..a66b1c4d2fe --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/kms.rs @@ -0,0 +1,63 @@ +use litellm_auth_aws::constants::AWS_REGION_NAME; +use std::sync::Arc; + +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Region}, + primitives::Blob, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsKms { + client: Client, +} + +impl AwsKms { + pub fn new(client: Client) -> Self { + Self { client } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .ciphertext_blob(Blob::new(ciphertext)) + .send() + .await + .map_err(|error| Error::Decrypt(Box::new(error)))?; + Ok(response + .plaintext + .ok_or(Error::MissingPlaintext)? + .into_inner()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + environment + .get(AWS_REGION_NAME) + .map(|_| ()) + .ok_or(Error::MissingRegion) +} + +pub fn load_aws_kms( + use_aws_kms: Option, + settings: &KeyManagementSettings, + environment: Arc, +) -> Result, Error> { + if use_aws_kms != Some(true) { + return Ok(None); + } + if settings.aws_region_name.is_none() { + validate_environment(environment.as_ref())?; + } + let config = aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(settings, environment)) + .build(); + Ok(Some(AwsKms::new(Client::from_conf(config)))) +} diff --git a/litellm-rust/crates/secrets-aws/src/lib.rs b/litellm-rust/crates/secrets-aws/src/lib.rs new file mode 100644 index 00000000000..d17eb38c7eb --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{AwsKms, load_aws_kms}; +pub use secret_manager::{AwsSecretWriteSettings, AwsSecretsManagerV2, RotationResponse}; diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs new file mode 100644 index 00000000000..97645ecac6d --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -0,0 +1,297 @@ +use litellm_auth_aws::constants::AWS_BEDROCK_RUNTIME_ENDPOINT; +use std::{collections::BTreeMap, sync::Arc}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Region}, + operation::{ + create_secret::CreateSecretOutput, delete_secret::DeleteSecretOutput, + put_secret_value::PutSecretValueOutput, + replicate_secret_to_regions::ReplicateSecretToRegionsOutput, + }, + types::{ReplicaRegionType, Tag}, +}; +use litellm_auth_aws::constants::{ + AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{ + BaseSecretManager, KeyManagementSettings, Secret, SecretValue, async_rotate_secret, +}; +use serde_json::Value; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsSecretsManagerV2 { + client: Client, + write_settings: AwsSecretWriteSettings, +} + +#[derive(Clone, Debug, Default)] +pub struct AwsSecretWriteSettings { + pub kms_key_id: Option, + pub tags: Option>, + pub replica_regions: Option>, +} + +impl From<&KeyManagementSettings> for AwsSecretWriteSettings { + fn from(settings: &KeyManagementSettings) -> Self { + Self { + kms_key_id: settings.kms_key_id.clone(), + tags: settings.tags.clone(), + replica_regions: settings.replica_regions.clone(), + } + } +} + +#[derive(Debug)] +pub enum RotationResponse { + Created(CreateSecretOutput), + Updated(PutSecretValueOutput), +} + +impl AwsSecretsManagerV2 { + pub fn new(client: Client, write_settings: AwsSecretWriteSettings) -> Self { + Self { + client, + write_settings, + } + } + + pub fn load_aws_secret_manager( + use_aws_secret_manager: Option, + settings: KeyManagementSettings, + environment: Arc, + ) -> Result, Error> { + if use_aws_secret_manager != Some(true) { + return Ok(None); + } + let builder = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(&settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(&settings, environment.clone())); + let config = match environment.get(AWS_BEDROCK_RUNTIME_ENDPOINT) { + Some(url) => builder + .endpoint_url(url.replace("bedrock-runtime", "secretsmanager")) + .build(), + None => builder.build(), + }; + Ok(Some(Self::new( + Client::from_conf(config), + (&settings).into(), + ))) + } + + pub async fn read_secret_for_resolver( + &self, + name: &str, + primary_name: Option<&str>, + environment: &(dyn Lookup + Sync), + ) -> Result, Error> { + if bootstrap_key(name) { + return Ok(environment + .get(name) + .map(SecretValue::new) + .map(Secret::String)); + } + match primary_name.filter(|name| !name.is_empty()) { + None => self + .async_read_secret(name) + .await + .map(|value| value.map(Secret::String)), + Some(primary) => { + let value = if bootstrap_key(primary) { + environment.get(primary).map(SecretValue::new) + } else { + self.async_read_secret(primary).await? + }; + let object: Value = serde_json::from_str( + value + .as_ref() + .map(SecretValue::expose) + .filter(|v| !v.is_empty()) + .unwrap_or("{}"), + ) + .map_err(|_| Error::PrimarySecret)?; + let object = object.as_object().ok_or(Error::PrimarySecret)?; + Ok(object.get(name).cloned().and_then(Secret::from_json)) + } + } + } + + pub async fn async_read_secret(&self, name: &str) -> Result, Error> { + match self.client.get_secret_value().secret_id(name).send().await { + Ok(response) => Ok(response.secret_string.map(SecretValue::new)), + Err(error) + if matches!( + &error, + aws_sdk_secretsmanager::error::SdkError::TimeoutError(_) + ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) => + { + Err(Error::Timeout) + } + Err(error) if request_preparation_failed(&error) => Err(Error::Read(Box::new(error))), + Err(_) => { + tracing::error!("AWS secret read failed"); + Ok(None) + } + } + } + + pub async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + let response = self + .client + .create_secret() + .name(name) + .secret_string(value.expose()) + .set_description(description.filter(|v| !v.is_empty()).map(str::to_owned)) + .set_kms_key_id( + self.write_settings + .kms_key_id + .clone() + .filter(|v| !v.is_empty()), + ) + .set_tags(self.write_settings.tags.as_ref().map(|tags| { + tags.iter() + .map(|(key, value)| Tag::builder().key(key).value(value).build()) + .collect() + })) + .send() + .await + .map_err(|error| Error::Create(Box::new(error)))?; + if let Some(regions) = &self.write_settings.replica_regions + && !regions.is_empty() + && self.async_replicate_secret(name, regions).await.is_err() + { + tracing::warn!("secret created but replication failed"); + } + Ok(response) + } + + pub async fn async_replicate_secret( + &self, + name: &str, + regions: &[String], + ) -> Result, Error> { + if regions.is_empty() { + return Ok(None); + } + self.client + .replicate_secret_to_regions() + .secret_id(name) + .set_add_replica_regions(Some( + regions + .iter() + .map(|region| ReplicaRegionType::builder().region(region).build()) + .collect(), + )) + .send() + .await + .map(Some) + .map_err(|error| Error::Replicate(Box::new(error))) + } + + pub async fn async_put_secret_value( + &self, + name: &str, + value: &SecretValue, + ) -> Result { + self.client + .put_secret_value() + .secret_id(name) + .secret_string(value.expose()) + .send() + .await + .map_err(|error| Error::Put(Box::new(error))) + } + + pub async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.client + .delete_secret() + .secret_id(name) + .recovery_window_in_days(recovery_window_in_days) + .send() + .await + .map_err(|error| Error::Delete(Box::new(error))) + } + + pub async fn async_rotate_secret( + &self, + current_name: &str, + new_name: &str, + value: &SecretValue, + ) -> Result { + if current_name == new_name { + return self + .async_put_secret_value(current_name, value) + .await + .map(RotationResponse::Updated); + } + async_rotate_secret(self, current_name, new_name, value) + .await + .map(RotationResponse::Created) + } +} + +impl BaseSecretManager for AwsSecretsManagerV2 { + type Error = Error; + type WriteResponse = CreateSecretOutput; + type DeleteResponse = DeleteSecretOutput; + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + self.async_read_secret(name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + self.async_write_secret(name, value, description).await + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.async_delete_secret(name, recovery_window_in_days) + .await + } +} + +fn bootstrap_key(name: &str) -> bool { + matches!( + name, + AWS_ACCESS_KEY_ID + | AWS_SECRET_ACCESS_KEY + | AWS_REGION_NAME + | AWS_REGION + | AWS_BEDROCK_RUNTIME_ENDPOINT + ) +} + +fn request_preparation_failed( + error: &aws_sdk_secretsmanager::error::SdkError< + aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError, + >, +) -> bool { + matches!( + error, + aws_sdk_secretsmanager::error::SdkError::ConstructionFailure(_) + ) || std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |error| { + error.source() + }) + .any(|source| source.is::()) +} diff --git a/litellm-rust/crates/secrets-aws/tests/kms.rs b/litellm-rust/crates/secrets-aws/tests/kms.rs new file mode 100644 index 00000000000..39a50297551 --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/kms.rs @@ -0,0 +1,59 @@ +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_aws::AwsKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header}, +}; + +#[tokio::test] +async fn kms_decrypt_calls_the_sdk_without_applying_lookup_policy() { + let server = MockServer::start().await; + let plaintext = " private-value\n"; + Mock::given(header("x-amz-target", "TrentService.Decrypt")) + .and(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext": STANDARD.encode(plaintext)})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + let manager = AwsKms::new(client); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + plaintext.as_bytes() + ); +} + +#[test] +fn disabled_kms_loader_does_not_require_environment_configuration() { + use litellm_secrets_aws::load_aws_kms; + use litellm_secrets_types::KeyManagementSettings; + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + load_aws_kms( + enabled, + &KeyManagementSettings::default(), + Arc::new(|_: &str| None) + ) + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs new file mode 100644 index 00000000000..56698b391ec --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -0,0 +1,292 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; +use litellm_secrets_types::{KeyManagementSettings, SecretValue}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header}, +}; + +fn manager(server: &MockServer, settings: KeyManagementSettings) -> AwsSecretsManagerV2 { + let client = Client::from_conf( + aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + AwsSecretsManagerV2::new(client, (&settings).into()) +} + +#[rstest::rstest] +#[case::string_value("KEY", Some("value"))] +#[case::missing_value("missing", None)] +#[case::non_string_value("BOOL", None)] +#[tokio::test] +async fn primary_lookup_preserves_read_semantics( + #[case] name: &str, + #[case] expected: Option<&str>, +) { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId":"primary"}))) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"SecretString":json!({"KEY":"value", "BOOL":true}).to_string()}), + ), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| None) + .await + .unwrap() + .and_then(|v| v.as_str().map(str::to_owned)) + .as_deref(), + expected + ); +} + +#[rstest::rstest] +#[case::access_key("AWS_ACCESS_KEY_ID")] +#[case::secret_access_key("AWS_SECRET_ACCESS_KEY")] +#[case::region_name("AWS_REGION_NAME")] +#[case::region("AWS_REGION")] +#[case::bedrock_endpoint("AWS_BEDROCK_RUNTIME_ENDPOINT")] +#[tokio::test] +async fn bootstrap_keys_bypass_primary_lookup(#[case] name: &str) { + let server = MockServer::start().await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| Some("bootstrap".into())) + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + "bootstrap" + ); +} + +#[tokio::test] +async fn failed_read_returns_none_but_invalid_primary_json_is_an_error() { + let server = MockServer::start().await; + Mock::given(body_partial_json(json!({"SecretId":"missing"}))) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"ResourceNotFoundException"})), + ) + .mount(&server) + .await; + Mock::given(body_partial_json(json!({"SecretId":"invalid"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString":"not-json"}))) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert!( + manager + .async_read_secret("missing") + .await + .unwrap() + .is_none() + ); + assert!(matches!( + manager + .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None) + .await, + Err(Error::PrimarySecret) + )); +} + +#[tokio::test] +async fn same_name_rotation_uses_put_and_returns_its_response() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.PutSecretValue")) + .and(body_partial_json( + json!({"SecretId":"key", "SecretString":"replacement"}), + )) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"Name":"key", "VersionId":"version"})), + ) + .expect(1) + .mount(&server) + .await; + let response = manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("key", "key", &SecretValue::new("replacement")) + .await + .unwrap(); + match response { + RotationResponse::Updated(output) => assert_eq!(output.version_id(), Some("version")), + _ => panic!("rotation created a second secret"), + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn renamed_rotation_reads_creates_verifies_then_deletes() { + let server = MockServer::start().await; + let step = AtomicUsize::new(0); + Mock::given(wiremock::matchers::method("POST")) + .respond_with(move |request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + let action = request + .headers + .get("x-amz-target") + .unwrap() + .to_str() + .unwrap(); + match step.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "old"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"old-value"})) + } + 1 => { + assert_eq!(action, "secretsmanager.CreateSecret"); + assert_eq!(body["Name"], "new"); + assert_eq!(body["Description"], "Rotated from old"); + assert_eq!(body["SecretString"], "replacement"); + ResponseTemplate::new(200).set_body_json(json!({"Name":"new"})) + } + 2 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "new"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"replacement"})) + } + 3 => { + assert_eq!(action, "secretsmanager.DeleteSecret"); + assert_eq!(body["SecretId"], "old"); + assert_eq!(body["RecoveryWindowInDays"], 7); + ResponseTemplate::new(200).set_body_json(json!({"Name":"old"})) + } + _ => panic!("unexpected request"), + } + }) + .expect(4) + .mount(&server) + .await; + assert!(matches!( + manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + RotationResponse::Created(_) + )); +} + +#[tokio::test] +async fn creation_passes_tags_and_kms_and_survives_replication_failure() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.CreateSecret")) + .and(body_partial_json(json!({"Name":"key", "SecretString":"value", "KmsKeyId":"kms-key", "Tags":[{"Key":"stage", "Value":"test"}]}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"Name":"key"}))).expect(1).mount(&server).await; + Mock::given(header( + "x-amz-target", + "secretsmanager.ReplicateSecretToRegions", + )) + .and(body_partial_json( + json!({"SecretId":"key", "AddReplicaRegions":[{"Region":"replica-region"}]}), + )) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"InvalidRequestException"})), + ) + .expect(1) + .mount(&server) + .await; + let settings = KeyManagementSettings { + kms_key_id: Some("kms-key".into()), + tags: Some(std::collections::BTreeMap::from([( + "stage".into(), + "test".into(), + )])), + replica_regions: Some(vec!["replica-region".into()]), + ..Default::default() + }; + let manager = manager(&server, settings); + assert_eq!( + manager + .async_write_secret("key", &SecretValue::new("value"), None) + .await + .unwrap() + .name(), + Some("key") + ); + assert!( + manager + .async_replicate_secret("key", &[]) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn credential_failures_are_not_swallowed_as_missing_secrets() { + use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; + #[derive(Debug)] + struct FailedCredentials; + impl ProvideCredentials for FailedCredentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::ready(Err(CredentialsError::provider_error( + "private-auth-detail", + ))) + } + } + let server = MockServer::start().await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(FailedCredentials) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + let error = manager.async_read_secret("key").await.unwrap_err(); + assert!(!format!("{error:?}").contains("private-auth-detail")); + assert!(matches!(error, Error::Read(_))); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { + use std::time::Duration; + let server = MockServer::start().await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({"SecretString":"late"})), + ) + .mount(&server) + .await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .timeout_config( + aws_sdk_secretsmanager::config::timeout::TimeoutConfig::builder() + .operation_timeout(Duration::from_millis(30)) + .build(), + ) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + assert!(matches!( + manager.async_read_secret("key").await, + Err(Error::Timeout) + )); +} diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml new file mode 100644 index 00000000000..daecf20ff9e --- /dev/null +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "litellm-secrets-google" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-gcp = { workspace = true, features = ["google-sdk"] } +litellm-secrets-types.workspace = true +litellm-auth-types.workspace = true +litellm-core-utils.workspace = true +base64.workspace = true +serde_json.workspace = true +thiserror.workspace = true +moka.workspace = true +veil.workspace = true +google-cloud-kms-v1 = "1.14.0" +google-cloud-gax = { version = "1.14.0", default-features = false } +percent-encoding = "2.3" +serde.workspace = true +reqwest.workspace = true + +[dev-dependencies] +google-cloud-auth.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs new file mode 100644 index 00000000000..45fc99d8d5f --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -0,0 +1,21 @@ +use std::sync::Arc; + +use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +pub(crate) fn credentials( + project: Option, + credentials: Option, + environment: Arc, +) -> GoogleCredentials { + GoogleCredentials::new( + VertexConfig::new( + credentials.map(|value| Sourced::new(value, InputSource::Environment)), + project, + None, + ), + Arc::new(move |name| environment.get(name)), + ) +} diff --git a/litellm-rust/crates/secrets-google/src/error.rs b/litellm-rust/crates/secrets-google/src/error.rs new file mode 100644 index 00000000000..a94cc8c2de6 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/error.rs @@ -0,0 +1,43 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("Google KMS client configuration failed")] + Client( + #[from] + #[redact] + google_cloud_gax::client_builder::Error, + ), + #[error("Google authentication failed")] + Auth( + #[from] + #[redact] + litellm_auth_types::Error, + ), + #[error("Google KMS request failed")] + Kms( + #[from] + #[redact] + google_cloud_gax::error::Error, + ), + #[error("Google Secret Manager HTTP request failed")] + Http( + #[from] + #[redact] + reqwest::Error, + ), + #[error("Google Secret Manager returned HTTP {0}")] + Status(u16), + #[error("Google Secret Manager returned no payload")] + MissingPayload, + #[error("required environment variable is missing: {0}")] + MissingEnvironment(&'static str), + #[error("invalid refresh interval")] + RefreshInterval, + #[error("payload is not valid base64")] + Base64(#[from] base64::DecodeError), + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("invalid Google Secret Manager endpoint")] + Endpoint, + #[error("Google Secret Manager requires an enterprise license")] + EnterpriseRequired, +} diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs new file mode 100644 index 00000000000..3a247edaa35 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +use crate::{Error, auth}; + +const GOOGLE_APPLICATION_CREDENTIALS: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +const GOOGLE_KMS_RESOURCE_NAME: &str = "GOOGLE_KMS_RESOURCE_NAME"; + +#[derive(Clone)] +pub struct GoogleKms { + client: KeyManagementService, + resource_name: String, +} + +impl GoogleKms { + pub fn new(client: KeyManagementService, resource_name: String) -> Self { + Self { + client, + resource_name, + } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .set_name(&self.resource_name) + .set_ciphertext(ciphertext) + .send() + .await?; + Ok(response.plaintext.to_vec()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + for key in [GOOGLE_APPLICATION_CREDENTIALS, GOOGLE_KMS_RESOURCE_NAME] { + if environment.get(key).is_none() { + return Err(Error::MissingEnvironment(key)); + } + } + Ok(()) +} + +pub async fn load_google_kms( + use_google_kms: Option, + environment: Arc, +) -> Result, Error> { + if use_google_kms != Some(true) { + return Ok(None); + } + validate_environment(environment.as_ref())?; + let credentials = environment + .get(GOOGLE_APPLICATION_CREDENTIALS) + .ok_or(Error::MissingEnvironment(GOOGLE_APPLICATION_CREDENTIALS))?; + let resource_name = environment + .get(GOOGLE_KMS_RESOURCE_NAME) + .ok_or(Error::MissingEnvironment(GOOGLE_KMS_RESOURCE_NAME))?; + let credentials = auth::credentials(None, Some(SecretValue::new(credentials)), environment); + let client = KeyManagementService::builder() + .with_credentials(credentials) + .build() + .await?; + Ok(Some(GoogleKms::new(client, resource_name))) +} diff --git a/litellm-rust/crates/secrets-google/src/lib.rs b/litellm-rust/crates/secrets-google/src/lib.rs new file mode 100644 index 00000000000..a664f11a018 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{GoogleKms, load_google_kms}; +pub use secret_manager::GoogleSecretManager; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs new file mode 100644 index 00000000000..3eb8475a367 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -0,0 +1,169 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{Secret, SecretValue}; +use moka::future::Cache; +use serde::Deserialize; + +use litellm_auth_gcp::GoogleCredentials; + +use crate::{Error, auth}; + +const GOOGLE_SECRET_MANAGER_PROJECT_ID: &str = "GOOGLE_SECRET_MANAGER_PROJECT_ID"; +const GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL: &str = "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL"; +const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL"; +const GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER: &str = + "GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER"; +const GCS_PATH_SERVICE_ACCOUNT: &str = "GCS_PATH_SERVICE_ACCOUNT"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400); +const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(600); +const CACHE_CAPACITY: u64 = 200; + +#[derive(Clone)] +pub struct GoogleSecretManager { + client: reqwest::Client, + credentials: Arc, + endpoint: reqwest::Url, + project: String, + cache: Cache>, + always_read: bool, +} + +#[derive(Deserialize)] +struct Response { + payload: Option, +} + +#[derive(Deserialize)] +struct Payload { + data: Option, +} + +impl GoogleSecretManager { + pub fn with_client( + client: reqwest::Client, + endpoint: reqwest::Url, + project: String, + environment: Arc, + refresh_interval: Option, + always_read: bool, + ) -> Result { + let credentials = auth::credentials( + Some(project.clone()), + environment + .get(GCS_PATH_SERVICE_ACCOUNT) + .map(SecretValue::new), + environment, + ); + let ttl = refresh_interval + .filter(|ttl| !ttl.is_zero()) + .unwrap_or(DEFAULT_CACHE_TTL); + let cache = Cache::builder() + .max_capacity(CACHE_CAPACITY) + .time_to_live(ttl) + .build(); + Ok(Self { + client, + credentials: Arc::new(credentials), + endpoint, + project, + cache, + always_read, + }) + } + + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let project = environment + .get(GOOGLE_SECRET_MANAGER_PROJECT_ID) + .ok_or(Error::MissingEnvironment(GOOGLE_SECRET_MANAGER_PROJECT_ID))?; + let ttl = environment + .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL) + .filter(|v| !v.is_empty()) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or( + environment + .get(SECRET_MANAGER_REFRESH_INTERVAL) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64), + ); + let always_read = environment + .get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER) + .is_some_and(|v| v.eq_ignore_ascii_case("true")); + Self::with_client( + reqwest::Client::new(), + reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"), + project, + environment, + Some(if ttl < 0 { + Duration::from_nanos(1) + } else { + Duration::from_secs(ttl as u64) + }), + always_read, + ) + } + + pub async fn get_secret_from_google_secret_manager( + &self, + name: &str, + ) -> Result, Error> { + if !self.always_read + && let Some(cached) = self.cache.get(name).await + { + return Ok(cached.and_then(cached_secret)); + } + let url = self + .endpoint + .join(&format!( + "/v1/projects/{}/secrets/{}/versions/latest:access", + percent_encoding::utf8_percent_encode( + &self.project, + percent_encoding::NON_ALPHANUMERIC + ), + percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) + )) + .map_err(|_| Error::Endpoint)?; + let response = self + .client + .get(url) + .headers(self.credentials.request_headers().await?) + .send() + .await?; + if response.status() != reqwest::StatusCode::OK { + self.cache.insert(name.to_owned(), None).await; + return Err(Error::Status(response.status().as_u16())); + } + let response: Response = response.json().await?; + let Some(data) = response.payload.and_then(|payload| payload.data) else { + self.cache.insert(name.to_owned(), None).await; + return Err(Error::MissingPayload); + }; + let filtered: String = data + .chars() + .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) + .collect(); + let bytes = STANDARD.decode(filtered)?; + let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; + let value = SecretValue::new(plaintext); + self.cache + .insert(name.to_owned(), Some(value.clone())) + .await; + Ok(Some(Secret::String(value))) + } +} + +fn cached_secret(value: SecretValue) -> Option { + match serde_json::from_str(value.expose()) { + Ok(json) => Secret::from_json(json), + Err(_) => Some(Secret::String(value)), + } +} diff --git a/litellm-rust/crates/secrets-google/tests/kms.rs b/litellm-rust/crates/secrets-google/tests/kms.rs new file mode 100644 index 00000000000..667ecd268c8 --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/kms.rs @@ -0,0 +1,49 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_secrets_google::GoogleKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, +}; + +#[tokio::test] +async fn google_kms_decrypts_using_the_configured_resource() { + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .with_retry_policy(google_cloud_gax::retry_policy::NeverRetry) + .build() + .await + .unwrap(); + let manager = GoogleKms::new(client, resource.into()); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + b" value\n" + ); +} + +#[tokio::test] +async fn disabled_google_kms_loader_does_not_require_environment_configuration() { + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + litellm_secrets_google::load_google_kms(enabled, Arc::new(|_: &str| None)) + .await + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs new file mode 100644 index 00000000000..66cc0068ebb --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -0,0 +1,176 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_google::{Error, GoogleSecretManager}; +use litellm_secrets_types::Secret; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, path}, +}; + +fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { + GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), + Some(ttl), + always_read, + ) + .unwrap() +} + +#[rstest::rstest] +#[case::nonempty("private-value")] +#[case::empty("")] +#[tokio::test] +async fn successful_reads_use_auth_latest_version_and_cache_including_empty_values( + #[case] value: &str, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .and(header("authorization", "Bearer token")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(value)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + value + ); + } +} + +#[rstest::rstest] +#[case::not_found(ResponseTemplate::new(404))] +#[case::missing_payload( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"payload":{}})) +)] +#[tokio::test] +async fn negative_cache_returns_none_after_initial_error(#[case] response: ResponseTemplate) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with(response) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + assert!(matches!( + manager.get_secret_from_google_secret_manager("key").await, + Err(Error::Status(404) | Error::MissingPayload) + )); + assert!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .is_none() + ); +} + +#[rstest::rstest] +#[case::always_read(true, Duration::from_secs(60))] +#[case::expired_cache(false, Duration::from_millis(1))] +#[tokio::test] +async fn always_read_and_expired_cache_fetch_again( + #[case] always_read: bool, + #[case] ttl: Duration, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("value")}})), + ) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, always_read, ttl); + for _ in 0..2 { + tokio::time::sleep(Duration::from_millis(5)).await; + assert!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .is_some() + ); + } +} + +#[test] +fn google_manager_requires_host_license_and_project_configuration() { + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), false), + Err(Error::EnterpriseRequired) + )); + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), true), + Err(Error::MissingEnvironment( + "GOOGLE_SECRET_MANAGER_PROJECT_ID" + )) + )); +} + +#[rstest::rstest] +#[case::boolean("true", Some(Secret::Bool(true)))] +#[case::null("null", None)] +#[case::string( + "\"text\"", + Some(Secret::String(litellm_secrets_types::SecretValue::new("text"))) +)] +#[case::object( + "{\"key\":1}", + Secret::from_json(serde_json::json!({"key":1})) +)] +#[tokio::test] +async fn cached_values_preserve_python_json_conversion( + #[case] raw: &str, + #[case] expected: Option, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(raw)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some(raw) + ); + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap(), + expected + ); +} diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml new file mode 100644 index 00000000000..acd29746722 --- /dev/null +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-secrets-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-types.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs new file mode 100644 index 00000000000..d71bce64221 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -0,0 +1,58 @@ +use crate::{Error, SecretValue}; + +pub fn validate_secret_name(name: &str) -> Result<(), Error> { + if name.split('/').any(|segment| segment == "..") + || name + .chars() + .any(|c| c.is_control() || matches!(c, '\u{2028}' | '\u{2029}')) + { + return Err(Error::UnsafeSecretName); + } + Ok(()) +} + +#[expect( + async_fn_in_trait, + reason = "closed backend dispatch does not require Send bounds on generic rotation" +)] +pub trait BaseSecretManager { + type Error: From; + type WriteResponse; + type DeleteResponse; + + async fn async_read_secret(&self, name: &str) -> Result, Self::Error>; + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result; + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result; +} + +pub async fn async_rotate_secret( + manager: &M, + current_name: &str, + new_name: &str, + value: &SecretValue, +) -> Result { + if manager.async_read_secret(current_name).await?.is_none() { + return Err(Error::CurrentSecretMissing.into()); + } + let response = manager + .async_write_secret( + new_name, + value, + Some(&format!("Rotated from {current_name}")), + ) + .await?; + if manager.async_read_secret(new_name).await?.is_none() { + return Err(Error::NewSecretMissing.into()); + } + manager.async_delete_secret(current_name, 7).await?; + Ok(response) +} diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs new file mode 100644 index 00000000000..36d319311a3 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -0,0 +1,92 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::SecretValue; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum KeyManagementSystem { + GoogleKms, + AzureKeyVault, + AwsSecretManager, + GoogleSecretManager, + HashicorpVault, + Cyberark, + Local, + AwsKms, + Custom, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AccessMode { + #[default] + ReadOnly, + WriteOnly, + ReadAndWrite, +} + +impl AccessMode { + pub fn readable(self) -> bool { + matches!(self, Self::ReadOnly | Self::ReadAndWrite) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(default)] +pub struct KeyManagementSettings { + pub hosted_keys: Option>, + pub store_virtual_keys: Option, + pub prefix_for_stored_virtual_keys: String, + pub access_mode: AccessMode, + pub primary_secret_name: Option, + pub description: Option, + pub tags: Option>, + pub kms_key_id: Option, + pub custom_secret_manager: Option, + pub aws_region_name: Option, + pub aws_role_name: Option, + pub aws_session_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_external_id: Option, + pub aws_profile_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_web_identity_token: Option, + pub aws_sts_endpoint: Option, + pub replica_regions: Option>, +} + +impl Default for KeyManagementSettings { + fn default() -> Self { + Self { + hosted_keys: None, + store_virtual_keys: Some(false), + prefix_for_stored_virtual_keys: "litellm/".into(), + access_mode: AccessMode::ReadOnly, + primary_secret_name: None, + description: None, + tags: None, + kms_key_id: None, + custom_secret_manager: None, + aws_region_name: None, + aws_role_name: None, + aws_session_name: None, + aws_external_id: None, + aws_profile_name: None, + aws_web_identity_token: None, + aws_sts_endpoint: None, + replica_regions: None, + } + } +} + +fn serialize_secret( + value: &Option, + serializer: S, +) -> Result { + value + .as_ref() + .map(SecretValue::expose) + .serialize(serializer) +} diff --git a/litellm-rust/crates/secrets-types/src/error.rs b/litellm-rust/crates/secrets-types/src/error.rs new file mode 100644 index 00000000000..cae9c7f4c69 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/error.rs @@ -0,0 +1,9 @@ +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("secret name contains an unsafe path segment or control character")] + UnsafeSecretName, + #[error("current secret was not found")] + CurrentSecretMissing, + #[error("new secret could not be verified")] + NewSecretMissing, +} diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs new file mode 100644 index 00000000000..0823ed13c06 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -0,0 +1,12 @@ +#![forbid(unsafe_code)] + +mod base_secret_manager; +mod config; +mod error; +mod value; + +pub use base_secret_manager::{BaseSecretManager, async_rotate_secret, validate_secret_name}; +pub use config::{AccessMode, KeyManagementSettings, KeyManagementSystem}; +pub use error::Error; +pub use litellm_auth_types::SecretValue; +pub use value::Secret; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs new file mode 100644 index 00000000000..524045ac007 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -0,0 +1,32 @@ +use crate::SecretValue; + +#[derive(Clone, PartialEq, Eq, veil::Redact)] +pub enum Secret { + String(SecretValue), + Bool(#[redact] bool), + Json(#[redact] serde_json::Value), +} + +impl From for Secret { + fn from(value: SecretValue) -> Self { + Self::String(value) + } +} + +impl Secret { + pub fn from_json(value: serde_json::Value) -> Option { + match value { + serde_json::Value::Null => None, + serde_json::Value::String(value) => Some(Self::String(SecretValue::new(value))), + serde_json::Value::Bool(value) => Some(Self::Bool(value)), + value => Some(Self::Json(value)), + } + } + + pub fn as_str(&self) -> Option<&str> { + match self { + Self::String(value) => Some(value.expose()), + Self::Bool(_) | Self::Json(_) => None, + } + } +} diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs new file mode 100644 index 00000000000..4a5f17bc68a --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -0,0 +1,60 @@ +use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +use serde_json::json; + +#[test] +fn config_preserves_defaults_nulls_and_serialized_names() { + let empty: KeyManagementSettings = serde_json::from_value(json!({})).unwrap(); + assert_eq!(empty, KeyManagementSettings::default()); + assert_eq!(empty.access_mode, AccessMode::ReadOnly); + assert_eq!(empty.store_virtual_keys, Some(false)); + assert_eq!(empty.prefix_for_stored_virtual_keys, "litellm/"); + let configured: KeyManagementSettings = serde_json::from_value(json!({ + "hosted_keys": [], "store_virtual_keys": null, "access_mode": "write_only", + "aws_web_identity_token": "private-token", "aws_external_id": "private-id", + "tags": {"stage": "test"}, "replica_regions": ["test-region"] + })) + .unwrap(); + assert!(!configured.access_mode.readable()); + assert_eq!(configured.store_virtual_keys, None); + assert_eq!(configured.hosted_keys.as_deref(), Some([].as_slice())); + assert!(!format!("{configured:?}").contains("private-")); + let serialized = serde_json::to_value(&configured).unwrap(); + assert_eq!(serialized["access_mode"], "write_only"); + assert_eq!(serialized["aws_web_identity_token"], "private-token"); + assert_eq!( + serde_json::from_value::(serialized).unwrap(), + configured + ); +} + +#[rstest::rstest] +#[case::aws_kms("aws_kms", KeyManagementSystem::AwsKms)] +#[case::aws_secret_manager("aws_secret_manager", KeyManagementSystem::AwsSecretManager)] +#[case::google_kms("google_kms", KeyManagementSystem::GoogleKms)] +#[case::google_secret_manager("google_secret_manager", KeyManagementSystem::GoogleSecretManager)] +#[case::azure_key_vault("azure_key_vault", KeyManagementSystem::AzureKeyVault)] +#[case::hashicorp_vault("hashicorp_vault", KeyManagementSystem::HashicorpVault)] +#[case::cyberark("cyberark", KeyManagementSystem::Cyberark)] +#[case::custom("custom", KeyManagementSystem::Custom)] +#[case::local("local", KeyManagementSystem::Local)] +fn key_management_system_serialization_round_trips( + #[case] name: &str, + #[case] system: KeyManagementSystem, +) { + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + system + ); + assert_eq!(serde_json::to_value(system).unwrap(), name); +} + +#[test] +fn secret_debug_never_exposes_values() { + assert!( + !format!("{:?}", Secret::String(SecretValue::new("sensitive-value"))) + .contains("sensitive-value") + ); + assert!(!format!("{:?}", Secret::Bool(true)).contains("true")); +} diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs new file mode 100644 index 00000000000..48a5304ece8 --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -0,0 +1,105 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_secrets_types::{ + BaseSecretManager, Error, SecretValue, async_rotate_secret, validate_secret_name, +}; + +struct Manager { + step: AtomicUsize, + absent_at: Option, +} + +impl BaseSecretManager for Manager { + type Error = Error; + type WriteResponse = &'static str; + type DeleteResponse = (); + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + let step = self.step.fetch_add(1, Ordering::SeqCst); + assert_eq!(name, if step == 0 { "old" } else { "new" }); + Ok((self.absent_at != Some(step)).then(|| SecretValue::new("value"))) + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 1); + assert_eq!(name, "new"); + assert_eq!(value.expose(), "replacement"); + assert_eq!(description, Some("Rotated from old")); + Ok("provider-response") + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result<(), Error> { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 3); + assert_eq!(name, "old"); + assert_eq!(recovery_window_in_days, 7); + Ok(()) + } +} + +#[tokio::test] +async fn rotation_verifies_before_deleting_and_returns_provider_response() { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: None, + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + "provider-response" + ); + assert_eq!(manager.step.load(Ordering::SeqCst), 4); +} + +#[rstest::rstest] +#[case::current_secret_missing(0, Error::CurrentSecretMissing, 1)] +#[case::new_secret_missing(2, Error::NewSecretMissing, 3)] +#[tokio::test] +async fn missing_old_or_new_value_stops_rotation_before_deletion( + #[case] absent_at: usize, + #[case] expected: Error, + #[case] calls: usize, +) { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: Some(absent_at), + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap_err(), + expected + ); + assert_eq!(manager.step.load(Ordering::SeqCst), calls); +} + +#[rstest::rstest] +#[case::parent("..")] +#[case::parent_prefix("../x")] +#[case::parent_segment("x/../y")] +#[case::parent_suffix("x/..")] +#[case::line_feed("line\n")] +#[case::next_line("\u{85}")] +#[case::line_separator("\u{2028}")] +#[case::paragraph_separator("\u{2029}")] +fn names_reject_path_traversal_and_control_characters(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Err(Error::UnsafeSecretName)); +} + +#[rstest::rstest] +#[case::embedded_double_dot("release-1.0..2")] +#[case::path_separator("folder/key")] +#[case::empty("")] +#[case::three_dots("...")] +fn names_allow_safe_values(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Ok(())); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml new file mode 100644 index 00000000000..3414470d234 --- /dev/null +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "litellm-secrets" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[features] +default = [] +aws = ["dep:litellm-secrets-aws"] +google = ["dep:litellm-secrets-google"] + +[dependencies] +litellm-secrets-types.workspace = true +litellm-secrets-aws = { workspace = true, optional = true } +litellm-secrets-google = { workspace = true, optional = true } +litellm-core-utils.workspace = true +base64.workspace = true +serde.workspace = true +strum.workspace = true +jsonwebtoken.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tracing = "0.1" +reqwest.workspace = true +moka.workspace = true +tokio = { workspace = true, features = ["fs"] } + +rustpython-parser = { version = "0.4.0", default-features = false, features = ["num-bigint"] } + +[dev-dependencies] +rstest.workspace = true +wiremock = "0.6.5" +tempfile = "3" +aws-sdk-kms = "1.120.0" +google-cloud-kms-v1 = "1.14.0" +google-cloud-auth.workspace = true diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs new file mode 100644 index 00000000000..d240adb6a67 --- /dev/null +++ b/litellm-rust/crates/secrets/src/error.rs @@ -0,0 +1,39 @@ +use crate::KeyManagementSystem; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("encrypted environment value is missing")] + MissingCiphertext, + #[error("ciphertext is not valid base64 for the configured manager")] + InvalidCiphertext, + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("secret manager backend is not compiled: {0:?}")] + UnsupportedBackend(KeyManagementSystem), + #[error("configured secret manager does not match its backend")] + BackendMismatch, + #[error("unsupported OIDC provider or missing build feature")] + UnsupportedOidc, + #[error("OIDC reference requires a provider and audience")] + InvalidOidc, + #[error("OIDC environment variable is missing")] + MissingEnvironment, + #[error("OIDC request failed")] + OidcHttp, + #[error("OIDC provider returned HTTP {0}")] + OidcStatus(u16), + #[error("OIDC response is invalid")] + OidcResponse, + #[error("OIDC file path must be absolute and within the credential allowlist")] + UnsafeOidcPath, + #[error("OIDC file could not be read")] + OidcFile, + #[error("secret manager returned no secret")] + MissingSecret, + #[cfg(feature = "aws")] + #[error(transparent)] + Aws(#[from] litellm_secrets_aws::Error), + #[cfg(feature = "google")] + #[error(transparent)] + Google(#[from] litellm_secrets_google::Error), +} diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs new file mode 100644 index 00000000000..6d19965605d --- /dev/null +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -0,0 +1,118 @@ +use litellm_core_utils::settings::Lookup; + +use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue}; + +#[derive(Clone)] +pub enum SecretManager { + Local, + #[cfg(feature = "aws")] + AwsKms(crate::aws::AwsKms), + #[cfg(feature = "aws")] + AwsSecretsManagerV2(crate::aws::AwsSecretsManagerV2), + #[cfg(feature = "google")] + GoogleKms(crate::google::GoogleKms), + #[cfg(feature = "google")] + GoogleSecretManager(crate::google::GoogleSecretManager), +} + +impl SecretManager { + pub fn system(&self) -> KeyManagementSystem { + match self { + Self::Local => KeyManagementSystem::Local, + #[cfg(feature = "aws")] + Self::AwsKms(_) => KeyManagementSystem::AwsKms, + #[cfg(feature = "aws")] + Self::AwsSecretsManagerV2(_) => KeyManagementSystem::AwsSecretManager, + #[cfg(feature = "google")] + Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, + #[cfg(feature = "google")] + Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + } + } +} + +pub async fn get_secret_from_manager( + client: &SecretManager, + secret_name: &str, + _settings: &KeyManagementSettings, + environment: &(dyn Lookup + Send + Sync), +) -> Result, Error> { + match client { + SecretManager::Local => Ok(environment + .get(secret_name) + .map(SecretValue::new) + .map(Secret::String)), + #[cfg(feature = "aws")] + SecretManager::AwsKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value.trim())))) + } + #[cfg(feature = "google")] + SecretManager::GoogleKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value)))) + } + #[cfg(feature = "aws")] + SecretManager::AwsSecretsManagerV2(client) => client + .read_secret_for_resolver( + secret_name, + _settings.primary_secret_name.as_deref(), + environment, + ) + .await + .map_err(Error::from), + #[cfg(feature = "google")] + SecretManager::GoogleSecretManager(client) => client + .get_secret_from_google_secret_manager(secret_name) + .await? + .map(Some) + .ok_or(Error::MissingSecret), + } +} + +#[cfg(any(feature = "aws", feature = "google"))] +#[derive(Clone, Copy)] +enum Base64Mode { + #[cfg(feature = "google")] + Canonical, + #[cfg(feature = "aws")] + Permissive, +} + +#[cfg(any(feature = "aws", feature = "google"))] +fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result, Error> { + use base64::{Engine, engine::general_purpose::STANDARD}; + let canonical = match mode { + #[cfg(feature = "google")] + Base64Mode::Canonical => true, + #[cfg(feature = "aws")] + Base64Mode::Permissive => false, + }; + let encoded = if canonical { + value.to_owned() + } else { + value + .chars() + .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) + .collect() + }; + let ciphertext = STANDARD + .decode(&encoded) + .map_err(|_| Error::InvalidCiphertext)?; + if canonical && STANDARD.encode(&ciphertext) != encoded { + return Err(Error::InvalidCiphertext); + } + Ok(ciphertext) +} diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs new file mode 100644 index 00000000000..c434cfaaf03 --- /dev/null +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -0,0 +1,21 @@ +#![forbid(unsafe_code)] + +mod error; +mod handler; +mod oidc; +mod resolver; +mod state; + +pub use error::Error; +pub use handler::{SecretManager, get_secret_from_manager}; +pub use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +pub use oidc::{OidcProvider, OidcReference, OidcResolver}; +pub use resolver::SecretResolver; +pub use state::{SecretManagerState, secret_manager_would_be_consulted}; + +#[cfg(feature = "aws")] +pub use litellm_secrets_aws as aws; +#[cfg(feature = "google")] +pub use litellm_secrets_google as google; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs new file mode 100644 index 00000000000..48e6a9bc0b0 --- /dev/null +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -0,0 +1,264 @@ +use std::{ + path::Path, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use jsonwebtoken::dangerous::insecure_decode_claims; +use litellm_core_utils::settings::Lookup; +use moka::future::Cache; +use serde::Deserialize; + +use crate::{Error, SecretValue}; + +const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); +const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); +const TOKEN_EXPIRY_MARGIN_SECONDS: f64 = 60.0; +const CIRCLE_OIDC_TOKEN: &str = "CIRCLE_OIDC_TOKEN"; +const CIRCLE_OIDC_TOKEN_V2: &str = "CIRCLE_OIDC_TOKEN_V2"; +const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE"; +const ACTIONS_ID_TOKEN_REQUEST_URL: &str = "ACTIONS_ID_TOKEN_REQUEST_URL"; +const ACTIONS_ID_TOKEN_REQUEST_TOKEN: &str = "ACTIONS_ID_TOKEN_REQUEST_TOKEN"; +const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS"; +const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets"; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)] +#[strum(serialize_all = "snake_case")] +pub enum OidcProvider { + Google, + #[strum(serialize = "circleci")] + CircleCi, + #[strum(serialize = "circleci_v2")] + CircleCiV2, + Github, + Azure, + File, + Env, + EnvPath, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct OidcReference<'a> { + pub provider: OidcProvider, + pub audience: &'a str, +} + +impl<'a> TryFrom<&'a str> for OidcReference<'a> { + type Error = Error; + + fn try_from(reference: &'a str) -> Result { + let (provider, audience) = reference + .strip_prefix("oidc/") + .and_then(|body| body.split_once('/')) + .ok_or(Error::InvalidOidc)?; + Ok(Self { + provider: provider.parse().map_err(|_| Error::UnsupportedOidc)?, + audience, + }) + } +} + +#[derive(Deserialize)] +struct OidcTokenClaims { + exp: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum NumericDate { + Number(f64), + String(String), +} + +impl NumericDate { + fn seconds(self) -> Option { + match self { + Self::Number(value) => Some(value), + Self::String(value) => value.parse().ok(), + } + .filter(|value| value.is_finite()) + } +} + +pub struct OidcResolver { + client: reqwest::Client, + google_identity_endpoint: reqwest::Url, + cache: Cache, + clock: fn() -> SystemTime, +} + +impl Default for OidcResolver { + fn default() -> Self { + Self::new( + reqwest::Client::builder().timeout(Duration::from_secs(600)).connect_timeout(Duration::from_secs(5)).build().expect("HTTP client configuration"), + reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), + ) + } +} + +impl OidcResolver { + pub fn new(client: reqwest::Client, google_identity_endpoint: reqwest::Url) -> Self { + Self { + client, + google_identity_endpoint, + cache: Cache::builder() + .max_capacity(200) + .time_to_live(GOOGLE_TOKEN_MAX_TTL) + .build(), + clock: SystemTime::now, + } + } + + pub fn with_clock(self, clock: fn() -> SystemTime) -> Self { + Self { clock, ..self } + } + + pub async fn resolve( + &self, + reference: &str, + environment: &(dyn Lookup + Send + Sync), + ) -> Result, Error> { + let OidcReference { provider, audience } = reference.try_into()?; + match provider { + OidcProvider::CircleCi => required_env(environment, CIRCLE_OIDC_TOKEN) + .map(SecretValue::new) + .map(Some), + OidcProvider::CircleCiV2 => required_env(environment, CIRCLE_OIDC_TOKEN_V2) + .map(SecretValue::new) + .map(Some), + OidcProvider::Env => required_env(environment, audience) + .map(SecretValue::new) + .map(Some), + OidcProvider::EnvPath => read_file(&required_env(environment, audience)?) + .await + .map(Some), + OidcProvider::File => read_allowed_file(audience, environment).await.map(Some), + OidcProvider::Azure => { + if let Some(path) = environment.get(AZURE_FEDERATED_TOKEN_FILE) { + return read_file(&path).await.map(Some); + } + Err(Error::UnsupportedOidc) + } + OidcProvider::Github => { + let url = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_URL)?; + let authorization = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_TOKEN)?; + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(url) + .query(&[("audience", audience)]) + .bearer_auth(authorization) + .header("Accept", "application/json; api-version=2.0") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + #[derive(Deserialize)] + struct Token { + value: Option, + } + let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?; + if let Some(value) = &token.value { + self.cache + .insert( + reference.to_owned(), + (value.clone(), (self.clock)() + GITHUB_TOKEN_TTL), + ) + .await; + } + Ok(token.value) + } + OidcProvider::Google => { + if !cfg!(feature = "google") { + return Err(Error::UnsupportedOidc); + } + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(self.google_identity_endpoint.clone()) + .query(&[("audience", audience)]) + .header("Metadata-Flavor", "Google") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + let token = response.text().await.map_err(|_| Error::OidcResponse)?; + let now = (self.clock)(); + let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL); + let value = SecretValue::new(token); + if let Some(ttl) = ttl.filter(|ttl| !ttl.is_zero()) { + self.cache + .insert(reference.to_owned(), (value.clone(), now + ttl)) + .await; + } + Ok(Some(value)) + } + } + } + + async fn cached(&self, reference: &str) -> Option { + self.cache + .get(reference) + .await + .and_then(|(value, expires)| ((self.clock)() < expires).then_some(value)) + } +} + +fn required_env(environment: &dyn Lookup, name: &str) -> Result { + environment.get(name).ok_or(Error::MissingEnvironment) +} + +async fn read_file(path: &str) -> Result { + tokio::fs::read_to_string(path) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile) +} + +async fn read_allowed_file( + path: &str, + environment: &(dyn Lookup + Sync), +) -> Result { + if !Path::new(path).is_absolute() { + return Err(Error::UnsafeOidcPath); + } + let resolved = tokio::fs::canonicalize(path) + .await + .map_err(|_| Error::OidcFile)?; + let allowed = environment + .get(OIDC_ALLOWED_CREDENTIAL_DIRS) + .filter(|v| !v.is_empty()) + .unwrap_or_else(|| DEFAULT_CREDENTIAL_DIRS.into()); + for directory in allowed.split(',').map(str::trim).filter(|d| !d.is_empty()) { + if let Ok(directory) = tokio::fs::canonicalize(directory).await + && resolved.starts_with(directory) + { + return tokio::fs::read_to_string(&resolved) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile); + } + } + Err(Error::UnsafeOidcPath) +} + +fn oidc_token_cache_ttl(token: &str, now: SystemTime, max_ttl: Duration) -> Option { + let fallback = Some(max_ttl); + let Ok(claims) = insecure_decode_claims::(token) else { + return fallback; + }; + let Some(exp) = claims.exp.and_then(NumericDate::seconds) else { + return fallback; + }; + let seconds = exp.trunc() + - now.duration_since(UNIX_EPOCH).ok()?.as_secs() as f64 + - TOKEN_EXPIRY_MARGIN_SECONDS; + (seconds > 0.0).then(|| Duration::from_secs_f64(seconds.min(max_ttl.as_secs_f64()))) +} diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs new file mode 100644 index 00000000000..93daddae0c4 --- /dev/null +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -0,0 +1,140 @@ +use std::sync::Arc; + +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; + +use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; + +use crate::state::{LookupTarget, normalize_secret_name}; + +pub struct SecretResolver { + state: Arc, + environment: Arc, + oidc: OidcResolver, +} + +impl Default for SecretResolver { + fn default() -> Self { + Self::new( + Arc::new(SecretManagerState::default()), + Arc::new(ProcessEnvironment), + OidcResolver::default(), + ) + } +} + +impl SecretResolver { + pub fn new( + state: Arc, + environment: Arc, + oidc: OidcResolver, + ) -> Self { + Self { + state, + environment, + oidc, + } + } + + pub async fn get_secret( + &self, + name: &str, + _default_value: Option, + ) -> Result, Error> { + let name = normalize_secret_name(name); + if name.starts_with("oidc/") { + return self + .oidc + .resolve(name, self.environment.as_ref()) + .await + .map(|value| value.map(Secret::String)); + } + if !self.state.readable() { + return Ok(self + .environment + .get(name) + .map(|value| match str_to_bool(&value) { + Some(value) => Secret::Bool(value), + None => Secret::String(SecretValue::new(value)), + })); + } + let result = match self.state.lookup_target(name) { + LookupTarget::Environment => Ok(self.environment_secret(name)), + LookupTarget::Manager { backend, settings } => { + crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) + .await + } + }; + let value = match result { + Ok(value) => value, + Err(_) => { + tracing::error!("secret manager lookup failed; falling back to environment"); + self.environment_secret(name) + } + }; + Ok(value.and_then(managed_secret)) + } + + fn environment_secret(&self, name: &str) -> Option { + self.environment + .get(name) + .map(SecretValue::new) + .map(Secret::String) + } + + pub async fn get_secret_str( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + Ok(match self.get_secret(name, default_value).await? { + Some(Secret::String(value)) => Some(value), + Some(Secret::Bool(_) | Secret::Json(_)) | None => None, + }) + } + + pub async fn get_secret_bool( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + Ok( + match self + .get_secret(name, default_value.map(Secret::Bool)) + .await? + { + Some(Secret::Bool(value)) => Some(value), + Some(Secret::String(value)) => str_to_bool(value.expose()), + Some(Secret::Json(_)) | None => None, + }, + ) + } +} + +fn str_to_bool(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Some(true), + "false" => Some(false), + _ => None, + } +} + +fn literal_bool(value: &str) -> Option { + use rustpython_parser::{Parse, ast}; + match ast::Expr::parse(value.trim_start_matches([' ', '\t']), "").ok()? { + ast::Expr::Constant(node) => match node.value { + ast::Constant::Bool(value) => Some(value), + _ => None, + }, + _ => None, + } +} + +fn managed_secret(value: Secret) -> Option { + match value { + Secret::String(value) => Some(match literal_bool(value.expose()) { + Some(boolean) => Secret::Bool(boolean), + None => Secret::String(value), + }), + Secret::Bool(_) | Secret::Json(_) => None, + } +} diff --git a/litellm-rust/crates/secrets/src/state.rs b/litellm-rust/crates/secrets/src/state.rs new file mode 100644 index 00000000000..ca942aa65db --- /dev/null +++ b/litellm-rust/crates/secrets/src/state.rs @@ -0,0 +1,105 @@ +use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; + +pub(crate) enum LookupTarget<'a> { + Environment, + Manager { + backend: &'a SecretManager, + settings: &'a KeyManagementSettings, + }, +} + +pub(crate) fn normalize_secret_name(name: &str) -> &str { + name.strip_prefix("os.environ/").unwrap_or(name) +} + +#[derive(Clone, Default)] +pub struct SecretManagerState { + system: Option, + settings: Option, + backend: Option, +} + +impl SecretManagerState { + pub fn new( + system: Option, + settings: Option, + backend: Option, + ) -> Result { + if let Some(system) = system { + let available = match system { + KeyManagementSystem::Local => true, + KeyManagementSystem::AwsKms | KeyManagementSystem::AwsSecretManager => { + cfg!(feature = "aws") + } + KeyManagementSystem::GoogleKms | KeyManagementSystem::GoogleSecretManager => { + cfg!(feature = "google") + } + KeyManagementSystem::AzureKeyVault + | KeyManagementSystem::HashicorpVault + | KeyManagementSystem::Cyberark + | KeyManagementSystem::Custom => false, + }; + if !available { + return Err(Error::UnsupportedBackend(system)); + } + if let Some(backend) = &backend + && system != backend.system() + { + return Err(Error::BackendMismatch); + } + } + Ok(Self { + system, + settings, + backend, + }) + } + + pub fn system(&self) -> Option { + self.system + } + pub fn settings(&self) -> Option<&KeyManagementSettings> { + self.settings.as_ref() + } + pub fn backend(&self) -> Option<&SecretManager> { + self.backend.as_ref() + } + + pub(crate) fn readable(&self) -> bool { + self.backend.is_some() + && self + .settings + .as_ref() + .is_some_and(|settings| settings.access_mode.readable()) + } + + pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> { + match (&self.backend, &self.settings) { + (Some(backend), Some(settings)) + if settings.access_mode.readable() + && hosts_secret(settings, name) + && self + .system + .is_some_and(|system| system != KeyManagementSystem::Local) => + { + LookupTarget::Manager { backend, settings } + } + _ => LookupTarget::Environment, + } + } +} + +pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool { + state.readable() + && state + .settings + .as_ref() + .is_some_and(|settings| hosts_secret(settings, normalize_secret_name(name))) +} + +fn hosts_secret(settings: &KeyManagementSettings, name: &str) -> bool { + settings + .hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == name)) +} diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs new file mode 100644 index 00000000000..a2cbbd843e1 --- /dev/null +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -0,0 +1,107 @@ +#[cfg(feature = "aws")] +#[tokio::test] +async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() { + use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region}, + }; + use base64::{Engine, engine::general_purpose::STANDARD}; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, aws::AwsKms, get_secret_from_manager, + }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_json}; + + let server = MockServer::start().await; + Mock::given(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .build(), + ); + let manager = SecretManager::AwsKms(AwsKms::new(client)); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|name: &str| { + assert_eq!(name, "KEY"); + Some(format!(" {}\n", STANDARD.encode("encrypted"))) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some("value")); + assert!(!format!("{value:?}").contains("value")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await, + Err(Error::InvalidCiphertext) + )); +} + +#[cfg(feature = "google")] +#[tokio::test] +async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whitespace() { + use base64::{Engine, engine::general_purpose::STANDARD}; + use google_cloud_kms_v1::client::KeyManagementService; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, get_secret_from_manager, google::GoogleKms, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, + }; + + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .build() + .await + .unwrap(); + let manager = SecretManager::GoogleKms(GoogleKms::new(client, resource.into())); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| { + Some(STANDARD.encode("encrypted")) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some(" value\n")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some(format!( + " {}", + STANDARD.encode("encrypted") + ))) + .await, + Err(Error::InvalidCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); +} diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs new file mode 100644 index 00000000000..b17e7de7f9d --- /dev/null +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -0,0 +1,295 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets::{Error, OidcResolver, Secret, SecretManagerState, SecretResolver}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method, path, query_param}, +}; + +fn environment(pairs: &[(&str, &str)]) -> Arc { + let values: BTreeMap = pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + Arc::new(move |name: &str| values.get(name).cloned()) +} + +#[rstest::rstest] +#[case::environment("oidc/env/TOKEN", "true")] +#[case::circleci("oidc/circleci/audience", "circle")] +#[case::circleci_v2("oidc/circleci_v2/audience", "circle-v2")] +#[tokio::test] +async fn environment_sources_resolve_expected_value( + #[case] reference: &str, + #[case] expected: &str, +) { + let env = environment(&[ + ("TOKEN", "true"), + ("CIRCLE_OIDC_TOKEN", "circle"), + ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), + ]); + assert_eq!( + OidcResolver::default() + .resolve(reference, env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + expected + ); +} + +#[tokio::test] +async fn environment_sources_bypass_boolean_conversion_and_defaults() { + let env = environment(&[("TOKEN", "true")]); + let oidc = OidcResolver::default(); + let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); + assert_eq!( + resolver + .get_secret_str("os.environ/oidc/env/TOKEN", None) + .await + .unwrap() + .unwrap() + .expose(), + "true" + ); + assert_eq!( + resolver + .get_secret_bool("oidc/env/TOKEN", None) + .await + .unwrap(), + Some(true) + ); + assert!(matches!( + resolver + .get_secret("oidc/env/MISSING", Some(Secret::Bool(true))) + .await, + Err(Error::MissingEnvironment) + )); + assert!(matches!( + resolver.get_secret("oidc/invalid", None).await, + Err(Error::InvalidOidc) + )); +} + +#[tokio::test] +async fn github_requests_are_authenticated_cached_and_revalidate_environment() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/token")) + .and(query_param("audience", "https://service/oidc/path")) + .and(header("authorization", "Bearer request-token")) + .and(header("accept", "application/json; api-version=2.0")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value":"identity-token"})), + ) + .expect(1) + .mount(&server) + .await; + let env = environment(&[ + ( + "ACTIONS_ID_TOKEN_REQUEST_URL", + &format!("{}/token", server.uri()), + ), + ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), + ]); + let oidc = OidcResolver::default(); + for _ in 0..2 { + assert_eq!( + oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "identity-token" + ); + } + assert!(matches!( + oidc.resolve( + "oidc/github/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await, + Err(Error::MissingEnvironment) + )); +} + +#[tokio::test] +async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explicit() { + let allowed = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + let token = allowed.path().join("token"); + let private = outside.path().join("private"); + std::fs::write(&token, "token\r\n").unwrap(); + std::fs::write(&private, "outside").unwrap(); + let env = environment(&[ + ( + "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", + allowed.path().to_str().unwrap(), + ), + ("PATH_TOKEN", private.to_str().unwrap()), + ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), + ]); + let oidc = OidcResolver::default(); + assert_eq!( + oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + assert!(matches!( + oidc.resolve("oidc/file/relative", env.as_ref()).await, + Err(Error::UnsafeOidcPath) + )); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", private.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + assert_eq!( + oidc.resolve("oidc/env_path/PATH_TOKEN", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "outside" + ); + assert_eq!( + oidc.resolve("oidc/azure/scope", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + #[cfg(unix)] + { + let link = allowed.path().join("link"); + std::os::unix::fs::symlink(&private, &link).unwrap(); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", link.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + } +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::at_refresh_boundary(serde_json::json!(1060), 2)] +#[case::beyond_refresh_boundary(serde_json::json!(1061), 1)] +#[case::already_expired(serde_json::json!(999), 2)] +#[case::string_expiry(serde_json::json!("999"), 2)] +#[case::fractional_expiry(serde_json::json!(1060.9), 2)] +#[case::negative_expiry(serde_json::json!(-1), 2)] +#[case::null_expiry(serde_json::Value::Null, 1)] +#[case::unreadable_expiry(serde_json::json!("invalid"), 1)] +#[case::nonfinite_expiry(serde_json::json!("NaN"), 1)] +#[tokio::test] +async fn google_expiry_caps_cache_and_preserves_audience( + #[case] expiry: serde_json::Value, + #[case] calls: u64, +) { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + fn now() -> SystemTime { + UNIX_EPOCH + Duration::from_secs(1000) + } + let server = MockServer::start().await; + let token = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(serde_json::json!({"alg":"RS256","typ":"JWT"}).to_string()), + URL_SAFE_NO_PAD.encode(serde_json::json!({"exp":expiry}).to_string()) + ); + Mock::given(method("GET")) + .and(header("metadata-flavor", "Google")) + .and(query_param("audience", "https://service/oidc/path")) + .respond_with(ResponseTemplate::new(200).set_body_string(&token)) + .expect(calls) + .mount(&server) + .await; + let oidc = + OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()).with_clock(now); + for _ in 0..2 { + assert_eq!( + oidc.resolve( + "oidc/google/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await + .unwrap() + .unwrap() + .expose(), + token + ); + } +} + +#[cfg(not(feature = "google"))] +#[tokio::test] +async fn google_oidc_requires_its_build_feature() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[tokio::test] +async fn azure_oidc_without_a_token_file_requires_an_unimplemented_backend() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/azure/scope", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[rstest::rstest] +#[case::missing_prefix("env/TOKEN", false)] +#[case::missing_audience_separator("oidc/env", false)] +#[case::unknown_provider("oidc/unknown/TOKEN", true)] +#[tokio::test] +async fn invalid_references_fail_before_environment_lookup( + #[case] reference: &str, + #[case] unsupported: bool, +) { + let error = OidcResolver::default() + .resolve(reference, &|_: &str| { + panic!("invalid reference reached environment lookup") + }) + .await + .unwrap_err(); + assert!(matches!(error, Error::UnsupportedOidc) == unsupported); + assert!(matches!(error, Error::InvalidOidc) != unsupported); +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::opaque("opaque-token")] +#[case::missing_expiry("header.e30.signature")] +#[tokio::test] +async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_string(token)) + .expect(1) + .mount(&server) + .await; + let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()); + for _ in 0..2 { + assert_eq!( + resolver + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + token, + ); + } +} diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs new file mode 100644 index 00000000000..e2b056f01cc --- /dev/null +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -0,0 +1,343 @@ +use std::sync::Arc; + +use litellm_secrets::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, OidcResolver, Secret, SecretManager, + SecretManagerState, SecretResolver, SecretValue, secret_manager_would_be_consulted, +}; + +fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { + let state = if readable { + SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings::default()), + Some(SecretManager::Local), + ) + .unwrap() + } else { + SecretManagerState::default() + }; + let value = value.map(str::to_owned); + SecretResolver::new( + Arc::new(state), + Arc::new(move |_: &str| value.clone()), + OidcResolver::default(), + ) +} + +#[rstest::rstest] +#[case::lowercase_true("true", Some(true), None)] +#[case::whitespace_lowercase_false(" FALSE ", Some(false), None)] +#[case::python_true("True", Some(true), Some(true))] +#[case::python_false("False", Some(false), Some(false))] +#[case::parenthesized_python_true("(True)", None, Some(true))] +#[case::commented_python_false("False # comment", None, Some(false))] +#[case::integer("1", None, None)] +#[case::yes("yes", None, None)] +#[case::plain_string("secret", None, None)] +#[tokio::test] +async fn boolean_conversion_preserves_local_and_manager_differences( + #[case] input: &str, + #[case] local: Option, + #[case] manager: Option, + #[values(false, true)] readable: bool, +) { + let boolean = if readable { manager } else { local }; + let resolver = resolver(Some(input), readable); + let expected = boolean + .map(Secret::Bool) + .unwrap_or_else(|| Secret::String(SecretValue::new(input))); + assert_eq!( + resolver.get_secret("key", None).await.unwrap(), + Some(expected) + ); + assert_eq!( + resolver + .get_secret_str("key", None) + .await + .unwrap() + .map(|v| v.expose().to_owned()), + boolean.is_none().then(|| input.to_owned()) + ); +} + +#[tokio::test] +async fn manager_boolean_conversion_trims_whitespace() { + assert_eq!( + resolver(Some(" true "), true) + .get_secret_bool("key", None) + .await + .unwrap(), + Some(true) + ); +} + +#[tokio::test] +async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { + let missing = resolver(None, false); + assert_eq!( + missing + .get_secret("missing", Some(Secret::Bool(true))) + .await + .unwrap(), + None + ); + assert_eq!( + missing + .get_secret_bool("missing", Some(true)) + .await + .unwrap(), + None + ); + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::default()), + Arc::new(|name: &str| (name == "KEY").then(|| "value".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str("os.environ/KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[rstest::rstest] +#[case::all_keys(None)] +#[case::no_keys(Some(Vec::new()))] +#[case::allowlisted_key(Some(vec!["KEY".into()]))] +fn manager_gating_requires_client_readable_settings_and_allowlisted_name( + #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] + access_mode: AccessMode, + #[values(false, true)] client: bool, + #[case] keys: Option>, +) { + let expected = client + && access_mode.readable() + && keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); + let state = SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings { + access_mode, + hosted_keys: keys, + ..Default::default() + }), + client.then_some(SecretManager::Local), + ) + .unwrap(); + assert_eq!( + secret_manager_would_be_consulted(&state, "os.environ/KEY"), + expected + ); +} + +#[test] +fn manager_gating_requires_settings() { + let no_settings = SecretManagerState::new(None, None, Some(SecretManager::Local)).unwrap(); + assert!(!secret_manager_would_be_consulted(&no_settings, "KEY")); +} + +#[cfg(feature = "aws")] +#[rstest::rstest] +#[case::missing_value(None, None)] +#[case::lookup_error(Some("primary".to_owned()), Some("environment-value"))] +#[tokio::test] +async fn aws_missing_values_do_not_fallback_but_lookup_errors_do( + #[case] primary: Option, + #[case] expected: Option<&str>, +) { + use litellm_secrets::aws::AwsSecretsManagerV2; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_partial_json}; + let server = MockServer::start().await; + Mock::given(body_partial_json(serde_json::json!({"SecretId":"KEY"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(u64::from(primary.is_none())) + .mount(&server) + .await; + Mock::given(body_partial_json(serde_json::json!({"SecretId":"primary"}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"SecretString":"invalid-json"})), + ) + .expect(u64::from(primary.is_some())) + .mount(&server) + .await; + let endpoint = server.uri(); + let environment: Arc = + Arc::new(move |name: &str| match name { + "AWS_REGION_NAME" => Some("us-east-1".into()), + "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), + "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), + "KEY" => Some("environment-value".into()), + _ => None, + }); + let settings = KeyManagementSettings { + primary_secret_name: primary, + ..Default::default() + }; + let manager = AwsSecretsManagerV2::load_aws_secret_manager( + Some(true), + settings.clone(), + environment.clone(), + ) + .unwrap() + .unwrap(); + let state = SecretManagerState::new( + Some(KeyManagementSystem::AwsSecretManager), + Some(settings), + Some(SecretManager::AwsSecretsManagerV2(manager)), + ) + .unwrap(); + let resolver = SecretResolver::new( + Arc::new(state), + environment.clone(), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str("os.environ/KEY", None) + .await + .unwrap() + .map(|v| v.expose().to_owned()) + .as_deref(), + expected + ); +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::hosted_filter(Some(Vec::new()), Some(KeyManagementSystem::GoogleSecretManager))] +#[case::negative_cache(None, Some(KeyManagementSystem::GoogleSecretManager))] +#[case::missing_system(None, None)] +#[case::hosted_nested_prefix(Some(vec!["os.environ/KEY".into()]), Some(KeyManagementSystem::GoogleSecretManager))] +#[tokio::test] +async fn google_negative_cache_still_falls_back_and_hosted_filter_avoids_io( + #[case] hosted_keys: Option>, + #[case] system: Option, +) { + use litellm_secrets::google::GoogleSecretManager; + use std::time::Duration; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path}; + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/os%2Eenviron%2FKEY/versions/latest:access", + )) + .respond_with(ResponseTemplate::new(404)) + .expect(u64::from( + hosted_keys.as_ref().is_none_or(|keys| !keys.is_empty()) && system.is_some(), + )) + .mount(&server) + .await; + let environment: Arc = + Arc::new(|name: &str| match name { + "VERTEX_AI_API_KEY" => Some("token".into()), + "os.environ/KEY" => Some("environment-value".into()), + _ => None, + }); + let manager = GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + environment.clone(), + Some(Duration::from_secs(60)), + false, + ) + .unwrap(); + let settings = KeyManagementSettings { + hosted_keys, + ..Default::default() + }; + let state = SecretManagerState::new( + system, + Some(settings), + Some(SecretManager::GoogleSecretManager(manager)), + ) + .unwrap(); + let resolver = SecretResolver::new( + Arc::new(state), + environment.clone(), + OidcResolver::default(), + ); + for _ in 0..2 { + assert_eq!( + resolver + .get_secret_str("os.environ/os.environ/KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "environment-value" + ); + } +} + +#[tokio::test] +async fn resolver_future_can_run_on_a_tokio_worker() { + let resolver = resolver(Some("worker-value"), false); + let result = tokio::spawn(async move { resolver.get_secret_str("KEY", None).await }) + .await + .unwrap() + .unwrap(); + assert_eq!(result.unwrap().expose(), "worker-value"); +} + +#[rstest::rstest] +#[case::nested_true("((True)) # comment", Some(true))] +#[case::commented_false("(False # comment\n)", Some(false))] +#[case::boolean_expression("True and False", None)] +#[case::string_literal("'True'", None)] +#[case::tuple("(True,)", None)] +#[case::unary_expression("not False", None)] +#[case::multiple_expressions("True\nFalse", None)] +#[case::incomplete_expression("(True", None)] +#[tokio::test] +async fn manager_boolean_literals_follow_python_syntax( + #[case] input: &str, + #[case] expected: Option, +) { + let value = resolver(Some(input), true) + .get_secret("key", None) + .await + .unwrap(); + assert_eq!( + value, + Some( + expected + .map(Secret::Bool) + .unwrap_or_else(|| Secret::String(SecretValue::new(input))) + ) + ); +} + +#[tokio::test] +async fn environment_prefix_is_removed_only_once_and_gating_uses_the_same_name() { + let name = "os.environ/folder/os.environ/KEY"; + let state = SecretManagerState::new( + Some(KeyManagementSystem::Local), + Some(KeyManagementSettings { + hosted_keys: Some(vec!["folder/os.environ/KEY".into()]), + ..Default::default() + }), + Some(SecretManager::Local), + ) + .unwrap(); + assert!(secret_manager_would_be_consulted(&state, name)); + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(|name: &str| (name == "folder/os.environ/KEY").then(|| "value".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str(name, None) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} From 15f368060ff0d420a86ec21bb818f0573c21e99b Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 23:30:54 +0000 Subject: [PATCH 112/127] chore(prices): sync OpenRouter prices: 3 models openrouter/~deepseek/deepseek-pro-latest: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/~z-ai/glm-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-pro-0813: off_peak_pricing, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 19 ++++++++++--------- model_prices_and_context_window.json | 19 ++++++++++--------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 39c418de524..172ec45147d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2404e-07, + "input_cost_per_token": 1.32e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7468e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 4.4e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7468e-08, - "input_cost_per_token": 5.2404e-07, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 39c418de524..172ec45147d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41579,22 +41579,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 5.2404e-07, + "input_cost_per_token": 1.32e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.7468e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.2404e-7,"output_cost_per_token":0.00000157212,"cache_read_input_token_cost":1.6716e-8}, + "cache_read_input_token_cost": 4.4e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -71317,14 +71317,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 1.7468e-08, - "input_cost_per_token": 5.2404e-07, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.57212e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "output_cost_per_token": 3.96e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71593,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 5ee75af90852d76da0bf6532c05368e8b984b609 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 23:34:40 +0000 Subject: [PATCH 113/127] fix(openrouter): keep glm-latest output limit at alias target value Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 172ec45147d..7089a6ffaa9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71594,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 172ec45147d..7089a6ffaa9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71594,8 +71594,8 @@ "input_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.86e-06, "source": "https://openrouter.ai/api/v1/models", From 7bcb01a40c7393b373ff8d4c517edde6ad68c865 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:00:57 +0000 Subject: [PATCH 114/127] chore(prices): sync OpenRouter prices: 4 models openrouter/~deepseek/deepseek-flash-latest: max_tokens, max_output_tokens openrouter/~z-ai/glm-flash-latest: max_tokens, max_output_tokens openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 20 +++++++++---------- model_prices_and_context_window.json | 20 +++++++++---------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7089a6ffaa9..e45e2db275a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 4.22298e-07, + "input_cost_per_token": 9.483e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.44596e-07, + "output_cost_per_token": 1.8966e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 3.51915e-08, + "cache_read_input_token_cost": 7.9025e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.556e-08, - "output_cost_per_token": 7.112e-08, - "cache_read_input_token_cost": 7.112e-09, + "input_cost_per_token": 8.9866e-08, + "output_cost_per_token": 1.79732e-07, + "cache_read_input_token_cost": 1.79732e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71300,8 +71300,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "output_cost_per_token": 1.2e-06, @@ -71574,8 +71574,8 @@ "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7089a6ffaa9..e45e2db275a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 4.22298e-07, + "input_cost_per_token": 9.483e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.44596e-07, + "output_cost_per_token": 1.8966e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 3.51915e-08, + "cache_read_input_token_cost": 7.9025e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 3.556e-08, - "output_cost_per_token": 7.112e-08, - "cache_read_input_token_cost": 7.112e-09, + "input_cost_per_token": 8.9866e-08, + "output_cost_per_token": 1.79732e-07, + "cache_read_input_token_cost": 1.79732e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -71300,8 +71300,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "output_cost_per_token": 1.2e-06, @@ -71574,8 +71574,8 @@ "input_cost_per_token": 9e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://openrouter.ai/api/v1/models", From a3ffc395b9399cb7b8ce63e9e66887dd376b5ad6 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:31:06 +0000 Subject: [PATCH 115/127] chore(prices): sync OpenRouter prices: 3 models openrouter/ibm-granite/granite-4.2-8b: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/meta-llama/llama-3.1-70b-instruct: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token openrouter/meta-llama/llama-4-maverick: input_cost_per_token, output_cost_per_token --- ...model_prices_and_context_window_backup.json | 18 +++++++++--------- model_prices_and_context_window.json | 18 +++++++++--------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e45e2db275a..a51cbdec882 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68617,8 +68617,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 1.875e-07, - "output_cost_per_token": 6.525e-07, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-07, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -72727,14 +72727,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e45e2db275a..a51cbdec882 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68617,8 +68617,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 1.875e-07, - "output_cost_per_token": 6.525e-07, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -69010,12 +69010,12 @@ "supports_reasoning": false }, "openrouter/meta-llama/llama-3.1-70b-instruct": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-07, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -72727,14 +72727,14 @@ "supports_web_search": false }, "openrouter/ibm-granite/granite-4.2-8b": { - "cache_read_input_token_cost": 1.5e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 117964, "max_tokens": 117964, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From acd303043b10d53cebb66002048afc06cabbd6f5 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:31:07 +0000 Subject: [PATCH 116/127] chore(prices): sync OpenRouter prices: 1 model openrouter/qwen/qwen3.8-27b: output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a51cbdec882..bdae7ede25c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66592,8 +66592,8 @@ }, "openrouter/qwen/qwen3.8-27b": { "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.55e-06, - "cache_read_input_token_cost": 8.5e-08, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a51cbdec882..bdae7ede25c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66592,8 +66592,8 @@ }, "openrouter/qwen/qwen3.8-27b": { "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.55e-06, - "cache_read_input_token_cost": 8.5e-08, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, From e3f69fc4d8e5348fdae1086d33c720b54df9a1be Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 19:08:25 -0700 Subject: [PATCH 117/127] refactor(rust): define consistent secret lookup contracts --- litellm-rust/Cargo.lock | 254 +------- litellm-rust/crates/secrets-aws/src/error.rs | 4 +- .../crates/secrets-aws/src/secret_manager.rs | 44 +- .../secrets-aws/tests/secret_manager.rs | 20 + .../secrets-google/src/secret_manager.rs | 26 +- .../secrets-google/tests/secret_manager.rs | 110 ++-- .../crates/secrets-types/src/value.rs | 9 +- litellm-rust/crates/secrets/Cargo.toml | 3 - litellm-rust/crates/secrets/README.md | 11 + litellm-rust/crates/secrets/src/error.rs | 10 +- litellm-rust/crates/secrets/src/handler.rs | 5 +- litellm-rust/crates/secrets/src/lib.rs | 2 +- litellm-rust/crates/secrets/src/oidc.rs | 7 +- litellm-rust/crates/secrets/src/resolver.rs | 135 ++--- litellm-rust/crates/secrets/src/state.rs | 90 +-- .../crates/secrets/tests/resolution.rs | 567 +++++++++--------- 16 files changed, 522 insertions(+), 775 deletions(-) create mode 100644 litellm-rust/crates/secrets/README.md diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 81fcabaf122..8c35a0be0b4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1515,15 +1515,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "getopts" -version = "0.2.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" -dependencies = [ - "unicode-width", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -2227,27 +2218,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "is-macro" -version = "0.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.119", -] - -[[package]] -name = "itertools" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.13.0" @@ -2394,12 +2364,6 @@ dependencies = [ "zeroize", ] -[[package]] -name = "lalrpop-util" -version = "0.20.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" - [[package]] name = "lazy_static" version = "1.5.0" @@ -2716,14 +2680,12 @@ dependencies = [ "moka", "reqwest 0.12.28", "rstest", - "rustpython-parser", "serde", "serde_json", "strum", "tempfile", "thiserror 2.0.19", "tokio", - "tracing", "wiremock", ] @@ -2809,7 +2771,7 @@ dependencies = [ "base64 0.22.1", "rand 0.8.7", "rstest", - "rustc-hash 2.1.3", + "rustc-hash", "serde", "serde_json", "thiserror 2.0.19", @@ -2985,16 +2947,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "num-bigint" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" -dependencies = [ - "num-integer", - "num-traits", -] - [[package]] name = "num-bigint" version = "0.5.1" @@ -3178,44 +3130,6 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_codegen" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" -dependencies = [ - "phf_generator", - "phf_shared", -] - -[[package]] -name = "phf_generator" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" -dependencies = [ - "phf_shared", - "rand 0.8.7", -] - -[[package]] -name = "phf_shared" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" -dependencies = [ - "siphasher", -] - [[package]] name = "pin-project" version = "1.1.13" @@ -3488,7 +3402,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.3", + "rustc-hash", "rustls 0.23.42", "socket2 0.6.5", "thiserror 2.0.19", @@ -3510,7 +3424,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash 2.1.3", + "rustc-hash", "rustls 0.23.42", "rustls-pki-types", "slab", @@ -3689,7 +3603,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint 0.5.1", + "num-bigint", "percent-encoding", "ryu", "sha1_smol", @@ -3918,12 +3832,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.3" @@ -4051,63 +3959,6 @@ dependencies = [ "untrusted", ] -[[package]] -name = "rustpython-ast" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" -dependencies = [ - "is-macro", - "num-bigint 0.4.8", - "rustpython-parser-core", - "static_assertions", -] - -[[package]] -name = "rustpython-parser" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" -dependencies = [ - "anyhow", - "is-macro", - "itertools 0.11.0", - "lalrpop-util", - "log", - "num-bigint 0.4.8", - "num-traits", - "phf", - "phf_codegen", - "rustc-hash 1.1.0", - "rustpython-ast", - "rustpython-parser-core", - "tiny-keccak", - "unic-emoji-char", - "unic-ucd-ident", - "unicode_names2", -] - -[[package]] -name = "rustpython-parser-core" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" -dependencies = [ - "is-macro", - "memchr", - "rustpython-parser-vendored", -] - -[[package]] -name = "rustpython-parser-vendored" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" -dependencies = [ - "memchr", - "once_cell", -] - [[package]] name = "rustversion" version = "1.0.23" @@ -4412,12 +4263,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" -[[package]] -name = "siphasher" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" - [[package]] name = "slab" version = "0.4.12" @@ -4648,7 +4493,7 @@ dependencies = [ "fancy-regex 0.17.0", "lazy_static", "regex", - "rustc-hash 2.1.3", + "rustc-hash", ] [[package]] @@ -4681,15 +4526,6 @@ dependencies = [ "time-core", ] -[[package]] -name = "tiny-keccak" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" -dependencies = [ - "crunchy", -] - [[package]] name = "tinystr" version = "0.8.3" @@ -5124,58 +4960,6 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" -[[package]] -name = "unic-char-property" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" -dependencies = [ - "unic-char-range", -] - -[[package]] -name = "unic-char-range" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" - -[[package]] -name = "unic-common" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" - -[[package]] -name = "unic-emoji-char" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-ident" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-version" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" -dependencies = [ - "unic-common", -] - [[package]] name = "unicase" version = "2.9.0" @@ -5203,40 +4987,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" -[[package]] -name = "unicode-width" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" - [[package]] name = "unicode_categories" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" -[[package]] -name = "unicode_names2" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" -dependencies = [ - "phf", - "unicode_names2_generator", -] - -[[package]] -name = "unicode_names2_generator" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" -dependencies = [ - "getopts", - "log", - "phf_codegen", - "rand 0.8.7", -] - [[package]] name = "untrusted" version = "0.9.0" diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs index 3f2c83a6a6d..23595397a13 100644 --- a/litellm-rust/crates/secrets-aws/src/error.rs +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -12,7 +12,7 @@ pub enum Error { Timeout, #[error("AWS KMS decrypt failed")] Decrypt(#[from] #[redact] Box>), - #[error("AWS Secrets Manager request preparation failed")] + #[error("AWS Secrets Manager read failed")] Read(#[from] #[redact] Box>), #[error("AWS Secrets Manager create failed")] Create(#[from] #[redact] Box>), @@ -22,6 +22,8 @@ pub enum Error { Delete(#[from] #[redact] Box>), #[error("AWS Secrets Manager replication failed")] Replicate(#[from] #[redact] Box>), + #[error("AWS Secrets Manager response has no string payload")] + MissingString, #[error("primary secret is not a JSON object")] PrimarySecret, #[error(transparent)] diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 97645ecac6d..493cb1d2e8f 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -106,23 +106,24 @@ impl AwsSecretsManagerV2 { } else { self.async_read_secret(primary).await? }; - let object: Value = serde_json::from_str( - value - .as_ref() - .map(SecretValue::expose) - .filter(|v| !v.is_empty()) - .unwrap_or("{}"), - ) - .map_err(|_| Error::PrimarySecret)?; + let Some(value) = value else { + return Ok(None); + }; + let object: Value = + serde_json::from_str(value.expose()).map_err(|_| Error::PrimarySecret)?; let object = object.as_object().ok_or(Error::PrimarySecret)?; - Ok(object.get(name).cloned().and_then(Secret::from_json)) + Ok(object.get(name).cloned().map(Secret::from_json)) } } } pub async fn async_read_secret(&self, name: &str) -> Result, Error> { match self.client.get_secret_value().secret_id(name).send().await { - Ok(response) => Ok(response.secret_string.map(SecretValue::new)), + Ok(response) => response + .secret_string + .map(SecretValue::new) + .map(Some) + .ok_or(Error::MissingString), Err(error) if matches!( &error, @@ -131,11 +132,14 @@ impl AwsSecretsManagerV2 { { Err(Error::Timeout) } - Err(error) if request_preparation_failed(&error) => Err(Error::Read(Box::new(error))), - Err(_) => { - tracing::error!("AWS secret read failed"); + Err(error) + if error + .as_service_error() + .is_some_and(|error| error.is_resource_not_found_exception()) => + { Ok(None) } + Err(error) => Err(Error::Read(Box::new(error))), } } @@ -281,17 +285,3 @@ fn bootstrap_key(name: &str) -> bool { | AWS_BEDROCK_RUNTIME_ENDPOINT ) } - -fn request_preparation_failed( - error: &aws_sdk_secretsmanager::error::SdkError< - aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError, - >, -) -> bool { - matches!( - error, - aws_sdk_secretsmanager::error::SdkError::ConstructionFailure(_) - ) || std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |error| { - error.source() - }) - .any(|source| source.is::()) -} diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs index 56698b391ec..a410767cb5a 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -290,3 +290,23 @@ async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { Err(Error::Timeout) )); } + +#[rstest::rstest] +#[case::denied(400, "AccessDeniedException")] +#[case::throttled(400, "ThrottlingException")] +#[case::unavailable(503, "ServiceUnavailableException")] +#[tokio::test] +async fn service_failures_remain_errors(#[case] status: u16, #[case] code: &str) { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(status).set_body_json(json!({"__type":code}))) + .expect(1) + .mount(&server) + .await; + assert!(matches!( + manager(&server, KeyManagementSettings::default()) + .async_read_secret("key") + .await, + Err(Error::Read(_)) + )); +} diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 3eb8475a367..3c34d9cbcc4 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -26,7 +26,7 @@ pub struct GoogleSecretManager { credentials: Arc, endpoint: reqwest::Url, project: String, - cache: Cache>, + cache: Cache, always_read: bool, } @@ -119,7 +119,7 @@ impl GoogleSecretManager { if !self.always_read && let Some(cached) = self.cache.get(name).await { - return Ok(cached.and_then(cached_secret)); + return Ok(Some(Secret::String(cached))); } let url = self .endpoint @@ -138,32 +138,20 @@ impl GoogleSecretManager { .headers(self.credentials.request_headers().await?) .send() .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } if response.status() != reqwest::StatusCode::OK { - self.cache.insert(name.to_owned(), None).await; return Err(Error::Status(response.status().as_u16())); } let response: Response = response.json().await?; let Some(data) = response.payload.and_then(|payload| payload.data) else { - self.cache.insert(name.to_owned(), None).await; return Err(Error::MissingPayload); }; - let filtered: String = data - .chars() - .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) - .collect(); - let bytes = STANDARD.decode(filtered)?; + let bytes = STANDARD.decode(data)?; let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; let value = SecretValue::new(plaintext); - self.cache - .insert(name.to_owned(), Some(value.clone())) - .await; + self.cache.insert(name.to_owned(), value.clone()).await; Ok(Some(Secret::String(value))) } } - -fn cached_secret(value: SecretValue) -> Option { - match serde_json::from_str(value.expose()) { - Ok(json) => Secret::from_json(json), - Err(_) => Some(Secret::String(value)), - } -} diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index 66cc0068ebb..b3b1d29e62c 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_secrets_google::{Error, GoogleSecretManager}; -use litellm_secrets_types::Secret; + use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{header, path}, @@ -55,32 +55,58 @@ async fn successful_reads_use_auth_latest_version_and_cache_including_empty_valu } #[rstest::rstest] -#[case::not_found(ResponseTemplate::new(404))] -#[case::missing_payload( - ResponseTemplate::new(200).set_body_json(serde_json::json!({"payload":{}})) -)] +#[case::not_found(404, serde_json::json!({}))] +#[case::unauthorized(401, serde_json::json!({}))] +#[case::forbidden(403, serde_json::json!({}))] +#[case::throttled(429, serde_json::json!({}))] +#[case::unavailable(503, serde_json::json!({}))] +#[case::missing_payload(200, serde_json::json!({"payload":{}}))] +#[case::invalid_base64(200, serde_json::json!({"payload":{"data":"%%%"}}))] #[tokio::test] -async fn negative_cache_returns_none_after_initial_error(#[case] response: ResponseTemplate) { +async fn failed_or_missing_reads_are_not_cached( + #[case] status: u16, + #[case] body: serde_json::Value, +) { let server = MockServer::start().await; + let manager = manager(&server, false, Duration::from_secs(60)); + let failing = Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount_as_scoped(&server) + .await; + let result = manager.get_secret_from_google_secret_manager("key").await; + match status { + 404 => assert_eq!(result.unwrap(), None), + 200 => assert!(matches!( + result, + Err(Error::MissingPayload | Error::Base64(_)) + )), + status => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)), + } + drop(failing); Mock::given(path( "/v1/projects/project/secrets/key/versions/latest:access", )) - .respond_with(response) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})), + ) .expect(1) .mount(&server) .await; - let manager = manager(&server, false, Duration::from_secs(60)); - assert!(matches!( - manager.get_secret_from_google_secret_manager("key").await, - Err(Error::Status(404) | Error::MissingPayload) - )); - assert!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap() - .is_none() - ); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some("recovered") + ); + } } #[rstest::rstest] @@ -130,21 +156,12 @@ fn google_manager_requires_host_license_and_project_configuration() { } #[rstest::rstest] -#[case::boolean("true", Some(Secret::Bool(true)))] -#[case::null("null", None)] -#[case::string( - "\"text\"", - Some(Secret::String(litellm_secrets_types::SecretValue::new("text"))) -)] -#[case::object( - "{\"key\":1}", - Secret::from_json(serde_json::json!({"key":1})) -)] +#[case("true")] +#[case("null")] +#[case("\"text\"")] +#[case("{\"key\":1}")] #[tokio::test] -async fn cached_values_preserve_python_json_conversion( - #[case] raw: &str, - #[case] expected: Option, -) { +async fn cache_preserves_raw_values(#[case] raw: &str) { let server = MockServer::start().await; Mock::given(path( "/v1/projects/project/secrets/key/versions/latest:access", @@ -157,20 +174,15 @@ async fn cached_values_preserve_python_json_conversion( .mount(&server) .await; let manager = manager(&server, false, Duration::from_secs(60)); - assert_eq!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap() - .unwrap() - .as_str(), - Some(raw) - ); - assert_eq!( - manager - .get_secret_from_google_secret_manager("key") - .await - .unwrap(), - expected - ); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some(raw) + ); + } } diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs index 524045ac007..087537fb3eb 100644 --- a/litellm-rust/crates/secrets-types/src/value.rs +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -14,12 +14,11 @@ impl From for Secret { } impl Secret { - pub fn from_json(value: serde_json::Value) -> Option { + pub fn from_json(value: serde_json::Value) -> Self { match value { - serde_json::Value::Null => None, - serde_json::Value::String(value) => Some(Self::String(SecretValue::new(value))), - serde_json::Value::Bool(value) => Some(Self::Bool(value)), - value => Some(Self::Json(value)), + serde_json::Value::String(value) => Self::String(SecretValue::new(value)), + serde_json::Value::Bool(value) => Self::Bool(value), + value => Self::Json(value), } } diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 3414470d234..a7e7ec80636 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -21,13 +21,10 @@ strum.workspace = true jsonwebtoken.workspace = true serde_json.workspace = true thiserror.workspace = true -tracing = "0.1" reqwest.workspace = true moka.workspace = true tokio = { workspace = true, features = ["fs"] } -rustpython-parser = { version = "0.4.0", default-features = false, features = ["num-bigint"] } - [dev-dependencies] rstest.workspace = true wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md new file mode 100644 index 00000000000..183a39e15bb --- /dev/null +++ b/litellm-rust/crates/secrets/README.md @@ -0,0 +1,11 @@ +# Secret resolution + +Construct `SecretManagerState::new(backend, settings)` for a configured manager or use `SecretManagerState::default()` for environment lookups. The configured backend determines its provider identity. Write-only settings and names excluded by `hosted_keys` use the environment directly. `secret_manager_would_be_consulted` follows the same routing decision as resolution + +`get_secret` returns `Ok(Some(value))` for a found value, `Ok(None)` when no source contains the value, and `Err(error)` when lookup fails. For managed names, resolution checks the manager, then the environment, then the caller's default. An empty string, `false`, or an explicitly stored JSON null is a found value + +Backend failures propagate by default. To allow fallback during a backend failure, construct the resolver with `.with_failure_policy(FailurePolicy::EnvironmentFallback)`. It then tries the environment and default, in that order. If neither exists, the original error is returned. This policy applies to manager lookups. Explicit OIDC references retain their own authentication errors and never fall back to environment secrets under the reference name + +`get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default + +Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index d240adb6a67..0c6e681b8aa 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -1,5 +1,3 @@ -use crate::KeyManagementSystem; - #[derive(Debug, thiserror::Error)] pub enum Error { #[error("encrypted environment value is missing")] @@ -8,10 +6,6 @@ pub enum Error { InvalidCiphertext, #[error("decrypted value is not UTF-8")] Utf8, - #[error("secret manager backend is not compiled: {0:?}")] - UnsupportedBackend(KeyManagementSystem), - #[error("configured secret manager does not match its backend")] - BackendMismatch, #[error("unsupported OIDC provider or missing build feature")] UnsupportedOidc, #[error("OIDC reference requires a provider and audience")] @@ -28,8 +22,8 @@ pub enum Error { UnsafeOidcPath, #[error("OIDC file could not be read")] OidcFile, - #[error("secret manager returned no secret")] - MissingSecret, + #[error("secret cannot be converted to {expected}")] + TypeMismatch { expected: &'static str }, #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 6d19965605d..943ffdf6158 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -76,9 +76,8 @@ pub async fn get_secret_from_manager( #[cfg(feature = "google")] SecretManager::GoogleSecretManager(client) => client .get_secret_from_google_secret_manager(secret_name) - .await? - .map(Some) - .ok_or(Error::MissingSecret), + .await + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index c434cfaaf03..ff2e95f7b2f 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -12,7 +12,7 @@ pub use litellm_secrets_types::{ AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, }; pub use oidc::{OidcProvider, OidcReference, OidcResolver}; -pub use resolver::SecretResolver; +pub use resolver::{FailurePolicy, SecretResolver}; pub use state::{SecretManagerState, secret_manager_would_be_consulted}; #[cfg(feature = "aws")] diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index 48e6a9bc0b0..fd477859bf6 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -88,8 +88,13 @@ pub struct OidcResolver { impl Default for OidcResolver { fn default() -> Self { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(600)) + .connect_timeout(Duration::from_secs(5)) + .build() + .expect("HTTP client configuration"); Self::new( - reqwest::Client::builder().timeout(Duration::from_secs(600)).connect_timeout(Duration::from_secs(5)).build().expect("HTTP client configuration"), + client, reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), ) } diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 93daddae0c4..89439893852 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -2,14 +2,21 @@ use std::sync::Arc; use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use crate::state::{LookupTarget, normalize_secret_name}; use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; -use crate::state::{LookupTarget, normalize_secret_name}; +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum FailurePolicy { + #[default] + Propagate, + EnvironmentFallback, +} pub struct SecretResolver { state: Arc, environment: Arc, oidc: OidcResolver, + failure_policy: FailurePolicy, } impl Default for SecretResolver { @@ -32,13 +39,21 @@ impl SecretResolver { state, environment, oidc, + failure_policy: FailurePolicy::default(), + } + } + + pub fn with_failure_policy(self, failure_policy: FailurePolicy) -> Self { + Self { + failure_policy, + ..self } } pub async fn get_secret( &self, name: &str, - _default_value: Option, + default_value: Option, ) -> Result, Error> { let name = normalize_secret_name(name); if name.starts_with("oidc/") { @@ -46,32 +61,26 @@ impl SecretResolver { .oidc .resolve(name, self.environment.as_ref()) .await - .map(|value| value.map(Secret::String)); + .map(|value| value.map(Secret::String).or(default_value)); } - if !self.state.readable() { - return Ok(self - .environment - .get(name) - .map(|value| match str_to_bool(&value) { - Some(value) => Secret::Bool(value), - None => Secret::String(SecretValue::new(value)), - })); + let LookupTarget::Manager { backend, settings } = self.state.lookup_target(name) else { + return Ok(self.environment_secret(name).or(default_value)); + }; + match crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) + .await + { + Ok(value) => Ok(value + .or_else(|| self.environment_secret(name)) + .or(default_value)), + Err(error) => match self.failure_policy { + FailurePolicy::Propagate => Err(error), + FailurePolicy::EnvironmentFallback => self + .environment_secret(name) + .or(default_value) + .map(Some) + .ok_or(error), + }, } - let result = match self.state.lookup_target(name) { - LookupTarget::Environment => Ok(self.environment_secret(name)), - LookupTarget::Manager { backend, settings } => { - crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref()) - .await - } - }; - let value = match result { - Ok(value) => value, - Err(_) => { - tracing::error!("secret manager lookup failed; falling back to environment"); - self.environment_secret(name) - } - }; - Ok(value.and_then(managed_secret)) } fn environment_secret(&self, name: &str) -> Option { @@ -84,12 +93,18 @@ impl SecretResolver { pub async fn get_secret_str( &self, name: &str, - default_value: Option, + default_value: Option, ) -> Result, Error> { - Ok(match self.get_secret(name, default_value).await? { - Some(Secret::String(value)) => Some(value), - Some(Secret::Bool(_) | Secret::Json(_)) | None => None, - }) + match self + .get_secret(name, default_value.map(Secret::String)) + .await? + { + Some(Secret::String(value)) => Ok(Some(value)), + None => Ok(None), + Some(Secret::Bool(_) | Secret::Json(_)) => { + Err(Error::TypeMismatch { expected: "string" }) + } + } } pub async fn get_secret_bool( @@ -97,44 +112,24 @@ impl SecretResolver { name: &str, default_value: Option, ) -> Result, Error> { - Ok( - match self - .get_secret(name, default_value.map(Secret::Bool)) - .await? - { - Some(Secret::Bool(value)) => Some(value), - Some(Secret::String(value)) => str_to_bool(value.expose()), - Some(Secret::Json(_)) | None => None, - }, - ) - } -} - -fn str_to_bool(value: &str) -> Option { - match value.trim().to_ascii_lowercase().as_str() { - "true" => Some(true), - "false" => Some(false), - _ => None, - } -} - -fn literal_bool(value: &str) -> Option { - use rustpython_parser::{Parse, ast}; - match ast::Expr::parse(value.trim_start_matches([' ', '\t']), "").ok()? { - ast::Expr::Constant(node) => match node.value { - ast::Constant::Bool(value) => Some(value), - _ => None, - }, - _ => None, - } -} - -fn managed_secret(value: Secret) -> Option { - match value { - Secret::String(value) => Some(match literal_bool(value.expose()) { - Some(boolean) => Secret::Bool(boolean), - None => Secret::String(value), - }), - Secret::Bool(_) | Secret::Json(_) => None, + match self + .get_secret(name, default_value.map(Secret::Bool)) + .await? + { + Some(Secret::Bool(value)) => Ok(Some(value)), + Some(Secret::String(value)) => { + match value.expose().trim().to_ascii_lowercase().as_str() { + "true" => Ok(Some(true)), + "false" => Ok(Some(false)), + _ => Err(Error::TypeMismatch { + expected: "boolean", + }), + } + } + Some(Secret::Json(_)) => Err(Error::TypeMismatch { + expected: "boolean", + }), + None => Ok(None), + } } } diff --git a/litellm-rust/crates/secrets/src/state.rs b/litellm-rust/crates/secrets/src/state.rs index ca942aa65db..7854d763ef2 100644 --- a/litellm-rust/crates/secrets/src/state.rs +++ b/litellm-rust/crates/secrets/src/state.rs @@ -1,4 +1,4 @@ -use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; +use crate::{KeyManagementSettings, KeyManagementSystem, SecretManager}; pub(crate) enum LookupTarget<'a> { Environment, @@ -14,73 +14,37 @@ pub(crate) fn normalize_secret_name(name: &str) -> &str { #[derive(Clone, Default)] pub struct SecretManagerState { - system: Option, - settings: Option, - backend: Option, + manager: Option<(SecretManager, KeyManagementSettings)>, } impl SecretManagerState { - pub fn new( - system: Option, - settings: Option, - backend: Option, - ) -> Result { - if let Some(system) = system { - let available = match system { - KeyManagementSystem::Local => true, - KeyManagementSystem::AwsKms | KeyManagementSystem::AwsSecretManager => { - cfg!(feature = "aws") - } - KeyManagementSystem::GoogleKms | KeyManagementSystem::GoogleSecretManager => { - cfg!(feature = "google") - } - KeyManagementSystem::AzureKeyVault - | KeyManagementSystem::HashicorpVault - | KeyManagementSystem::Cyberark - | KeyManagementSystem::Custom => false, - }; - if !available { - return Err(Error::UnsupportedBackend(system)); - } - if let Some(backend) = &backend - && system != backend.system() - { - return Err(Error::BackendMismatch); - } + pub fn new(backend: SecretManager, settings: KeyManagementSettings) -> Self { + Self { + manager: Some((backend, settings)), } - Ok(Self { - system, - settings, - backend, - }) } pub fn system(&self) -> Option { - self.system - } - pub fn settings(&self) -> Option<&KeyManagementSettings> { - self.settings.as_ref() - } - pub fn backend(&self) -> Option<&SecretManager> { - self.backend.as_ref() + self.backend().map(SecretManager::system) } - pub(crate) fn readable(&self) -> bool { - self.backend.is_some() - && self - .settings - .as_ref() - .is_some_and(|settings| settings.access_mode.readable()) + pub fn settings(&self) -> Option<&KeyManagementSettings> { + self.manager.as_ref().map(|(_, settings)| settings) + } + + pub fn backend(&self) -> Option<&SecretManager> { + self.manager.as_ref().map(|(backend, _)| backend) } pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> { - match (&self.backend, &self.settings) { - (Some(backend), Some(settings)) - if settings.access_mode.readable() - && hosts_secret(settings, name) - && self - .system - .is_some_and(|system| system != KeyManagementSystem::Local) => + match &self.manager { + Some((backend, settings)) + if backend.system() != KeyManagementSystem::Local + && settings.access_mode.readable() + && settings + .hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == name)) => { LookupTarget::Manager { backend, settings } } @@ -90,16 +54,6 @@ impl SecretManagerState { } pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool { - state.readable() - && state - .settings - .as_ref() - .is_some_and(|settings| hosts_secret(settings, normalize_secret_name(name))) -} - -fn hosts_secret(settings: &KeyManagementSettings, name: &str) -> bool { - settings - .hosted_keys - .as_ref() - .is_none_or(|keys| keys.iter().any(|key| key == name)) + let name = normalize_secret_name(name); + !name.starts_with("oidc/") && matches!(state.lookup_target(name), LookupTarget::Manager { .. }) } diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index e2b056f01cc..3a826092d72 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -1,18 +1,13 @@ use std::sync::Arc; use litellm_secrets::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, OidcResolver, Secret, SecretManager, - SecretManagerState, SecretResolver, SecretValue, secret_manager_would_be_consulted, + Error, KeyManagementSettings, OidcResolver, Secret, SecretManager, SecretManagerState, + SecretResolver, SecretValue, secret_manager_would_be_consulted, }; -fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { - let state = if readable { - SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings::default()), - Some(SecretManager::Local), - ) - .unwrap() +fn resolver(value: Option<&str>, configured: bool) -> SecretResolver { + let state = if configured { + SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()) } else { SecretManagerState::default() }; @@ -25,77 +20,107 @@ fn resolver(value: Option<&str>, readable: bool) -> SecretResolver { } #[rstest::rstest] -#[case::lowercase_true("true", Some(true), None)] -#[case::whitespace_lowercase_false(" FALSE ", Some(false), None)] -#[case::python_true("True", Some(true), Some(true))] -#[case::python_false("False", Some(false), Some(false))] -#[case::parenthesized_python_true("(True)", None, Some(true))] -#[case::commented_python_false("False # comment", None, Some(false))] -#[case::integer("1", None, None)] -#[case::yes("yes", None, None)] -#[case::plain_string("secret", None, None)] +#[case("true", Some(true))] +#[case(" FALSE ", Some(false))] +#[case("(True)", None)] +#[case("False # comment", None)] +#[case("1", None)] +#[case("secret", None)] #[tokio::test] -async fn boolean_conversion_preserves_local_and_manager_differences( +async fn conversion_is_explicit_and_independent_of_manager_configuration( #[case] input: &str, - #[case] local: Option, - #[case] manager: Option, - #[values(false, true)] readable: bool, + #[case] boolean: Option, + #[values(false, true)] configured: bool, ) { - let boolean = if readable { manager } else { local }; - let resolver = resolver(Some(input), readable); - let expected = boolean - .map(Secret::Bool) - .unwrap_or_else(|| Secret::String(SecretValue::new(input))); + let resolver = resolver(Some(input), configured); assert_eq!( resolver.get_secret("key", None).await.unwrap(), - Some(expected) + Some(Secret::String(SecretValue::new(input))) ); assert_eq!( resolver .get_secret_str("key", None) .await .unwrap() - .map(|v| v.expose().to_owned()), - boolean.is_none().then(|| input.to_owned()) + .unwrap() + .expose(), + input + ); + match boolean { + Some(value) => assert_eq!( + resolver.get_secret_bool("key", None).await.unwrap(), + Some(value) + ), + None => assert!(matches!( + resolver.get_secret_bool("key", Some(true)).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), + } +} + +#[rstest::rstest] +#[tokio::test] +async fn defaults_apply_only_to_absence(#[values(false, true)] configured: bool) { + let missing = resolver(None, configured); + assert_eq!(missing.get_secret("key", None).await.unwrap(), None); + assert_eq!( + missing.get_secret_bool("key", Some(false)).await.unwrap(), + Some(false) + ); + assert_eq!( + missing + .get_secret_str("key", Some(SecretValue::new("default"))) + .await + .unwrap() + .unwrap() + .expose(), + "default" + ); + for value in [ + Secret::Bool(false), + Secret::from_json(serde_json::json!({"key":1})), + Secret::from_json(serde_json::Value::Null), + ] { + assert_eq!( + missing + .get_secret("key", Some(value.clone())) + .await + .unwrap(), + Some(value) + ); + } + assert_eq!( + resolver(Some(""), configured) + .get_secret_str("key", Some(SecretValue::new("default"))) + .await + .unwrap() + .unwrap() + .expose(), + "" ); } #[tokio::test] -async fn manager_boolean_conversion_trims_whitespace() { +async fn prefix_is_removed_once_and_local_manager_is_not_consulted() { + let state = SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()); assert_eq!( - resolver(Some(" true "), true) - .get_secret_bool("key", None) - .await - .unwrap(), - Some(true) - ); -} - -#[tokio::test] -async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { - let missing = resolver(None, false); - assert_eq!( - missing - .get_secret("missing", Some(Secret::Bool(true))) - .await - .unwrap(), - None - ); - assert_eq!( - missing - .get_secret_bool("missing", Some(true)) - .await - .unwrap(), - None + state.system(), + Some(litellm_secrets::KeyManagementSystem::Local) ); + assert!(!secret_manager_would_be_consulted( + &state, + "os.environ/os.environ/KEY" + )); let resolver = SecretResolver::new( - Arc::new(SecretManagerState::default()), - Arc::new(|name: &str| (name == "KEY").then(|| "value".into())), + Arc::new(state), + Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), OidcResolver::default(), ); assert_eq!( resolver - .get_secret_str("os.environ/KEY", None) + .get_secret_str("os.environ/os.environ/KEY", None) .await .unwrap() .unwrap() @@ -104,177 +129,6 @@ async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() { ); } -#[rstest::rstest] -#[case::all_keys(None)] -#[case::no_keys(Some(Vec::new()))] -#[case::allowlisted_key(Some(vec!["KEY".into()]))] -fn manager_gating_requires_client_readable_settings_and_allowlisted_name( - #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] - access_mode: AccessMode, - #[values(false, true)] client: bool, - #[case] keys: Option>, -) { - let expected = client - && access_mode.readable() - && keys - .as_ref() - .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); - let state = SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings { - access_mode, - hosted_keys: keys, - ..Default::default() - }), - client.then_some(SecretManager::Local), - ) - .unwrap(); - assert_eq!( - secret_manager_would_be_consulted(&state, "os.environ/KEY"), - expected - ); -} - -#[test] -fn manager_gating_requires_settings() { - let no_settings = SecretManagerState::new(None, None, Some(SecretManager::Local)).unwrap(); - assert!(!secret_manager_would_be_consulted(&no_settings, "KEY")); -} - -#[cfg(feature = "aws")] -#[rstest::rstest] -#[case::missing_value(None, None)] -#[case::lookup_error(Some("primary".to_owned()), Some("environment-value"))] -#[tokio::test] -async fn aws_missing_values_do_not_fallback_but_lookup_errors_do( - #[case] primary: Option, - #[case] expected: Option<&str>, -) { - use litellm_secrets::aws::AwsSecretsManagerV2; - use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_partial_json}; - let server = MockServer::start().await; - Mock::given(body_partial_json(serde_json::json!({"SecretId":"KEY"}))) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) - .expect(u64::from(primary.is_none())) - .mount(&server) - .await; - Mock::given(body_partial_json(serde_json::json!({"SecretId":"primary"}))) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"SecretString":"invalid-json"})), - ) - .expect(u64::from(primary.is_some())) - .mount(&server) - .await; - let endpoint = server.uri(); - let environment: Arc = - Arc::new(move |name: &str| match name { - "AWS_REGION_NAME" => Some("us-east-1".into()), - "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), - "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), - "KEY" => Some("environment-value".into()), - _ => None, - }); - let settings = KeyManagementSettings { - primary_secret_name: primary, - ..Default::default() - }; - let manager = AwsSecretsManagerV2::load_aws_secret_manager( - Some(true), - settings.clone(), - environment.clone(), - ) - .unwrap() - .unwrap(); - let state = SecretManagerState::new( - Some(KeyManagementSystem::AwsSecretManager), - Some(settings), - Some(SecretManager::AwsSecretsManagerV2(manager)), - ) - .unwrap(); - let resolver = SecretResolver::new( - Arc::new(state), - environment.clone(), - OidcResolver::default(), - ); - assert_eq!( - resolver - .get_secret_str("os.environ/KEY", None) - .await - .unwrap() - .map(|v| v.expose().to_owned()) - .as_deref(), - expected - ); -} - -#[cfg(feature = "google")] -#[rstest::rstest] -#[case::hosted_filter(Some(Vec::new()), Some(KeyManagementSystem::GoogleSecretManager))] -#[case::negative_cache(None, Some(KeyManagementSystem::GoogleSecretManager))] -#[case::missing_system(None, None)] -#[case::hosted_nested_prefix(Some(vec!["os.environ/KEY".into()]), Some(KeyManagementSystem::GoogleSecretManager))] -#[tokio::test] -async fn google_negative_cache_still_falls_back_and_hosted_filter_avoids_io( - #[case] hosted_keys: Option>, - #[case] system: Option, -) { - use litellm_secrets::google::GoogleSecretManager; - use std::time::Duration; - use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path}; - let server = MockServer::start().await; - Mock::given(path( - "/v1/projects/project/secrets/os%2Eenviron%2FKEY/versions/latest:access", - )) - .respond_with(ResponseTemplate::new(404)) - .expect(u64::from( - hosted_keys.as_ref().is_none_or(|keys| !keys.is_empty()) && system.is_some(), - )) - .mount(&server) - .await; - let environment: Arc = - Arc::new(|name: &str| match name { - "VERTEX_AI_API_KEY" => Some("token".into()), - "os.environ/KEY" => Some("environment-value".into()), - _ => None, - }); - let manager = GoogleSecretManager::with_client( - reqwest::Client::new(), - server.uri().parse().unwrap(), - "project".into(), - environment.clone(), - Some(Duration::from_secs(60)), - false, - ) - .unwrap(); - let settings = KeyManagementSettings { - hosted_keys, - ..Default::default() - }; - let state = SecretManagerState::new( - system, - Some(settings), - Some(SecretManager::GoogleSecretManager(manager)), - ) - .unwrap(); - let resolver = SecretResolver::new( - Arc::new(state), - environment.clone(), - OidcResolver::default(), - ); - for _ in 0..2 { - assert_eq!( - resolver - .get_secret_str("os.environ/os.environ/KEY", None) - .await - .unwrap() - .unwrap() - .expose(), - "environment-value" - ); - } -} - #[tokio::test] async fn resolver_future_can_run_on_a_tokio_worker() { let resolver = resolver(Some("worker-value"), false); @@ -285,59 +139,230 @@ async fn resolver_future_can_run_on_a_tokio_worker() { assert_eq!(result.unwrap().expose(), "worker-value"); } -#[rstest::rstest] -#[case::nested_true("((True)) # comment", Some(true))] -#[case::commented_false("(False # comment\n)", Some(false))] -#[case::boolean_expression("True and False", None)] -#[case::string_literal("'True'", None)] -#[case::tuple("(True,)", None)] -#[case::unary_expression("not False", None)] -#[case::multiple_expressions("True\nFalse", None)] -#[case::incomplete_expression("(True", None)] -#[tokio::test] -async fn manager_boolean_literals_follow_python_syntax( - #[case] input: &str, - #[case] expected: Option, -) { - let value = resolver(Some(input), true) - .get_secret("key", None) - .await - .unwrap(); - assert_eq!( - value, - Some( - expected - .map(Secret::Bool) - .unwrap_or_else(|| Secret::String(SecretValue::new(input))) +#[cfg(feature = "aws")] +mod aws { + use super::*; + use litellm_secrets::{AccessMode, FailurePolicy, aws::AwsSecretsManagerV2}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + + fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState { + let endpoint = server.uri(); + let environment = Arc::new(move |name: &str| match name { + "AWS_REGION_NAME" => Some("us-east-1".into()), + "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()), + "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()), + _ => None, + }); + let manager = + AwsSecretsManagerV2::load_aws_secret_manager(Some(true), settings.clone(), environment) + .unwrap() + .unwrap(); + SecretManagerState::new(SecretManager::AwsSecretsManagerV2(manager), settings) + } + + #[rstest::rstest] + #[case::missing(400, serde_json::json!({"__type":"ResourceNotFoundException"}), false)] + #[case::denied(400, serde_json::json!({"__type":"AccessDeniedException"}), true)] + #[case::malformed(200, serde_json::json!({}), true)] + #[tokio::test] + async fn failure_policy_preserves_errors_and_fallback_precedence( + #[case] status: u16, + #[case] body: serde_json::Value, + #[case] fails: bool, + #[values(FailurePolicy::Propagate, FailurePolicy::EnvironmentFallback)] + policy: FailurePolicy, + #[values(None, Some("environment"))] environment: Option<&'static str>, + #[values(None, Some("default"))] default: Option<&str>, + ) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount(&server) + .await; + let resolver = SecretResolver::new( + Arc::new(state(&server, KeyManagementSettings::default())), + Arc::new(move |_: &str| environment.map(str::to_owned)), + OidcResolver::default(), ) - ); + .with_failure_policy(policy); + let result = resolver + .get_secret_str("KEY", default.map(SecretValue::new)) + .await; + let fallback = environment.or(default); + if fails && (policy == FailurePolicy::Propagate || fallback.is_none()) { + assert!(matches!(result, Err(Error::Aws(_)))); + } else { + assert_eq!(result.unwrap().as_ref().map(SecretValue::expose), fallback); + } + } + + #[rstest::rstest] + #[case::boolean(serde_json::json!(false))] + #[case::object(serde_json::json!({"key":1}))] + #[case::null(serde_json::Value::Null)] + #[case::string(serde_json::json!("true"))] + #[tokio::test] + async fn typed_values_survive_resolution_and_accessors_reject_wrong_types( + #[case] value: serde_json::Value, + ) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"SecretString":serde_json::json!({"KEY":value}).to_string()}), + )) + .expect(3) + .mount(&server) + .await; + let settings = KeyManagementSettings { + primary_secret_name: Some("primary".into()), + ..Default::default() + }; + let resolver = SecretResolver::new( + Arc::new(state(&server, settings)), + Arc::new(|_: &str| Some("fallback".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret("KEY", Some(Secret::Bool(true))) + .await + .unwrap(), + Some(Secret::from_json(value.clone())) + ); + match &value { + serde_json::Value::String(text) => assert_eq!( + resolver + .get_secret_str("KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + text + ), + _ => assert!(matches!( + resolver.get_secret_str("KEY", None).await, + Err(Error::TypeMismatch { expected: "string" }) + )), + } + match value { + serde_json::Value::Bool(boolean) => assert_eq!( + resolver.get_secret_bool("KEY", None).await.unwrap(), + Some(boolean) + ), + serde_json::Value::String(_) => assert_eq!( + resolver.get_secret_bool("KEY", None).await.unwrap(), + Some(true) + ), + _ => assert!(matches!( + resolver.get_secret_bool("KEY", None).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn gating_prediction_matches_actual_lookup( + #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)] + access_mode: AccessMode, + #[values(None, Some(vec![]), Some(vec!["KEY".into()]))] hosted_keys: Option>, + #[values("os.environ/KEY", "os.environ/oidc/env/KEY")] name: &str, + ) { + let server = MockServer::start().await; + let expected = name == "os.environ/KEY" + && access_mode.readable() + && hosted_keys + .as_ref() + .is_none_or(|keys| keys.iter().any(|key| key == "KEY")); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"SecretString":"remote"})), + ) + .expect(u64::from(expected)) + .mount(&server) + .await; + let state = state( + &server, + KeyManagementSettings { + access_mode, + hosted_keys, + ..Default::default() + }, + ); + assert!(state.backend().is_some()); + assert_eq!(state.settings().unwrap().access_mode, access_mode); + assert_eq!(secret_manager_would_be_consulted(&state, name), expected); + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(|_: &str| Some("environment".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str(name, None) + .await + .unwrap() + .unwrap() + .expose(), + if expected { "remote" } else { "environment" } + ); + } } +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::missing(404)] +#[case::failure(503)] #[tokio::test] -async fn environment_prefix_is_removed_only_once_and_gating_uses_the_same_name() { - let name = "os.environ/folder/os.environ/KEY"; - let state = SecretManagerState::new( - Some(KeyManagementSystem::Local), - Some(KeyManagementSettings { - hosted_keys: Some(vec!["folder/os.environ/KEY".into()]), - ..Default::default() - }), - Some(SecretManager::Local), +async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) { + use litellm_secrets::{FailurePolicy, google::GoogleSecretManager}; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(status)) + .expect(2) + .mount(&server) + .await; + let environment: Arc = + Arc::new(|name: &str| match name { + "VERTEX_AI_API_KEY" => Some("token".into()), + "KEY" => Some("environment".into()), + _ => None, + }); + let manager = GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + environment.clone(), + None, + false, ) .unwrap(); - assert!(secret_manager_would_be_consulted(&state, name)); - let resolver = SecretResolver::new( - Arc::new(state), - Arc::new(|name: &str| (name == "folder/os.environ/KEY").then(|| "value".into())), - OidcResolver::default(), + let state = SecretManagerState::new( + SecretManager::GoogleSecretManager(manager), + KeyManagementSettings::default(), ); + let resolver = SecretResolver::new(Arc::new(state), environment, OidcResolver::default()); + let result = resolver.get_secret_str("KEY", None).await; + if status == 404 { + assert_eq!(result.unwrap().unwrap().expose(), "environment"); + } else { + assert!( + matches!(result, Err(Error::Google(litellm_secrets::google::Error::Status(actual))) if actual == status) + ); + } assert_eq!( resolver - .get_secret_str(name, None) + .with_failure_policy(FailurePolicy::EnvironmentFallback) + .get_secret_str("KEY", None) .await .unwrap() .unwrap() .expose(), - "value" + "environment" ); } From 2bacd6caa0154c46461ec3f790e4b438052156af Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:30:52 +0000 Subject: [PATCH 118/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~moonshotai/kimi-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/moonshotai/kimi-k3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bdae7ede25c..1095ec9e467 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.675e-06, + "output_cost_per_token": 9.38e-06, + "cache_read_input_token_cost": 1.943e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.675e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 9.38e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bdae7ede25c..1095ec9e467 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.7e-06, - "output_cost_per_token": 8.5e-06, - "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.675e-06, + "output_cost_per_token": 9.38e-06, + "cache_read_input_token_cost": 1.943e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.7e-07, - "input_cost_per_token": 1.7e-06, + "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.675e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 8.5e-06, + "output_cost_per_token": 9.38e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 4258bd366c7c6342d1e29641d060823c7cea7a40 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:00:39 +0000 Subject: [PATCH 119/127] chore(prices): sync OpenRouter prices: 2 models openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1095ec9e467..ba001cd3923 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.483e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8966e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9025e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.9866e-08, - "output_cost_per_token": 1.79732e-07, - "cache_read_input_token_cost": 1.79732e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1095ec9e467..ba001cd3923 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.483e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8966e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9025e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.9866e-08, - "output_cost_per_token": 1.79732e-07, - "cache_read_input_token_cost": 1.79732e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From d5921713a8105804917134614a28457bc7fc552c Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:30:42 +0000 Subject: [PATCH 120/127] chore(prices): sync OpenRouter prices: 2 models openrouter/~moonshotai/kimi-latest: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/moonshotai/kimi-k3: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ba001cd3923..1976437f900 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.675e-06, - "output_cost_per_token": 9.38e-06, - "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.943e-07, - "input_cost_per_token": 1.675e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 9.38e-06, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ba001cd3923..1976437f900 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66774,9 +66774,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 1.675e-06, - "output_cost_per_token": 9.38e-06, - "cache_read_input_token_cost": 1.943e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -71407,14 +71407,14 @@ "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 1.943e-07, - "input_cost_per_token": 1.675e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 9.38e-06, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, From 83223885e6d4f9dd536f73214b446e1f14b64244 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 04:09:59 +0000 Subject: [PATCH 121/127] fix(responses): forward safety_identifier through the chat completion bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 1 + tests/e2e/llm_translation/endpoints_client.py | 3 + .../e2e/llm_translation/test_responses_e2e.py | 89 ++++++++++++++++++- tests/e2e/models.py | 2 + tests/e2e/provider_edge.py | 10 ++- .../test_litellm_completion_responses.py | 10 +++ 6 files changed, 109 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index cf3075ee28d..3ca2cc28c9a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -454,6 +454,7 @@ class LiteLLMCompletionResponsesConfig: "stream": stream, "metadata": kwargs.get("metadata"), "service_tier": kwargs.get("service_tier"), + "safety_identifier": responses_api_request.get("safety_identifier"), "web_search_options": web_search_options, "response_format": response_format, "reasoning_effort": reasoning.effort, diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 4d2c73e7078..eb7bae2220c 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -75,6 +75,7 @@ class ResponsesRequest(BaseModel): stream: bool = False tools: list[ResponsesFunctionTool] | None = None guardrails: list[str] | None = None + safety_identifier: str | None = None cache: dict[str, bool] | None = {"no-cache": True} @@ -316,6 +317,7 @@ class EndpointsClient: *, stream: bool = False, guardrails: list[str] | None = None, + safety_identifier: str | None = None, ) -> StreamingResponse: return self._send( "/v1/responses", @@ -326,6 +328,7 @@ class EndpointsClient: instructions="You are a helpful assistant", stream=stream, guardrails=guardrails, + safety_identifier=safety_identifier, ), stream=stream, ) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 3fcf2d1ac05..083bc66b3f2 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -8,13 +8,18 @@ litellm-regression-tests/tests/test_inference_endpoints.py. from __future__ import annotations import json -from typing import cast +import threading +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Final, cast import pytest -from e2e_config import unique_marker +from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, unique_marker from e2e_http import ( assert_client_error, require_successful_call, + unwrap, ) from endpoints_client import ( EndpointsClient, @@ -26,7 +31,9 @@ from endpoints_client import ( ResponsesStreamEventType, ) from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import ChatBody, ChatMessage, LiteLLMParamsBody +from provider_edge import LiveEdge, start_provider_edge +from provider_edge_bedrock import bedrock_signer from pydantic import BaseModel, ValidationError pytestmark = pytest.mark.e2e @@ -39,6 +46,33 @@ class _OptionalResponsesBody(BaseModel): BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +BEDROCK_EDGE_REGION: Final = "us-east-1" +BEDROCK_EDGE_MOUNT: Final = f"bedrock/{BEDROCK_EDGE_REGION}" + + +class ConverseRequestBody(BaseModel): + additionalModelRequestFields: dict[str, str] | None = None + + +@dataclass(slots=True) +class ConverseRequestCapture: + """The Converse bodies the proxy actually sent upstream, as seen by a live + edge sitting between the proxy and Bedrock.""" + + _bodies: list[ConverseRequestBody] = field(default_factory=list) + _lock: threading.Lock = field(default_factory=threading.Lock) + + def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None: + if body is None or "/converse" not in url: + return + with self._lock: + self._bodies.append(ConverseRequestBody.model_validate_json(body)) + + @property + def bodies(self) -> tuple[ConverseRequestBody, ...]: + with self._lock: + return tuple(self._bodies) + WEATHER_TOOL = ResponsesFunctionTool( name="get_weather", @@ -295,6 +329,55 @@ class TestResponses: arguments = WeatherArguments.model_validate(raw_arguments) assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) + def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( + self, endpoints_client: EndpointsClient, resources: ResourceManager, endpoint: str + ) -> None: + capture: Final = ConverseRequestCapture() + edge: Final = start_provider_edge( + LiveEdge(observe_request=capture.observe, sign=bedrock_signer(BEDROCK_EDGE_REGION)), + mounts=MappingProxyType({BEDROCK_EDGE_MOUNT: f"https://bedrock-runtime.{BEDROCK_EDGE_REGION}.amazonaws.com"}), + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + ) + resources.defer(edge.shutdown) + model: Final = f"e2e-responses-{unique_marker()}" + model_id: Final = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + api_base=edge.edge.api_base(BEDROCK_EDGE_MOUNT), + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name=BEDROCK_EDGE_REGION, + allowed_openai_params=["safety_identifier"], + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key: Final = resources.key() + safety_identifier: Final = f"end-user-{unique_marker()}" + + if endpoint == "/v1/responses": + responses_result: Final = endpoints_client.responses( + key, model, "reply with one word", safety_identifier=safety_identifier + ) + require_successful_call(responses_result) + else: + unwrap( + endpoints_client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="reply with one word")], + safety_identifier=safety_identifier, + ), + ) + ) + + assert [body.additionalModelRequestFields for body in capture.bodies] == [ + {"safety_identifier": safety_identifier} + ], f"{endpoint} did not forward safety_identifier to Bedrock Converse: {capture.bodies}" + @pytest.mark.skip(reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") def test_missing_input_returns_error( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 4b202e3c663..47ef672ebec 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -298,6 +298,7 @@ class ChatBody(BaseModel): max_completion_tokens: int | None = None temperature: float | None = None user: str | None = None + safety_identifier: str | None = None metadata: ChatMetadata | None = None reasoning_effort: str | None = None thinking: ThinkingParam | None = None @@ -976,6 +977,7 @@ class LiteLLMParamsBody(BaseModel): api_base: str | None = None api_version: str | None = None realtime_protocol: str | None = None + allowed_openai_params: list[str] | None = None aws_access_key_id: str | None = None aws_secret_access_key: str | None = None aws_region_name: str | None = None diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 136b00208f7..fc10dde2a77 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -99,6 +99,7 @@ from provider_cache import ( SIGNATURE_HEADERS, CacheEdge, MountPolicy, + RequestSigner, is_bedrock, scoped_edge_base, split_test_segment, @@ -539,6 +540,7 @@ class ReplayEdge: @dataclass(frozen=True, slots=True) class LiveEdge: observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None + sign: RequestSigner | None = None type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge @@ -788,14 +790,16 @@ def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None, observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None, + sign: RequestSigner | None = None, ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } if observe_request is not None: observe_request(url, forwarded, body) + outbound: Final = forwarded if sign is None else sign(method, url, forwarded, body) head: Final = ( - forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) + forward_stream(method, url, headers=outbound, body=body, timeout=timeout) if cache is None else cache.forward(mount, method, url, forwarded, body, timeout, test_key=test_key) ) match head: @@ -871,10 +875,10 @@ def handle_edge_request( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, mount, test_key, ) - case LiveEdge(observe_request=observe_request): + case LiveEdge(observe_request=observe_request, sign=sign): return _handle_live( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, - observe_request=observe_request, + observe_request=observe_request, sign=sign, ) case RecordEdge(): return _handle_record( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 6077f281a81..7b9de4644b4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1248,6 +1248,16 @@ class TestFunctionCallTransformation: assert "tool_choice" not in result assert "tools" not in result + def test_safety_identifier_forwarded_to_chat_completion_request(self) -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="bedrock/global.openai.gpt-5.6-luna", + input="hi", + responses_api_request={"safety_identifier": "user-7f3a"}, + custom_llm_provider="bedrock", + ) + + assert result["safety_identifier"] == "user-7f3a" + def test_parallel_tool_calls_dropped_when_no_chat_tools_remain(self) -> None: transform: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request codex_tool_search: Final = { From ba93c7402a029ac7efb41148e31b3a6b1d186f8b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 04:53:10 +0000 Subject: [PATCH 122/127] test(e2e): tolerate provider retries in safety_identifier capture assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/llm_translation/test_responses_e2e.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 083bc66b3f2..a5dc48a015f 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -374,9 +374,11 @@ class TestResponses: ) ) - assert [body.additionalModelRequestFields for body in capture.bodies] == [ - {"safety_identifier": safety_identifier} - ], f"{endpoint} did not forward safety_identifier to Bedrock Converse: {capture.bodies}" + forwarded: Final = tuple(body.additionalModelRequestFields for body in capture.bodies) + assert forwarded, f"{endpoint} produced no Bedrock Converse request" + assert forwarded == ({"safety_identifier": safety_identifier},) * len(forwarded), ( + f"{endpoint} did not forward safety_identifier to Bedrock Converse on every attempt: {capture.bodies}" + ) @pytest.mark.skip(reason="stage red: product gap, /v1/responses 500s (aresponses TypeError) on missing input instead of 400") @pytest.mark.covers("llm.responses.openai.input_validation.nonstream.works") From 7e0fa40fe3b260efe8447077fa4726e8e0080729 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:17:03 +0000 Subject: [PATCH 123/127] test(e2e): gate the Bedrock edge capture behind a provider_edge_host opt-in The Buildkite ephemeral stack runs the gateway in another pod, so it cannot reach the pytest host's provider edge. The GitHub changed-e2e lane runs gateways on the runner and sets E2E_PROVIDER_EDGE_HOST_REACHABLE Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-e2e-changed.yml | 1 + tests/e2e/AGENTS.md | 2 +- tests/e2e/conftest.py | 7 +++++++ tests/e2e/e2e_config.py | 1 + tests/e2e/llm_translation/test_responses_e2e.py | 1 + tests/e2e/pytest.ini | 1 + 6 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index c9f08deb36e..f659afad7e3 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -175,6 +175,7 @@ jobs: env: TESTS: ${{ needs.detect.outputs.tests }} E2E_FIXTURE_MODE: live + E2E_PROVIDER_EDGE_HOST_REACHABLE: '1' run: | umask 077 read -r -a test_files <<< "${TESTS}" diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 78ba2ec4ed1..c25e958242f 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -122,7 +122,7 @@ E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days. CI records and replays this lane on a schedule in `.github/workflows/e2e_record_replay.yml`, publishing the bundle as a private `e2e-fixtures-bundle` artifact instead of committing it, selecting the tests with the `@pytest.mark.replayable` marker, and proving the bogus-credentials replay hermetic by counting provider egress with `.github/scripts/e2e_egress_sentinel.py` -Current limits: Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode +Current limits: Bedrock cannot be mounted in record or replay (SigV4 signs the Host header, so a rewritten api_base fails signature verification); a test that needs to observe the Converse body registers its own `LiveEdge` with `provider_edge_bedrock.bedrock_signer` re-signing the forwarded request, and carries the `provider_edge_host` opt-in marker because the gateway must reach the pytest host, which the Buildkite ephemeral stack cannot (the GitHub changed-e2e lane, whose gateways run on the runner, sets `E2E_PROVIDER_EDGE_HOST_REACHABLE`). Deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode ## Typing diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 52a634693c5..8776d00d502 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -31,6 +31,7 @@ from e2e_config import ( MANAGED_FILES_OPT_IN_ENV, MCP_OAUTH_LIVE_OPT_IN_ENV, PROMPT_CACHING_OPT_IN_ENV, + PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, REDIS_CHAOS_OPT_IN_ENV, WEEKLY_ANOMALY_OPT_IN_ENV, @@ -59,6 +60,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, + "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, } ) @@ -143,6 +145,11 @@ def pytest_configure(config: pytest.Config) -> None: "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " "E2E_MCP_OAUTH_LIVE is set", ) + config.addinivalue_line( + "markers", + "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " + "gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index a79c158f9c4..14de4619664 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -146,6 +146,7 @@ PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" +PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index a5dc48a015f..1c2336bd166 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -329,6 +329,7 @@ class TestResponses: arguments = WeatherArguments.model_validate(raw_arguments) assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + @pytest.mark.provider_edge_host @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( self, endpoints_client: EndpointsClient, resources: ResourceManager, endpoint: str diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 97acb9ec52b..c6acd449884 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,3 +13,4 @@ markers = cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set + provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set From 0317903a440117c0ffd4faba3639049147a2645e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:16:02 +0000 Subject: [PATCH 124/127] ci(e2e-changed): surface failed test ids from the pytest log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-e2e-changed.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index f659afad7e3..6da16a33ea3 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -176,6 +176,7 @@ jobs: TESTS: ${{ needs.detect.outputs.tests }} E2E_FIXTURE_MODE: live E2E_PROVIDER_EDGE_HOST_REACHABLE: '1' + COLUMNS: '400' run: | umask 077 read -r -a test_files <<< "${TESTS}" @@ -190,6 +191,7 @@ jobs: uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" verified=$? set -e + grep -E '^(FAILED|ERROR) ' "${log}" || true grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1 echo "::endgroup::" if [ "${status}" = "5" ]; then From b38504b5a6f646c7d4eaedb1847b32b4b18647b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:37:28 +0000 Subject: [PATCH 125/127] test(e2e): assert the forwarded Converse body without requiring the model to accept safety_identifier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../e2e/llm_translation/test_responses_e2e.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 1c2336bd166..525231de917 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -19,7 +19,6 @@ from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, un from e2e_http import ( assert_client_error, require_successful_call, - unwrap, ) from endpoints_client import ( EndpointsClient, @@ -359,20 +358,15 @@ class TestResponses: safety_identifier: Final = f"end-user-{unique_marker()}" if endpoint == "/v1/responses": - responses_result: Final = endpoints_client.responses( - key, model, "reply with one word", safety_identifier=safety_identifier - ) - require_successful_call(responses_result) + endpoints_client.responses(key, model, "reply with one word", safety_identifier=safety_identifier) else: - unwrap( - endpoints_client.proxy.chat( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content="reply with one word")], - safety_identifier=safety_identifier, - ), - ) + endpoints_client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="reply with one word")], + safety_identifier=safety_identifier, + ), ) forwarded: Final = tuple(body.additionalModelRequestFields for body in capture.bodies) From 8795be0a65754221bee0b318ad355843d5e3ca34 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 10:51:46 +0000 Subject: [PATCH 126/127] refactor(types): replace Any with proven types in 32 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_logging.py | 8 +++++--- litellm/a2a_protocol/streaming_iterator.py | 12 ++++++------ litellm/integrations/focus/focus_logger.py | 4 +++- .../generic_prompt_manager.py | 2 +- .../mavvrik_focus/mavvrik_focus_logger.py | 17 ++++++++++++++--- litellm/integrations/otel/presets/agentops.py | 7 +++++-- litellm/integrations/otel/runtime.py | 13 ++++++++----- .../bounded_prometheus_series_tracker.py | 16 +++++++++++----- .../vector_store_pre_call_hook.py | 2 +- litellm/integrations/weights_biases.py | 9 +++++---- .../usage_object_transformation.py | 6 +++--- litellm/litellm_core_utils/safe_json_dumps.py | 2 +- litellm/llms/azure/fine_tuning/handler.py | 10 +++++----- .../llms/codestral/completion/transformation.py | 2 +- litellm/llms/databricks/common_utils.py | 4 ++-- litellm/llms/gemini/realtime/transformation.py | 4 ++-- .../llms/jina_ai/embedding/transformation.py | 2 +- .../llms/openrouter/embedding/transformation.py | 4 +++- litellm/llms/reducto/common.py | 4 +++- .../embedding/transformation.py | 3 ++- .../vertex_ai/agent_engine/transformation.py | 2 +- litellm/proxy/a2a/discovery.py | 7 ++++--- .../proxy/agent_endpoints/databricks_oauth.py | 11 ++++++----- litellm/proxy/client/cli/main.py | 13 +++++++++++-- .../generic_guardrail_api/__init__.py | 2 +- .../guardrails/guardrail_hooks/onyx/onyx.py | 4 ++-- litellm/proxy/hooks/responses_id_security.py | 6 +++--- .../callback_logs_endpoints.py | 5 +++-- .../management_v1/spend_logs.py | 8 ++++---- .../object_permission_utils.py | 4 ++-- .../gemini_passthrough_logging_handler.py | 2 +- .../proxy/public_endpoints/public_endpoints.py | 10 +++++++--- 32 files changed, 127 insertions(+), 78 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 5ba0c080364..644a79d8cbd 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -5,6 +5,7 @@ import logging import os import re import sys +from collections.abc import Sequence from datetime import datetime from logging import Formatter from typing import Any, Final, TextIO @@ -186,7 +187,8 @@ class SecretRedactionFilter(logging.Filter): record.stack_info = _redact_string(record.stack_info) # rebind-ok: a Filter scrubs records in place # Redact extra fields passed via logger.debug("msg", extra={...}) - for key, value in list(record.__dict__.items()): + record_items: Final[Sequence[tuple[str, object]]] = list(record.__dict__.items()) + for key, value in record_items: if key in _STANDARD_RECORD_ATTRS: continue if isinstance(value, str): @@ -507,7 +509,7 @@ handler.addFilter(_secret_filter) handler.addFilter(_correlation_filter) -def _try_parse_json_message(message: str) -> dict[str, Any] | None: +def _try_parse_json_message(message: str) -> dict[str, object] | None: """ Try to parse a log message as JSON. Returns parsed dict if valid, else None. Handles messages that are entirely valid JSON (e.g. json.dumps output). @@ -585,7 +587,7 @@ class JsonFormatter(Formatter): def format(self, record): message_str: Final = record.getMessage() - json_record: Final[dict[str, Any]] = { + json_record: Final[dict[str, object]] = { "message": message_str, "level": record.levelname, "timestamp": self.formatTime(record), diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index d936caeb75e..8232d7cf2d8 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -5,7 +5,7 @@ A2A Streaming Iterator with token tracking and logging support. import asyncio from collections.abc import AsyncIterator from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_logger @@ -15,7 +15,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj if TYPE_CHECKING: - from a2a.types import SendStreamingMessageRequest, SendStreamingMessageResponse + from a2a.compat.v0_3.types import SendStreamingMessageRequest, SendStreamingMessageResponse class A2AStreamingIterator: @@ -39,9 +39,9 @@ class A2AStreamingIterator: self.start_time = datetime.now() # Collect chunks for token counting - self.chunks: list[Any] = [] + self.chunks: list[SendStreamingMessageResponse] = [] self.collected_text_parts: list[str] = [] - self.final_chunk: Any | None = None + self.final_chunk: SendStreamingMessageResponse | None = None def __aiter__(self): return self @@ -69,7 +69,7 @@ class A2AStreamingIterator: await self._handle_stream_complete() raise - def _collect_text_from_chunk(self, chunk: Any) -> None: + def _collect_text_from_chunk(self, chunk: "SendStreamingMessageResponse") -> None: """Extract text from a streaming chunk and add to collected parts.""" try: chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} @@ -79,7 +79,7 @@ class A2AStreamingIterator: except Exception: verbose_logger.debug("Failed to extract text from A2A streaming chunk") - def _is_completed_chunk(self, chunk: Any) -> bool: + def _is_completed_chunk(self, chunk: "SendStreamingMessageResponse") -> bool: """Check if chunk indicates stream completion.""" try: chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index c9b47835948..dce51b8190b 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -15,6 +15,8 @@ from .destinations import FocusTimeWindow if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from .export_engine import FocusExportEngine else: AsyncIOScheduler = Any @@ -111,7 +113,7 @@ class FocusLogger(CustomLogger): """Entry point for scheduler jobs to run export cycle with locking.""" from litellm.proxy.proxy_server import proxy_logging_obj - pod_lock_manager = None + pod_lock_manager: PodLockManager | None = None if proxy_logging_obj is not None: writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 77d315d0cee..de01b2bb02c 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -58,7 +58,7 @@ class GenericPromptManager(CustomPromptManagement): api_key: str | None = None, timeout: int = 30, prompt_id: str | None = None, - additional_provider_specific_query_params: dict[str, Any] | None = None, + additional_provider_specific_query_params: Mapping[str, object] | None = None, **kwargs, ): """ diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 3c189b4d53e..7e3c4cc3ce8 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -21,7 +21,7 @@ from __future__ import annotations import os from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import litellm from litellm._logging import verbose_proxy_logger @@ -35,6 +35,17 @@ else: AsyncIOScheduler = Any +class _PodLockManager(Protocol): + """The subset of PodLockManager this logger drives to serialize the export across pods.""" + + @property + def redis_cache(self) -> object: ... + + async def acquire_lock(self, cronjob_id: str) -> bool | None: ... + + async def release_lock(self, cronjob_id: str) -> None: ... + + def _parse_metrics_marker( marker: object | None, ) -> datetime | None: @@ -226,9 +237,9 @@ class MavvrikFocusLogger(FocusLogger): """Scheduler entry point — uses Mavvrik-specific pod-lock key.""" from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415 - pod_lock_manager = None + pod_lock_manager: _PodLockManager | None = None if proxy_logging_obj is not None: - writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) + writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: pod_lock_manager = getattr(writer, "pod_lock_manager", None) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 965213f2ee4..58123656caa 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,9 +9,12 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ +from collections.abc import Sequence from typing import Any, Final import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -71,7 +74,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> Any: +def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -106,7 +109,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Any) -> Any: + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: self._ensure_authenticated() return super().export(spans) diff --git a/litellm/integrations/otel/runtime.py b/litellm/integrations/otel/runtime.py index c6eaecd108b..13903597e1a 100644 --- a/litellm/integrations/otel/runtime.py +++ b/litellm/integrations/otel/runtime.py @@ -8,13 +8,16 @@ identity unconditionally. """ from collections.abc import Callable, Iterator -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager from functools import cache -from typing import Any, Final +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from opentelemetry.trace import Span @cache -def _otel_runtime() -> "tuple[Callable[[str], Any], Callable[..., None]] | None": +def _otel_runtime() -> "tuple[Callable[[str], AbstractContextManager[Span | None]], Callable[..., None]] | None": """Resolve the SDK-backed hooks once and cache the outcome, absence included. CPython never caches a failed import, so without this memoization every call @@ -29,7 +32,7 @@ def _otel_runtime() -> "tuple[Callable[[str], Any], Callable[..., None]] | None" @contextmanager -def phase_span(name: str) -> "Iterator[Any]": +def phase_span(name: str) -> "Iterator[Span | None]": """Run a request phase inside a live active span so its DB/service calls nest. Yields ``None`` (a plain no-op) when the OTel SDK is unavailable or V2 is not @@ -43,7 +46,7 @@ def phase_span(name: str) -> "Iterator[Any]": yield span -def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None: +def seed_request_identity(user_api_key_dict: object, model: object = None) -> None: """Seed request-identity Baggage at the auth boundary (no-op without V2).""" runtime: Final = _otel_runtime() if runtime is None: diff --git a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py index c1ccf09d5d6..ba7d54fafea 100644 --- a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py +++ b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py @@ -3,7 +3,13 @@ from __future__ import annotations import time from collections import OrderedDict from threading import RLock -from typing import Any, Final +from typing import Final, Protocol + + +class _RemovableMetric(Protocol): + """The one prometheus-client metric method this tracker calls.""" + + def remove(self, *labelvalues: object) -> None: ... class BoundedPrometheusSeriesTracker: @@ -21,7 +27,7 @@ class BoundedPrometheusSeriesTracker: def track_series( self, - metric: Any, + metric: _RemovableMetric, metric_name: str, label_values: tuple[str | None, ...], max_series: int | None, @@ -60,7 +66,7 @@ class BoundedPrometheusSeriesTracker: break del series[tracked_label_values] - def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool: + def remove_series(self, metric: _RemovableMetric, label_values: tuple[str | None, ...]) -> bool: """Drop one child series, True when it is gone (removed or never existed).""" return self._remove_metric_child(metric, label_values) @@ -82,7 +88,7 @@ class BoundedPrometheusSeriesTracker: def _remove_metric_series( self, - metric: Any, + metric: _RemovableMetric, series: OrderedDict[tuple[str | None, ...], float], label_values: tuple[str | None, ...], ) -> None: @@ -90,7 +96,7 @@ class BoundedPrometheusSeriesTracker: series.pop(label_values, None) @staticmethod - def _remove_metric_child(metric: Any, label_values: tuple[str | None, ...]) -> bool: + def _remove_metric_child(metric: _RemovableMetric, label_values: tuple[str | None, ...]) -> bool: """ Remove the Prometheus child for ``label_values`` and report whether the tracker should commit the matching state change. diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index b2243060c6c..74fb8a8d6a3 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -406,7 +406,7 @@ class VectorStorePreCallHook(CustomLogger): request_data: dict, response_chunk: Any, call_type: CallTypes | None, - ) -> Any | None: + ) -> object | None: """ Add search results to the final streaming chunk. diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index d1a8ec098cf..9bc070a1f9a 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -4,6 +4,7 @@ imported_openAIResponse = True try: import io import logging + from collections.abc import Mapping from typing import Any, Literal, Protocol, TypeVar from wandb.sdk.data_types import trace_tree @@ -43,7 +44,7 @@ try: @staticmethod def results_to_trace_tree( - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, results: list[trace_tree.Result], time_elapsed: float, @@ -73,7 +74,7 @@ try: def _resolve_edit( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: @@ -91,7 +92,7 @@ try: def _resolve_completion( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: @@ -134,7 +135,7 @@ try: def _request_response_result_to_trace( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, request_str: str, choices: list[str], diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index f11f6d46fb2..a02c40b7611 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -50,7 +50,7 @@ _INTERACTIONS_MODALITY_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) -def _modality_field(entry: Mapping[str, Any]) -> str | None: +def _modality_field(entry: Mapping[str, object]) -> str | None: return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) @@ -58,7 +58,7 @@ def _token_count(value: object) -> int: return value if isinstance(value, int) else 0 -def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, int]: +def _modality_token_sums(entries: Sequence[Mapping[str, object]]) -> Mapping[str, int]: fields: Final = frozenset(field for entry in entries if (field := _modality_field(entry)) is not None) return MappingProxyType( { @@ -68,7 +68,7 @@ def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, i ) -def _google_search_query_count(usage_object: Mapping[str, Any]) -> int: +def _google_search_query_count(usage_object: Mapping[str, object]) -> int: entries: Final = usage_object.get("grounding_tool_count") if not isinstance(entries, Sequence): return 0 diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 4f9ac82d57d..63242a580e7 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -85,7 +85,7 @@ def safe_json_structure( def safe_dumps( - data: Any, + data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, ) -> str: diff --git a/litellm/llms/azure/fine_tuning/handler.py b/litellm/llms/azure/fine_tuning/handler.py index ac1e430e063..36c4fae04c7 100644 --- a/litellm/llms/azure/fine_tuning/handler.py +++ b/litellm/llms/azure/fine_tuning/handler.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine -from typing import Any, Final, cast +from typing import Final, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI @@ -19,7 +19,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): """ @staticmethod - def _ensure_training_type(create_fine_tuning_job_data: dict[str, Any]) -> None: + def _ensure_training_type(create_fine_tuning_job_data: dict[str, object]) -> None: """ Azure requires trainingType in extra_body. Default to 1 (supervised) if omitted. """ @@ -66,7 +66,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: self._ensure_training_type(create_fine_tuning_job_data) openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( @@ -109,7 +109,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -149,7 +149,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index e3c3fd1231c..baa134bb398 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -29,7 +29,7 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): random_seed: int | None = None, stop: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[dict[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index a4ec2c5378b..f2f9df422e8 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -12,7 +12,7 @@ Authentication priority: import os import re -from typing import Any, Final, Literal +from typing import Final, Literal from urllib.parse import urlsplit, urlunsplit from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -48,7 +48,7 @@ class DatabricksBase: ] @classmethod - def redact_sensitive_data(cls, data: Any) -> Any: + def redact_sensitive_data(cls, data: object) -> object: """ Redact sensitive information (tokens, secrets) from data before logging. diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 79985569c5f..e64cbf88d95 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -453,7 +453,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return normalized @staticmethod - def _finalize_gemini_live_setup(model: str, setup: dict[str, Any]) -> dict[str, Any]: + def _finalize_gemini_live_setup(model: str, setup: dict[str, object]) -> dict[str, object]: generation_config: Final = setup.get("generationConfig") if isinstance(generation_config, dict): modalities: Final = generation_config.get("responseModalities") @@ -1172,7 +1172,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): def map_openai_event( self, key: str, - value: Any, + value: object, current_delta_type: ALL_DELTA_TYPES | None, ) -> OpenAIRealtimeEventTypes | ResponsesAPIStreamEvents: if isinstance(value, dict): diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 26f512979e5..260d9e6e494 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -31,7 +31,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): def __init__( self, ) -> None: - locals_: Final = locals().copy() + locals_: Final[dict[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 29d0c8c1c56..14b0e462ea7 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -170,7 +170,9 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> OpenRouterException: """ Get the error class for OpenRouter errors. """ diff --git a/litellm/llms/reducto/common.py b/litellm/llms/reducto/common.py index 9b9efd24b72..b194590fdb9 100644 --- a/litellm/llms/reducto/common.py +++ b/litellm/llms/reducto/common.py @@ -3,6 +3,8 @@ import binascii from collections import defaultdict from typing import TYPE_CHECKING, Any, Final, NoReturn +import httpx + from litellm.constants import request_timeout REDUCTO_API_BASE: Final = "https://platform.reducto.ai" @@ -62,7 +64,7 @@ def extract_file_id_or_bytes( return None, raw_bytes, mime -def _extract_file_id_from_upload_response(response: Any) -> str: +def _extract_file_id_from_upload_response(response: httpx.Response) -> str: try: payload: Final = response.json() except ValueError as exc: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index 3f228c0881d..fc9c6bcc19f 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllEmbeddingInputValues @@ -160,7 +161,7 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: + def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: """ Get the error class for Vercel AI Gateway errors. """ diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index e430d9e2280..c5ca9f38144 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -205,7 +205,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): session_id: Final = self._get_session_id(optional_params) # Build the input - input_data: Final[dict[str, Any]] = { + input_data: Final[dict[str, str]] = { "message": prompt, "user_id": user_id, } diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index e08c938f195..ff58c9c85ec 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -14,6 +14,7 @@ fetcher dispatches by ``discovery_mode``: pure-A2A fallback strategy returns 404 for these deployments. """ +from collections.abc import Mapping from enum import Enum from typing import Any, Final from urllib.parse import urlencode @@ -55,7 +56,7 @@ def _normalize_base_url(base_url: str) -> str: def _build_langgraph_platform_paths( - params: dict[str, Any] | None, + params: Mapping[str, object] | None, ) -> tuple[str, ...]: """Build the paths to try for LangGraph Platform discovery. @@ -71,7 +72,7 @@ def _build_langgraph_platform_paths( return tuple(f"{path}?{query}" for path in AGENT_CARD_WELL_KNOWN_PATHS) -def _paths_for_mode(mode: DiscoveryMode, params: dict[str, Any] | None) -> tuple[str, ...]: +def _paths_for_mode(mode: DiscoveryMode, params: Mapping[str, object] | None) -> tuple[str, ...]: if mode == DiscoveryMode.WELL_KNOWN_FALLBACK: return AGENT_CARD_WELL_KNOWN_PATHS if mode == DiscoveryMode.LANGGRAPH_PLATFORM: @@ -83,7 +84,7 @@ async def fetch_well_known_card( base_url: str, *, discovery_mode: DiscoveryMode = DiscoveryMode.WELL_KNOWN_FALLBACK, - params: dict[str, Any] | None = None, + params: Mapping[str, object] | None = None, timeout: float = DEFAULT_DISCOVERY_TIMEOUT_SECONDS, headers: dict[str, str] | None = None, ) -> dict[str, Any]: diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 4c3b1bc084d..38a76ea6890 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -25,8 +25,9 @@ Config example:: import asyncio import base64 import hashlib +from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Final +from typing import Final import httpx @@ -43,7 +44,7 @@ _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 _DEFAULT_TTL_SECONDS: Final = 3600 -def _resolve_secret(value: Any) -> str | None: +def _resolve_secret(value: object) -> str | None: """Resolve a config value, expanding ``os.environ/`` references.""" if not isinstance(value, str): return None @@ -75,7 +76,7 @@ class DatabricksAppOAuthConfig: def parse_databricks_oauth_config( - litellm_params: dict[str, Any] | None, + litellm_params: Mapping[str, object] | None, ) -> DatabricksAppOAuthConfig | None: """Build a Databricks App OAuth config from an agent's ``litellm_params``. @@ -191,7 +192,7 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): except httpx.HTTPError as exc: raise ValueError(f"Databricks App OAuth token request failed: {exc}") from exc - body: Final = response.json() + body: Final[object] = response.json() if not isinstance(body, dict): raise ValueError( f"Databricks App OAuth token response returned non-object JSON (got {type(body).__name__})" @@ -215,7 +216,7 @@ databricks_app_oauth_token_cache: Final = DatabricksAppOAuthTokenCache() async def resolve_databricks_app_auth_header( - litellm_params: dict[str, Any] | None, + litellm_params: Mapping[str, object] | None, ) -> dict[str, str] | None: """Return ``{"Authorization": "Bearer "}`` for a Databricks App agent. diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 63e38c93221..6d63acc7479 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,7 +9,15 @@ from litellm._version import version as litellm_version from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands -from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, login, logout, whoami +from .commands.auth import ( + CliContextObj, + auth_group, + context_secret_vault, + get_stored_api_key, + login, + logout, + whoami, +) from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names @@ -126,7 +134,8 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s @click.pass_context def version(ctx: click.Context): """Show the LiteLLM Proxy CLI and server version.""" - print_version(ctx.obj.get("base_url"), ctx.obj.get("api_key")) + ctx_obj: Final[CliContextObj] = ctx.obj + print_version(ctx_obj.get("base_url"), ctx_obj.get("api_key")) # Add authentication commands as top-level commands diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index d1576b68813..e3511d46544 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def _get_config_value(litellm_params: Any, optional_params: Any, attribute_name: str) -> Any | None: +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> Any | None: if optional_params is not None: value: Final = ( optional_params.get(attribute_name) diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 7529c4ce3f3..1a6feb47215 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os import uuid -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException @@ -63,7 +63,7 @@ class OnyxGuardrail(CustomGuardrail): async def _validate_with_guard_server( self, - payload: Any, + payload: object, input_type: Literal["request", "response"], conversation_id: str, ) -> dict: diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index d9050489095..bdf7e2ab53d 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -40,7 +40,7 @@ _UNMANAGED_RESPONSE_ID_DETAIL: Final = ( _PROXY_ADMIN_ROLES: Final = frozenset({LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value}) -def _proxy_general_settings() -> Mapping[str, Any]: +def _proxy_general_settings() -> Mapping[str, object]: from litellm.proxy.proxy_server import general_settings return general_settings @@ -107,7 +107,7 @@ def _is_responses_api_create_route(request_route: str | None) -> bool: class ResponsesIDSecurity(CustomLogger): def __init__( self, - general_settings_reader: Callable[[], Mapping[str, Any]] = _proxy_general_settings, + general_settings_reader: Callable[[], Mapping[str, object]] = _proxy_general_settings, signing_key_reader: Callable[[], str | None] = _proxy_signing_key, ) -> None: self._general_settings_reader: Final = general_settings_reader @@ -307,7 +307,7 @@ class ResponsesIDSecurity(CustomLogger): data: dict, user_api_key_dict: "UserAPIKeyAuth", response: LLMResponseTypes, - ) -> Any: + ) -> LLMResponseTypes: """ Queue response IDs for batch processing instead of writing directly to DB. diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index cecadc03d71..4a1079871b0 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -15,6 +15,7 @@ self-describing `StandardLoggingPayload`, so completions/responses can use it to """ import uuid +from collections.abc import Mapping from datetime import datetime, timezone from typing import Any, Final @@ -48,7 +49,7 @@ class CallbackLogsReplayer: """ @staticmethod - def _epoch_to_datetime(value: Any) -> datetime: + def _epoch_to_datetime(value: object) -> datetime: """`StandardLoggingPayload` stores startTime/endTime as float epoch seconds.""" if isinstance(value, (int, float)): return datetime.fromtimestamp(float(value), tz=timezone.utc) @@ -114,7 +115,7 @@ class CallbackLogsReplayer: return logging_obj @staticmethod - def _response_obj_from_payload(payload: dict[str, Any]) -> dict[str, Any]: + def _response_obj_from_payload(payload: Mapping[str, object]) -> dict[str, object]: """Minimal response object so usage-derived spend-log fields resolve.""" return { "id": payload.get("id"), diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index f6907a7f87a..1cbc454ca5e 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,7 +1,7 @@ """`/management/v1/spend_logs` facets.""" from datetime import datetime, timezone -from typing import Annotated, Any, Final, Literal +from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, Query, Request @@ -39,7 +39,7 @@ async def _spend_log_scope_clause( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, next_param_index: int, -) -> tuple[str | None, tuple[Any, ...]]: +) -> tuple[str | None, tuple[str | list[str], ...]]: """SQL predicate restricting the facet to spend logs this caller may read. Returns ``(None, ())`` for a proxy admin. Mirrors the scoping ``/spend/logs/ui`` @@ -101,8 +101,8 @@ async def _list_spend_log_facet( ) column_sql: Final = "end_user" if column == "end_user" else '"user"' - window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) - search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () + window_params: Final[tuple[datetime, datetime]] = (_as_utc(start_time), _as_utc(end_time)) + search_params: Final[tuple[str, ...]] = (f"%{escape_like(q)}%",) if q else () search_clause: Final = (f"{column_sql} ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () scope_clause, scope_params = await _spend_log_scope_clause( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index daab38d3662..437e6763502 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -230,7 +230,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index a95ee87fd31..d97ddb9a909 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -147,7 +147,7 @@ class GeminiPassthroughLoggingHandler: - Creates standard logging object - Logs in litellm callbacks """ - kwargs: dict[str, Any] = {} + kwargs: dict[str, object] = {} model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) complete_streaming_response: Final = GeminiPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=all_chunks, diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index e395f56194f..26a5c44fce1 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -199,9 +199,13 @@ def _build_endpoints(raw: _ProvidersFile) -> list[_EndpointEntry]: return result +_PROVIDERS_FILE_ADAPTER: Final = TypeAdapter(_ProvidersFile) +_PROVIDER_CREATE_FIELDS_ADAPTER: Final = TypeAdapter(list[ProviderCreateInfo]) + + def _load_endpoints() -> list[_EndpointEntry]: - raw: Final[_ProvidersFile] = json.loads( - files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8") + raw: Final = _PROVIDERS_FILE_ADAPTER.validate_python( + json.loads(files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8")) ) return _build_endpoints(raw) @@ -398,7 +402,7 @@ async def get_provider_fields() -> list[ProviderCreateInfo]: ) with open(provider_create_fields_path, "r") as f: - provider_create_fields: Final = json.load(f) + provider_create_fields: Final = _PROVIDER_CREATE_FIELDS_ADAPTER.validate_python(json.load(f)) return provider_create_fields From 85e72a2e2fac84e9be99cbac4ede9d4991b834af Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:00:50 +0000 Subject: [PATCH 127/127] chore(prices): sync OpenRouter prices: 2 models openrouter/deepseek/deepseek-v4-flash: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1976437f900..d3053dee25d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.5526e-07, + "input_cost_per_token": 9.53172e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.91052e-06, + "output_cost_per_token": 1.906344e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9605e-08, + "cache_read_input_token_cost": 7.9431e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 5.852e-08, + "output_cost_per_token": 1.1704e-07, + "cache_read_input_token_cost": 1.1704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1976437f900..d3053dee25d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41537,21 +41537,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.5526e-07, + "input_cost_per_token": 9.53172e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.91052e-06, + "output_cost_per_token": 1.906344e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.9605e-08, + "cache_read_input_token_cost": 7.9431e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -67259,9 +67259,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 5.852e-08, + "output_cost_per_token": 1.1704e-07, + "cache_read_input_token_cost": 1.1704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000,