From 02bb3f9310e6633caefd8f739d2340ebd0d56cff Mon Sep 17 00:00:00 2001 From: Timik232 <100406268+Timik232@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:56:54 +0300 Subject: [PATCH 01/92] fix(streaming): keep response id stable across streamed chunks Providers that stream via GenericStreamingChunk (e.g. GigaChat) do not propagate an upstream response id, so every chunk of one streamed response got a freshly generated id. Pin CustomStreamWrapper.response_id from the first chunk it creates, mirroring the existing 'created' pinning (#11437). Clients that merge deltas by chunk id (e.g. goose) split one reply into one message per chunk. Fixes #38098 --- .../litellm_core_utils/streaming_handler.py | 2 + .../test_streaming_handler.py | 79 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..fb693943b2b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -816,6 +816,8 @@ class CustomStreamWrapper: model_response: Final = ModelResponseStream(**args) if self.response_id is not None: model_response.id = self.response_id + elif model_response.id: + self.response_id = model_response.id if self.system_fingerprint is not None: model_response.system_fingerprint = self.system_fingerprint diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index b5e33a4e421..a76d427495c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4460,3 +4460,82 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp finally: trace_id_var.set("") session_id_var.set("") + + +class TestStableStreamingResponseId: + """ + All chunks of one streamed response must share the same top-level id + (OpenAI streaming contract). Providers streaming via GenericStreamingChunk + (e.g. GigaChat) do not propagate an upstream response id, so + CustomStreamWrapper must pin the id from the first chunk it creates, + mirroring the existing `created` pinning (issue #11437). + + Clients such as goose merge streamed deltas into one assistant message by + chunk id; per-chunk ids split a single reply into many messages. + """ + + def test_generic_chunks_share_one_id(self): + def _generic_chunks(): + return iter( + [ + { + "text": "Hello", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": " world", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": "", + "tool_use": None, + "is_finished": True, + "finish_reason": "stop", + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + }, + "index": 0, + }, + ] + ) + + wrapper = CustomStreamWrapper( + completion_stream=_generic_chunks(), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + ids = [chunk.id for chunk in wrapper if chunk.id] + assert ids, "no chunks emitted" + assert len(set(ids)) == 1, f"chunk ids differ across one stream: {ids}" + + def test_creator_pins_id_from_first_chunk(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + first = wrapper.model_response_creator() + assert wrapper.response_id == first.id + assert wrapper.model_response_creator().id == first.id + + def test_provider_supplied_id_still_wins(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + wrapper.response_id = "chatcmpl-from-provider" + assert wrapper.model_response_creator().id == "chatcmpl-from-provider" From 4f6fd85ab1e36af0cb3f44b045c3767d0ebe89b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:10:05 -0700 Subject: [PATCH 02/92] feat(proxy): default to the v2 migration resolver, keep v1 as an opt-out The v2 resolver skips the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contend for one database during a rolling deploy. The standalone migration Job already defaulted to v2; this aligns the proxy-server path. v1 stays reachable two ways: --use_legacy_migration_resolver on the CLI, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, where prisma_migration.py calls run_server with a fixed argv and the env var is the only route in. --use_v2_migration_resolver still parses, so existing commands do not die on an unknown option. Because v2 fails fast where v1 retried every failed deploy, a database that is not accepting connections yet, or another instance holding the migration advisory lock, would now kill a boot that used to ride it out. Those two failures are retried, with Prisma's stderr logged each round, and still raise once the attempts are spent. Moves the resolver tests from litellm-proxy-extras/tests, which no CI job runs, into tests/litellm-proxy-extras, and repoints the dedicated Postgres CircleCI job at the legacy path so v1 keeps real-DB and proxy-boot coverage. --- .circleci/config.yml | 15 +- CLAUDE.md | 2 +- .../litellm_proxy_extras/utils.py | 81 ++++-- litellm-proxy-extras/tests/__init__.py | 0 litellm/proxy/proxy_cli.py | 22 +- .../test_setup_database_fail_fast.py | 238 +++++++++++++++++- .../test_basic_python_version.py | 14 +- tests/test_litellm/proxy/test_proxy_cli.py | 83 +++++- 8 files changed, 409 insertions(+), 46 deletions(-) delete mode 100644 litellm-proxy-extras/tests/__init__.py rename {litellm-proxy-extras/tests => tests/litellm-proxy-extras}/test_setup_database_fail_fast.py (50%) diff --git a/.circleci/config.yml b/.circleci/config.yml index 4615a6a5a7e..0fbbd5ec7f7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1483,7 +1483,7 @@ jobs: - run: name: Run tests command: | - uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver" + uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver" installing_litellm_on_python_3_13: docker: @@ -1507,9 +1507,9 @@ jobs: - run: name: Run tests command: | - uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver" + uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver" - installing_litellm_on_python_v2_migration_resolver: + installing_litellm_on_python_legacy_migration_resolver: docker: - *python312_image - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 @@ -1536,10 +1536,10 @@ jobs: url: tcp://localhost:5432 timeout: "60" - run: - name: Run v2 migration resolver proxy smoke test + name: Run legacy migration resolver proxy smoke test command: | uv run --no-sync python -m pytest -vv \ - tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver + tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver helm_chart_testing: machine: @@ -2918,7 +2918,8 @@ jobs: command: | if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \ (grep -q "Database setup failed after multiple retries" docker_output.log || \ - grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then + grep -q "ERROR: Application startup failed. Exiting." docker_output.log || \ + grep -q "Database migration cannot proceed" docker_output.log); then echo "Expected error found. Test passed." else echo "Expected error not found. Test failed." @@ -3050,7 +3051,7 @@ workflows: filters: *main_branches - installing_litellm_on_python_3_13: filters: *main_branches - - installing_litellm_on_python_v2_migration_resolver: + - installing_litellm_on_python_legacy_migration_resolver: filters: *main_branches - helm_chart_testing: requires: diff --git a/CLAUDE.md b/CLAUDE.md index 930825aeb89..819f65059b7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -35,7 +35,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b27221c9beb..22da9b834ab 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -7,7 +7,8 @@ import subprocess import tempfile import time from pathlib import Path -from typing import Optional +from types import MappingProxyType +from typing import Final, Optional from litellm_proxy_extras._logging import logger from litellm_proxy_extras.replica_identity import ( @@ -50,6 +51,35 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile( re.IGNORECASE, ) +_MIGRATE_DEPLOY_ATTEMPTS: Final = 4 + +_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( + { + "deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)", + "P1001": "an unreachable database server", + "P1002": "a database server that timed out", + } +) + + +def _transient_deploy_failure(stderr: str) -> str | None: + """Describe why a failed `prisma migrate deploy` is worth retrying, or None. + + These are environment failures, not migration failures: the database is not + up yet, or another instance holds the migration lock. v1 retried every + failed deploy and absorbed them; failing fast on them instead would turn a + database that is ten seconds late into a dead proxy. + """ + return next( + ( + reason + for marker, reason in _TRANSIENT_DEPLOY_FAILURES.items() + if marker in stderr + ), + None, + ) + + PARTITIONED_SPEND_LOGS_PUSH_ERROR = ( "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), " "so its primary key must include the partition key (\"startTime\"). `prisma db push` " @@ -648,7 +678,7 @@ class ProxyExtrasDBManager: @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ - v2 migration resolver (opt-in via --use_v2_migration_resolver). + v2 migration resolver (what the proxy CLI selects by default). Runs `prisma migrate deploy` and handles standard recovery paths (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does @@ -692,7 +722,7 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - for attempt in range(4): + for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -807,16 +837,36 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" - ) from e + transient = _transient_deploy_failure(stderr) + if transient is None: + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from e + + if attempt == _MIGRATE_DEPLOY_ATTEMPTS - 1: + raise RuntimeError( + f"Database migration failed after " + f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. " + "Check database connectivity and load." + f"\n\nPrisma error:\n{stderr}" + ) from e + + logger.info( + "prisma migrate deploy attempt %s failed on %s, retrying. " + "Prisma error:\n%s", + attempt + 1, + transient, + stderr, + ) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( - "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts or repeated idempotent-recovery " - "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} " + "attempts (retry loop exhausted by timeouts or repeated " + "idempotent-recovery continues). Check database connectivity, " + "load, and _prisma_migrations ledger state." ) finally: os.chdir(original_dir) @@ -864,10 +914,11 @@ class ProxyExtrasDBManager: Args: use_migrate: Whether to use prisma migrate instead of db push - use_v2_resolver: Opt into the v2 migration resolver (safer during + use_v2_resolver: Run the v2 migration resolver (safer during rolling deploys; does not run the diff-and-force recovery - that causes schema thrashing). Defaults to False for - backwards compatibility. + that causes schema thrashing). Defaults to False here so + direct callers keep the old behavior; the proxy CLI passes + True, so the proxy's runtime default is v2. Returns: bool: True if setup was successful, False otherwise @@ -885,7 +936,7 @@ class ProxyExtrasDBManager: @staticmethod def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool: if use_v2_resolver: - logger.info("Using v2 migration resolver (--use_v2_migration_resolver)") + logger.info("Using v2 migration resolver") return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate) schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" diff --git a/litellm-proxy-extras/tests/__init__.py b/litellm-proxy-extras/tests/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8ac63ba25c9..c1d86344a1e 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -913,13 +913,14 @@ class ProxyInitializationHelpers: envvar="ENFORCE_PRISMA_MIGRATION_CHECK", ) @click.option( - "--use_v2_migration_resolver", - is_flag=True, - default=False, + "--use_v2_migration_resolver/--use_legacy_migration_resolver", + default=True, help=( - "Opt into the v2 migration resolver. Avoids the diff-and-force recovery " - "path that can cause schema thrashing during rolling deploys where two " - "LiteLLM versions contend for the same DB. Default is the v1 resolver." + "Which database migration resolver to run at startup. The default v2 " + "resolver avoids the diff-and-force recovery path that can cause schema " + "thrashing during rolling deploys where two LiteLLM versions contend for " + "the same DB. Pass --use_legacy_migration_resolver, or set " + "USE_V2_MIGRATION_RESOLVER=false, to fall back to v1." ), envvar="USE_V2_MIGRATION_RESOLVER", ) @@ -1310,10 +1311,11 @@ def run_server( else: if not use_v2_migration_resolver: print( - "\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. " - "If your deployment has seen schema thrashing during rolling " - "deploys, try --use_v2_migration_resolver (safer: avoids the " - "diff-and-force recovery that caused the thrash).\033[0m" + "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration resolver. " + "The default v2 resolver is safer: it avoids the diff-and-force " + "recovery that caused schema thrashing during rolling deploys. " + "Remove --use_legacy_migration_resolver / " + "USE_V2_MIGRATION_RESOLVER=false to switch back to it.\033[0m" ) try: setup_ok: Final = PrismaManager.setup_database( diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py similarity index 50% rename from litellm-proxy-extras/tests/test_setup_database_fail_fast.py rename to tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 8d66bf872de..f637e139e86 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -1,15 +1,26 @@ -"""Regression tests for ProxyExtrasDBManager v2 migration resolver. +"""Regression tests for ProxyExtrasDBManager's v2 migration resolver. -The v2 resolver is opt-in via `--use_v2_migration_resolver` / the -`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1 -(default) behavior is unchanged from pre-fix. +v2 is what the proxy CLI selects by default; v1 stays reachable via +`--use_legacy_migration_resolver` or `USE_V2_MIGRATION_RESOLVER=false`. At the +library level the resolver is picked with the `use_v2_resolver` kwarg, which +still defaults to False so `migrations/run.py` and any direct caller keep their +own explicit choice. """ +import os import subprocess +import sys from unittest.mock import patch import pytest +sys.path.insert( + 0, + os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") + ), +) + from litellm_proxy_extras.utils import ( ProxyExtrasDBManager, _max_migration_timestamp, @@ -240,3 +251,222 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): 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_STDERR = ( + "Error: ERROR: deadlock detected\n" + "DETAIL: Process 277 waits for ExclusiveLock on advisory lock " + "[17556,0,72707369,1]; blocked by process 278.\n" + "Process 278 waits for ShareLock on virtual transaction 3/1041; " + "blocked by process 277." +) + + +class _DeployApplied: + stdout = "All migrations have been successfully applied." + stderr = "" + returncode = 0 + + +def _deploy_only(deploy_side_effect): + """subprocess.run stand-in that only intercepts `prisma migrate deploy`. + + Everything else the resolver shells out to, the Prisma toolchain check + above all, succeeds untouched, so a mock meant for the deploy call cannot + be silently consumed by an earlier subprocess call. + """ + deploys = {"n": 0} + + def _run(*args, **kwargs): + cmd = args[0] if args else kwargs.get("args", []) + if list(cmd)[-2:] == ["migrate", "deploy"]: + deploys["n"] += 1 + return deploy_side_effect(deploys["n"], cmd) + return _DeployApplied() + + return _run, deploys + + +def _prepare_v2_resolver(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, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda *_a, **_k: None) + + +def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): + """A deadlock on Prisma's migration advisory lock is transient and retried. + + Several proxy replicas booting against one database race `migrate deploy`, + and Postgres aborts one side. v1 retried any failed deploy, so it rode this + out; v2 classifies unrecognised stderr as unrecoverable and raises, which + with v2 as the default would take a replica's whole boot down. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output="" + ) + return _DeployApplied() + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert deploys["n"] == 2, "the deadlocked deploy must be retried, not raised" + + +def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path): + """The deadlock retry stays bounded: a deadlock that never clears still + raises rather than looping forever or reporting a successful migration.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output="" + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 4 + + +@pytest.mark.parametrize( + "stderr", + [ + "Error: P1001: Can't reach database server at `db`:`5432`", + "Error: P1002: The database server was reached but timed out.", + ], +) +def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr): + """A database that is not accepting connections yet is retried, not fatal. + + A proxy and its database starting together race routinely, and v1 rode that + out by retrying every failed deploy. v2 treats unrecognised stderr as + unrecoverable, so without this the default flip would turn a database that + is a few seconds late into a dead proxy instead of a slow boot. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + return _DeployApplied() + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert deploys["n"] == 2, "an unreachable database must be retried, not raised" + + +def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path): + """Retrying connectivity errors must not turn a genuinely unreachable + database into a silent success: after the attempts are spent it still + raises, so the proxy exits instead of serving without its database.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 4 + + +def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog): + """Retrying must not swallow why the database was unreachable. + + Prisma's stderr is captured, so if the retry path neither logs it nor puts + it in the final error, an operator (and CI's bad-DATABASE_URL job, which + greps the boot log for the P1001 line) sees four silent retries and no + cause. + """ + _prepare_v2_resolver(monkeypatch, tmp_path) + stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`" + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + + run, _ = _deploy_only(_side_effect) + with caplog.at_level("INFO", logger="litellm_proxy_extras"): + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError) as exc_info: + ProxyExtrasDBManager.setup_database( + use_migrate=True, use_v2_resolver=True + ) + + assert "P1001" in str(exc_info.value) + assert "P1001" in caplog.text + + +def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path): + """The transient classification must stay narrow: a genuinely broken + migration still fails fast on the first attempt rather than being retried + into the same error four times.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "The `20260101000000_genuinely_broken` migration failed to apply.\n" + 'Reason: syntax error at or near "BRKN" LINE 42' + ) + + def _side_effect(n, cmd): + raise subprocess.CalledProcessError( + returncode=1, cmd=cmd, stderr=stderr, output="" + ) + + run, deploys = _deploy_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert deploys["n"] == 1 + + +def test_v1_still_runs_the_diff_and_force_recovery(monkeypatch, tmp_path): + """v1 remains the pre-existing diff-and-force resolver, unchanged by the + default flip: it still calls _resolve_all_migrations after a deploy that + applied something. Operators opting back in must get exactly the old path. + """ + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + resolve_called = {"n": 0} + + def fake_resolve(*args, **kwargs): + resolve_called["n"] += 1 + + monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=False) + assert ok is True + assert resolve_called["n"] == 1 diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index fb06ed6b69d..506c58d26b4 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -305,14 +305,16 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None): def test_litellm_proxy_server_config_no_general_settings(): - """Exercises the default (v1) migration resolver.""" + """Exercises the default (v2) migration resolver.""" _run_proxy_server_smoke_test() -def test_litellm_proxy_server_config_no_general_settings_v2_resolver(): - """Exercises the opt-in v2 migration resolver. +def test_litellm_proxy_server_config_no_general_settings_legacy_resolver(): + """Exercises the legacy (v1) migration resolver against a real database. - Runs in a separate CI job against a local Postgres to avoid collisions - with the v1 variant when they share a database. + v2 is the default, so the no-arg test above already covers it. This one is + the only place the v1 opt-out gets real-DB migration plus proxy-boot + coverage, and it runs in a separate CI job against its own Postgres to + avoid collisions with the default variant. """ - _run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"]) + _run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"]) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6ea6f208bb5..ac1defe8d79 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1787,7 +1787,7 @@ class TestRunServerDbSetup: # use_prisma_db_push should be False (default), so use_migrate should be True run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) mock_setup_database.assert_called_with( - use_migrate=True, use_v2_resolver=False + use_migrate=True, use_v2_resolver=True ) # Reset mocks @@ -1802,7 +1802,7 @@ class TestRunServerDbSetup: standalone_mode=False, ) mock_setup_database.assert_called_with( - use_migrate=False, use_v2_resolver=False + use_migrate=False, use_v2_resolver=True ) @patch("subprocess.run") @@ -1869,7 +1869,7 @@ class TestRunServerDbSetup: ) assert exc_info.value.code == 1 mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=False + use_migrate=True, use_v2_resolver=True ) @patch("subprocess.run") @@ -1981,6 +1981,83 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=True ) + @pytest.mark.parametrize( + "argv, env_value, expected_v2", + [ + ([], None, True), + (["--use_v2_migration_resolver"], None, True), + (["--use_legacy_migration_resolver"], None, False), + ([], "false", False), + ([], "true", True), + (["--use_v2_migration_resolver"], "false", True), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_migration_resolver_default_and_opt_out( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + argv, + env_value, + expected_v2, + ): + """The proxy defaults to the v2 resolver, and v1 stays reachable. + + Both opt-out routes matter: --use_legacy_migration_resolver for a CLI + boot, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, + where litellm/proxy/prisma_migration.py calls run_server with a fixed + argv and an env var is the only way in. The deprecated + --use_v2_migration_resolver must still parse so existing commands do + not die on an unknown option, and an explicit flag still beats the env. + """ + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + mock_setup_database.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k + not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + if env_value is not None: + clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + run_server.main( + ["--local", "--skip_server_startup", *argv], standalone_mode=False + ) + + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=expected_v2 + ) + # --- Module-level helpers for worker startup hook tests --- From 7b36bfb96720651c756c911ad5bcdb7a819e2006 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:31:44 -0700 Subject: [PATCH 03/92] fix(proxy-extras): retry transient db push failures, drop a vacuous test `prisma db push` under v2 raised on the first failure while v1 retried it four times, so making v2 the default silently cost --use_prisma_db_push its retries. It now uses the same transient classification as migrate deploy. The classifier moves onto ProxyExtrasDBManager next to _is_permission_error and _is_idempotent_error, which do the same kind of stderr matching. Replaces a test that claimed to pin the transient classification but fed it a P3009 stderr, which an earlier branch catches, so it passed even when the classifier was mutated to treat everything as transient. The replacement uses an unclassified error and fails on that mutant. Drops a v1 test that duplicated test_v1_default_still_calls_resolve_all_migrations. --- .../litellm_proxy_extras/utils.py | 98 +++++++++------ .../test_setup_database_fail_fast.py | 113 ++++++++---------- tests/test_litellm/proxy/test_proxy_cli.py | 12 +- 3 files changed, 110 insertions(+), 113 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 22da9b834ab..a31016e0f17 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -51,9 +51,9 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile( re.IGNORECASE, ) -_MIGRATE_DEPLOY_ATTEMPTS: Final = 4 +_PRISMA_ATTEMPTS: Final = 4 -_TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( +_TRANSIENT_PRISMA_FAILURES: Final = MappingProxyType( { "deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)", "P1001": "an unreachable database server", @@ -62,24 +62,6 @@ _TRANSIENT_DEPLOY_FAILURES: Final = MappingProxyType( ) -def _transient_deploy_failure(stderr: str) -> str | None: - """Describe why a failed `prisma migrate deploy` is worth retrying, or None. - - These are environment failures, not migration failures: the database is not - up yet, or another instance holds the migration lock. v1 retried every - failed deploy and absorbed them; failing fast on them instead would turn a - database that is ten seconds late into a dead proxy. - """ - return next( - ( - reason - for marker, reason in _TRANSIENT_DEPLOY_FAILURES.items() - if marker in stderr - ), - None, - ) - - PARTITIONED_SPEND_LOGS_PUSH_ERROR = ( "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), " "so its primary key must include the partition key (\"startTime\"). `prisma db push` " @@ -304,6 +286,23 @@ class ProxyExtrasDBManager: env=prisma_env, ) + @staticmethod + def _transient_prisma_failure(stderr: str) -> str | None: + """Why a failed prisma command is worth retrying, or None. + + v1 retried every failure, so it absorbed a database that was not up yet + or another instance holding the migration lock. v2 fails fast, which is + right for a broken migration and wrong for these. + """ + return next( + ( + reason + for marker, reason in _TRANSIENT_PRISMA_FAILURES.items() + if marker in stderr + ), + None, + ) + @staticmethod def _is_permission_error(error_message: str) -> bool: """ @@ -699,20 +698,43 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - subprocess.run( - [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=prisma_command_timeout(), - check=True, - env=_get_prisma_env(), + for attempt in range(_PRISMA_ATTEMPTS): + try: + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=prisma_command_timeout(), + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + return True + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as e: + stderr = e.stderr or "" + transient = ProxyExtrasDBManager._transient_prisma_failure( + stderr + ) + # Re-raise as RuntimeError so proxy_cli.py's + # `except RuntimeError` catches it and exits cleanly. + if transient is None or attempt == _PRISMA_ATTEMPTS - 1: + raise RuntimeError( + f"prisma db push failed.\n\nDetail: {e}" + f"\n\nPrisma error:\n{stderr}" + ) from e + logger.info( + "prisma db push attempt %s failed on %s, retrying. " + "Prisma error:\n%s", + attempt + 1, + transient, + stderr, + ) + time.sleep(random.randrange(5, 15)) + raise RuntimeError( + f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." ) - return True - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as e: - # Re-raise as RuntimeError so proxy_cli.py's - # `except RuntimeError` catches it and exits cleanly. - raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e finally: os.chdir(original_dir) @@ -722,7 +744,7 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - for attempt in range(_MIGRATE_DEPLOY_ATTEMPTS): + for attempt in range(_PRISMA_ATTEMPTS): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -837,17 +859,17 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - transient = _transient_deploy_failure(stderr) + transient = ProxyExtrasDBManager._transient_prisma_failure(stderr) if transient is None: raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - if attempt == _MIGRATE_DEPLOY_ATTEMPTS - 1: + if attempt == _PRISMA_ATTEMPTS - 1: raise RuntimeError( f"Database migration failed after " - f"{_MIGRATE_DEPLOY_ATTEMPTS} attempts on {transient}. " + f"{_PRISMA_ATTEMPTS} attempts on {transient}. " "Check database connectivity and load." f"\n\nPrisma error:\n{stderr}" ) from e @@ -863,7 +885,7 @@ class ProxyExtrasDBManager: continue raise RuntimeError( - f"Database migration failed after {_MIGRATE_DEPLOY_ATTEMPTS} " + f"Database migration failed after {_PRISMA_ATTEMPTS} " "attempts (retry loop exhausted by timeouts or repeated " "idempotent-recovery continues). Check database connectivity, " "load, and _prisma_migrations ledger state." diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index f637e139e86..deb78dcdb30 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -1,10 +1,7 @@ """Regression tests for ProxyExtrasDBManager's v2 migration resolver. -v2 is what the proxy CLI selects by default; v1 stays reachable via -`--use_legacy_migration_resolver` or `USE_V2_MIGRATION_RESOLVER=false`. At the -library level the resolver is picked with the `use_v2_resolver` kwarg, which -still defaults to False so `migrations/run.py` and any direct caller keep their -own explicit choice. +v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver` +kwarg, which still defaults to False for direct callers. """ import os @@ -271,9 +268,7 @@ class _DeployApplied: def _deploy_only(deploy_side_effect): """subprocess.run stand-in that only intercepts `prisma migrate deploy`. - Everything else the resolver shells out to, the Prisma toolchain check - above all, succeeds untouched, so a mock meant for the deploy call cannot - be silently consumed by an earlier subprocess call. + Scoped by argv so the Prisma toolchain check cannot consume the mock first. """ deploys = {"n": 0} @@ -298,13 +293,8 @@ def _prepare_v2_resolver(monkeypatch, tmp_path): def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): - """A deadlock on Prisma's migration advisory lock is transient and retried. - - Several proxy replicas booting against one database race `migrate deploy`, - and Postgres aborts one side. v1 retried any failed deploy, so it rode this - out; v2 classifies unrecognised stderr as unrecoverable and raises, which - with v2 as the default would take a replica's whole boot down. - """ + """v2: replicas racing `migrate deploy` deadlock on Prisma's advisory + lock, which is transient and must be retried rather than kill the boot.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -323,8 +313,8 @@ def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path): def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path): - """The deadlock retry stays bounded: a deadlock that never clears still - raises rather than looping forever or reporting a successful migration.""" + """v2: the deadlock retry is bounded, so a deadlock that never clears + still raises instead of looping or reporting success.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -348,13 +338,7 @@ def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp ], ) def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr): - """A database that is not accepting connections yet is retried, not fatal. - - A proxy and its database starting together race routinely, and v1 rode that - out by retrying every failed deploy. v2 treats unrecognised stderr as - unrecoverable, so without this the default flip would turn a database that - is a few seconds late into a dead proxy instead of a slow boot. - """ + """v2: a database not accepting connections yet is retried, not fatal.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -373,9 +357,8 @@ def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path): - """Retrying connectivity errors must not turn a genuinely unreachable - database into a silent success: after the attempts are spent it still - raises, so the proxy exits instead of serving without its database.""" + """v2: a genuinely unreachable database still raises once the attempts + are spent, rather than passing as a successful migration.""" _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): @@ -395,13 +378,8 @@ def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog): - """Retrying must not swallow why the database was unreachable. - - Prisma's stderr is captured, so if the retry path neither logs it nor puts - it in the final error, an operator (and CI's bad-DATABASE_URL job, which - greps the boot log for the P1001 line) sees four silent retries and no - cause. - """ + """v2: retrying must not swallow Prisma's stderr, which is captured and is + the only place the cause appears for an operator or a boot-log grep.""" _prepare_v2_resolver(monkeypatch, tmp_path) stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`" @@ -422,21 +400,47 @@ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, cap assert "P1001" in caplog.text -def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path): - """The transient classification must stay narrow: a genuinely broken - migration still fails fast on the first attempt rather than being retried - into the same error four times.""" +def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): + """v2: `prisma db push` retries a transient failure like v1 did, so the + default flip does not cost --use_prisma_db_push its retries.""" _prepare_v2_resolver(monkeypatch, tmp_path) - stderr = ( - "Error: P3009\n" - "The `20260101000000_genuinely_broken` migration failed to apply.\n" - 'Reason: syntax error at or near "BRKN" LINE 42' + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + if pushes["n"] == 1: + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + return _DeployApplied() + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False ) + with patch("subprocess.run", side_effect=_run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert ok is True + assert pushes["n"] == 2 + + +def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): + """v2: an unrecognised deploy failure still raises on the first attempt.""" + _prepare_v2_resolver(monkeypatch, tmp_path) def _side_effect(n, cmd): raise subprocess.CalledProcessError( - returncode=1, cmd=cmd, stderr=stderr, output="" + returncode=1, + cmd=cmd, + stderr="Error: relation \"LiteLLM_SpendLogs\" does not exist", + output="", ) run, deploys = _deploy_only(_side_effect) @@ -447,26 +451,3 @@ def test_v2_migration_failure_is_not_treated_as_transient(monkeypatch, tmp_path) assert deploys["n"] == 1 -def test_v1_still_runs_the_diff_and_force_recovery(monkeypatch, tmp_path): - """v1 remains the pre-existing diff-and-force resolver, unchanged by the - default flip: it still calls _resolve_all_migrations after a deploy that - applied something. Operators opting back in must get exactly the old path. - """ - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - (tmp_path / "schema.prisma").write_text("// stub") - - class FakeResult: - stdout = "Applied migration.\n" - stderr = "" - - resolve_called = {"n": 0} - - def fake_resolve(*args, **kwargs): - resolve_called["n"] += 1 - - monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) - monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) - - ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=False) - assert ok is True - assert resolve_called["n"] == 1 diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index ac1defe8d79..dac36965a6c 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2008,15 +2008,9 @@ class TestRunServerDbSetup: env_value, expected_v2, ): - """The proxy defaults to the v2 resolver, and v1 stays reachable. - - Both opt-out routes matter: --use_legacy_migration_resolver for a CLI - boot, and USE_V2_MIGRATION_RESOLVER=false for containerised deploys, - where litellm/proxy/prisma_migration.py calls run_server with a fixed - argv and an env var is the only way in. The deprecated - --use_v2_migration_resolver must still parse so existing commands do - not die on an unknown option, and an explicit flag still beats the env. - """ + """The proxy defaults to v2, and both v1 opt-out routes work: the + flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that + cannot pass one. An explicit flag beats the env var.""" from litellm.proxy.proxy_cli import run_server mock_subprocess_run.return_value = MagicMock(returncode=0) From fbd1339993252ca3c134494ef431a22ce591c83f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:38:27 -0700 Subject: [PATCH 04/92] fix(tests): satisfy the tests-tree ruff config and correct a stale comment Moving the resolver tests under tests/ brings them under ruff-tests.toml, which the package-internal directory they came from was never linted by, so a pre-existing pytest.raises pattern now needs to be a raw string (RUF043). Also corrects the comment on proxy_cli's RuntimeError handler: both resolvers raise on permission failures, not just v2. --- litellm/proxy/proxy_cli.py | 8 ++++---- .../litellm-proxy-extras/test_setup_database_fail_fast.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c1d86344a1e..86f6853a625 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1323,10 +1323,10 @@ def run_server( use_v2_resolver=use_v2_migration_resolver, ) except RuntimeError as e: - # Raised on unrecoverable migration errors: the v2 - # resolver's non-idempotent failures and permission - # issues, and any `prisma db push` against a - # partitioned LiteLLM_SpendLogs. + # Raised on unrecoverable migration errors: permission + # failures from either resolver, the v2 resolver's + # non-idempotent failures, and any `prisma db push` + # against a partitioned LiteLLM_SpendLogs. print( f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m", file=sys.stderr, diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index deb78dcdb30..0b59c5f9430 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -219,7 +219,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error( ) with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises( - RuntimeError, match="Failed to mark migration .* as applied" + RuntimeError, match=r"Failed to mark migration .* as applied" ): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) From c3fc86869d45e2b5057dce53ec7bec0617f84149 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:52:15 -0700 Subject: [PATCH 05/92] test: keep the resolver tests inside the test-quality ceilings The moved fail-fast test carried a sys.path.insert that the uv workspace makes unnecessary, and one pre-existing case asserted nothing beyond "did not raise", so it could not tell a swallowed error from a skipped query. Give it a liveness gate on the connect count instead. Fold the resolver default/opt-out matrix into the existing db-push flag test rather than standing up another patched test, so the flag pair, the env var, and their precedence are covered without new mock scaffolding. --- .../test_setup_database_fail_fast.py | 16 +-- tests/test_litellm/proxy/test_proxy_cli.py | 105 ++++++------------ 2 files changed, 38 insertions(+), 83 deletions(-) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 0b59c5f9430..964355492a1 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -4,20 +4,11 @@ v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver` kwarg, which still defaults to False for direct callers. """ -import os import subprocess -import sys from unittest.mock import patch import pytest -sys.path.insert( - 0, - os.path.abspath( - os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") - ), -) - from litellm_proxy_extras.utils import ( ProxyExtrasDBManager, _max_migration_timestamp, @@ -175,13 +166,16 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): # Simulate an InsufficientPrivilege (subclass of DatabaseError). raise psycopg.errors.InsufficientPrivilege("permission denied") + connects = {"n": 0} + def _fake_connect(*a, **kw): + connects["n"] += 1 return _FakeConn() monkeypatch.setattr("psycopg.connect", _fake_connect) - # Must not raise. - ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) + assert ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) is None + assert connects["n"] == 1, "the failing query must actually have been reached" def test_v2_resolve_specific_migration_failure_raises_runtime_error( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index dac36965a6c..4f205ada609 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1805,6 +1805,39 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=True ) + # Test 3+: the resolver default and both routes back to v1. The flag + # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that + # cannot pass one, where prisma_migration.py fixes the argv. An + # explicit flag beats the env var. + for argv, env_value, expected_v2 in ( + ([], None, True), + (["--use_v2_migration_resolver"], None, True), + (["--use_legacy_migration_resolver"], None, False), + ([], "false", False), + ([], "true", True), + (["--use_v2_migration_resolver"], "false", True), + (["--use_legacy_migration_resolver"], "true", False), + ): + mock_setup_database.reset_mock() + mock_should_update_schema.reset_mock() + mock_should_update_schema.return_value = True + + resolver_env = ( + {"USE_V2_MIGRATION_RESOLVER": env_value} + if env_value is not None + else {} + ) + os.environ.pop("USE_V2_MIGRATION_RESOLVER", None) + with patch.dict(os.environ, resolver_env): + run_server.main( + ["--local", "--skip_server_startup", *argv], + standalone_mode=False, + ) + assert mock_setup_database.call_args.kwargs == { + "use_migrate": True, + "use_v2_resolver": expected_v2, + }, f"argv={argv} env={env_value}" + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") @@ -1981,78 +2014,6 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=True ) - @pytest.mark.parametrize( - "argv, env_value, expected_v2", - [ - ([], None, True), - (["--use_v2_migration_resolver"], None, True), - (["--use_legacy_migration_resolver"], None, False), - ([], "false", False), - ([], "true", True), - (["--use_v2_migration_resolver"], "false", True), - ], - ) - @patch("subprocess.run") - @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") - def test_migration_resolver_default_and_opt_out( - self, - mock_should_update_schema, - mock_check_schema_diff, - mock_setup_database, - mock_atexit_register, - mock_subprocess_run, - argv, - env_value, - expected_v2, - ): - """The proxy defaults to v2, and both v1 opt-out routes work: the - flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that - cannot pass one. An explicit flag beats the env var.""" - from litellm.proxy.proxy_cli import run_server - - mock_subprocess_run.return_value = MagicMock(returncode=0) - mock_should_update_schema.return_value = True - mock_setup_database.return_value = True - - mock_proxy_module = MagicMock( - app=MagicMock(), - ProxyConfig=MagicMock(), - KeyManagementSettings=MagicMock(), - save_worker_config=MagicMock(), - ) - - clean_env = { - k: v - for k, v in os.environ.items() - if k - not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") - } - clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" - if env_value is not None: - clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value - - with ( - patch.dict(os.environ, clean_env, clear=True), - patch.dict( - "sys.modules", - { - "proxy_server": mock_proxy_module, - "litellm.proxy.proxy_server": mock_proxy_module, - }, - ), - ): - run_server.main( - ["--local", "--skip_server_startup", *argv], standalone_mode=False - ) - - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=expected_v2 - ) - - # --- Module-level helpers for worker startup hook tests --- _dummy_hook_called = False From 6b3e30e6601ac4fb285e97525061ccdd93aaac16 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 01:56:15 -0700 Subject: [PATCH 06/92] fix(proxy-extras): bound the db push retries in the reachable branch The retry loop already raises on the final attempt, so the raise that followed the loop could never run. Drop it and cover the exhaustion path with a test that pins the attempt count and keeps the prisma error in the message, which is the only thing that tells an operator why the boot stopped. --- .../litellm_proxy_extras/utils.py | 3 -- .../test_setup_database_fail_fast.py | 35 +++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index a31016e0f17..f751c4087eb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -732,9 +732,6 @@ class ProxyExtrasDBManager: stderr, ) time.sleep(random.randrange(5, 15)) - raise RuntimeError( - f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." - ) finally: os.chdir(original_dir) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 964355492a1..ed82037590a 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -10,6 +10,7 @@ from unittest.mock import patch import pytest from litellm_proxy_extras.utils import ( + _PRISMA_ATTEMPTS, ProxyExtrasDBManager, _max_migration_timestamp, _migration_timestamp, @@ -425,6 +426,40 @@ def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): assert pushes["n"] == 2 +def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error( + monkeypatch, tmp_path +): + """v2: a database that never comes back stops after _PRISMA_ATTEMPTS and + surfaces the prisma error, rather than retrying the boot forever.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr="Error: P1001: Can't reach database server at `db`:`5432`", + output="", + ) + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + with patch("subprocess.run", side_effect=_run): + with pytest.raises(RuntimeError) as exc: + ProxyExtrasDBManager.setup_database( + use_migrate=False, use_v2_resolver=True + ) + + assert pushes["n"] == _PRISMA_ATTEMPTS + assert "P1001" in str(exc.value) + + def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): """v2: an unrecognised deploy failure still raises on the first attempt.""" _prepare_v2_resolver(monkeypatch, tmp_path) From 2495673d01ee26c73d54adb6b78e30ff87cf0bc2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 02:14:57 -0700 Subject: [PATCH 07/92] fix(proxy-extras): stop a db push timeout crashing the migration job subprocess.run leaves stderr as bytes on TimeoutExpired even under text=True, unlike CalledProcessError. Classifying both in one handler meant a real `prisma db push` timeout died on a TypeError, which proxy_cli.py's `except RuntimeError` does not catch, so the migrations Job container ended on an unhandled traceback instead of a clean exit. Give the timeout its own handler and retry it, matching what the migrate deploy loop beside it already does. That puts a fallthrough back into the loop, so the trailing raise removed in the previous commit is reachable again and comes back with it. Also drop a comment restating why the resolver cases exist and widen the db push test's docstring, which had stopped describing what it covers. --- .../litellm_proxy_extras/utils.py | 14 ++- .../test_setup_database_fail_fast.py | 86 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 7 +- 3 files changed, 98 insertions(+), 9 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f751c4087eb..22e30dfa897 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -709,10 +709,13 @@ class ProxyExtrasDBManager: env=_get_prisma_env(), ) return True - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as e: + except subprocess.TimeoutExpired: + logger.info( + "prisma db push attempt %s timed out, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + except subprocess.CalledProcessError as e: stderr = e.stderr or "" transient = ProxyExtrasDBManager._transient_prisma_failure( stderr @@ -732,6 +735,9 @@ class ProxyExtrasDBManager: stderr, ) time.sleep(random.randrange(5, 15)) + raise RuntimeError( + f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." + ) finally: os.chdir(original_dir) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index ed82037590a..3996b91c2a6 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -460,6 +460,92 @@ def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error( assert "P1001" in str(exc.value) +def _db_push_only(push_side_effect): + """subprocess.run stand-in that only intercepts `prisma db push`.""" + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + return push_side_effect(pushes["n"], cmd) + + return _run, pushes + + +def _timed_out_for_real(): + """Capture what subprocess.run really puts on a TimeoutExpired. + + Under text=True it still leaves stderr as bytes, unlike CalledProcessError, + so hardcoding a str here would test a shape production never sees. Derived + at import, before any test patches subprocess.run. + """ + try: + subprocess.run( + ["sh", "-c", "echo 'Error: P1001 unreachable' >&2; sleep 5"], + timeout=0.2, + check=True, + capture_output=True, + text=True, + ) + except subprocess.TimeoutExpired as e: + return e + raise AssertionError("the helper command was supposed to time out") + + +_TIMEOUT_TEMPLATE = _timed_out_for_real() + + +def _real_timeout_expired(cmd): + return subprocess.TimeoutExpired( + cmd=cmd, + timeout=_TIMEOUT_TEMPLATE.timeout, + output=_TIMEOUT_TEMPLATE.stdout, + stderr=_TIMEOUT_TEMPLATE.stderr, + ) + + +def test_v2_db_push_retries_a_timeout(monkeypatch, tmp_path): + """v2: a `prisma db push` that times out is retried, not turned into a + TypeError by classifying its bytes stderr as if it were text.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise _real_timeout_expired(cmd) + return _DeployApplied() + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert ok is True + assert pushes["n"] == 2 + + +def test_v2_db_push_timeouts_are_bounded(monkeypatch, tmp_path): + """v2: a `prisma db push` that never stops timing out gives up as a + RuntimeError, which is the only exception proxy_cli.py exits cleanly on.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise _real_timeout_expired(cmd) + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match=r"prisma db push failed after \d+"): + ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert pushes["n"] == _PRISMA_ATTEMPTS + + def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): """v2: an unrecognised deploy failure still raises on the first attempt.""" _prepare_v2_resolver(monkeypatch, tmp_path) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4f205ada609..5e2dd358d75 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1737,7 +1737,8 @@ class TestRunServerDbSetup: mock_atexit_register, mock_subprocess_run, ): - """Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter""" + """Which resolver and which migration mode run_server hands setup_database, + across the db push flag, the v2/legacy flag pair and USE_V2_MIGRATION_RESOLVER.""" from litellm.proxy.proxy_cli import run_server # Mock subprocess.run to simulate prisma being available @@ -1805,10 +1806,6 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=True ) - # Test 3+: the resolver default and both routes back to v1. The flag - # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that - # cannot pass one, where prisma_migration.py fixes the argv. An - # explicit flag beats the env var. for argv, env_value, expected_v2 in ( ([], None, True), (["--use_v2_migration_resolver"], None, True), From 5dfe32c889ca778b682bcab6e7772a4ebb14d47f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 02:16:20 -0700 Subject: [PATCH 08/92] docs(tests): name the entrypoint that actually reaches the db push branch The proxy CLI's --use_prisma_db_push never gets here; PrismaManager keeps its own db push loop and only delegates when use_migrate is true. The caller this covers is the migrations Job with USE_PRISMA_DB_PUSH=true. --- .../litellm-proxy-extras/test_setup_database_fail_fast.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index 3996b91c2a6..ef447315a8c 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -396,8 +396,11 @@ def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, cap def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path): - """v2: `prisma db push` retries a transient failure like v1 did, so the - default flip does not cost --use_prisma_db_push its retries.""" + """v2: `prisma db push` retries a transient failure like v1 did. + + Reached from the migrations Job (USE_PRISMA_DB_PUSH=true), not from the + proxy CLI, whose --use_prisma_db_push has its own loop in prisma_client. + """ _prepare_v2_resolver(monkeypatch, tmp_path) pushes = {"n": 0} From e22744c4393858357d34b2abfffb0249528880ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:06:29 -0700 Subject: [PATCH 09/92] fix(embeddings): omit encoding_format when the client omits it on OpenAI-compatible calls When no encoding_format is set on the call, the model config, or LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT, leave the field out of the upstream request instead of defaulting to float, and bypass the OpenAI SDK's own base64 default so nothing re-adds it on the wire. Downstreams that reject encoding_format, such as a second LiteLLM proxy fronting Bedrock Titan embeddings, now work when the client omits the field. Fixes #38661 --- litellm/llms/hosted_vllm/embedding/README.md | 5 +- litellm/llms/openai/openai.py | 72 ++++++++++++-------- litellm/main.py | 19 +++--- tests/test_litellm/test_main.py | 61 +++++++++++++++++ 4 files changed, 115 insertions(+), 42 deletions(-) diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index 2c58e16fc23..32b7ea5c560 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -9,8 +9,7 @@ For OpenAI-compatible embedding calls (including `openai/...` with a custom `api 1. Explicit value on the embedding call (`encoding_format=...`). 2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). 3. Environment variable `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT` (e.g. in `.env` or container env). -4. Default **`float`**. -That avoids forwarding `encoding_format=None` to the provider/SDK where some servers behave poorly. +If none of those is set, or the winning value is the literal string `none`, the field is omitted from the upstream request entirely (LiteLLM also bypasses the OpenAI SDK's own base64 default), so OpenAI-compatible servers that reject `encoding_format` keep working. -To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). \ No newline at end of file +To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e66c998acf..7b71e532b61 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -12,9 +12,14 @@ if TYPE_CHECKING: import openai from openai import AsyncOpenAI, OpenAI +from openai._base_client import make_request_options +from openai._constants import RAW_RESPONSE_HEADER +from openai._legacy_response import LegacyAPIResponse +from openai._types import RequestOptions +from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import overload import litellm @@ -322,6 +327,24 @@ class OpenAIChatCompletionResponseIterator(BaseModelResponseIterator): raise e +_EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None) +_EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _embedding_request_without_sdk_defaults( + data: Mapping[str, object], timeout: float | httpx.Timeout +) -> tuple[dict[str, object], RequestOptions]: + body: Final = {k: v for k, v in data.items() if k not in ("extra_headers", "extra_query", "extra_body")} + extra_headers: Final = _EXTRA_HEADERS_ADAPTER.validate_python(data.get("extra_headers")) + options: Final = make_request_options( + extra_headers={**(extra_headers or {}), RAW_RESPONSE_HEADER: "true"}, + extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")), + extra_body=data.get("extra_body"), + timeout=timeout, + ) + return body, options + + class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): def __init__(self) -> None: super().__init__() @@ -1148,19 +1171,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> tuple[dict[str, str], CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = await openai_aclient.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return dict(bypass_response.headers), bypass_response.parse(to=CreateEmbeddingResponse) + raw_response: Final = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) + return dict(raw_response.headers), raw_response.parse() @track_llm_api_timing() def make_sync_openai_embedding_request( @@ -1169,20 +1189,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) - - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> tuple[dict[str, str], CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = openai_client.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return dict(bypass_response.headers), bypass_response.parse(to=CreateEmbeddingResponse) + raw_response: Final = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) + return dict(raw_response.headers), raw_response.parse() async def aembedding( self, diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..f40b0e5c227 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6289,18 +6289,15 @@ def embedding( if headers is not None and headers != {}: optional_params["extra_headers"] = headers - if encoding_format is not None: - optional_params["encoding_format"] = encoding_format + requested_encoding_format: Final = ( + encoding_format + or optional_params.get("encoding_format") + or get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") + ) + if requested_encoding_format is None or requested_encoding_format.strip().lower() == "none": + optional_params.pop("encoding_format", None) else: - env_fmt: Final = get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") - if env_fmt is not None and env_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - _default_fmt: Final = optional_params.get("encoding_format") or env_fmt or "float" - if _default_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - optional_params["encoding_format"] = _default_fmt + optional_params["encoding_format"] = requested_encoding_format api_version = None diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..4c591be0f61 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,3 +3181,64 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None + + +def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +def test_embedding_openai_env_var_sets_default_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "float" + + +@pytest.mark.asyncio +async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] From c254605e924455062e72305fdfaa2bd9c9feb089 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:10:21 -0700 Subject: [PATCH 10/92] test(embeddings): move encoding_format default coverage to wire-level assertions Consolidate the new regression tests into test_openai_embedding_encoding_format_default.py, replacing mocks that pinned the old float default with respx captures of the request body, and update the stale local_testing default-float test to assert omission --- tests/local_testing/test_embedding.py | 64 +++--- tests/test_litellm/test_main.py | 60 ------ ...penai_embedding_encoding_format_default.py | 188 ++++++++---------- 3 files changed, 109 insertions(+), 203 deletions(-) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index aed2849f056..ee2ac14f498 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -3,6 +3,8 @@ import os import re import traceback +import httpx + import openai import pytest from dotenv import load_dotenv @@ -1255,56 +1257,42 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): assert sent_data["input"] == expected_payload_input -def test_encoding_format_defaults_to_float_for_openai_sdk(monkeypatch): +def test_encoding_format_omitted_by_default_for_openai_sdk(monkeypatch): """ - When encoding_format is not provided, LiteLLM sends `float` for OpenAI-path embeddings. + When encoding_format is not provided, LiteLLM leaves it out of the upstream request. Optional global override: `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT`. """ monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - # Create a mock client instance - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance + captured_bodies = [] - # Mock the embeddings.with_raw_response.create method - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-ada-002", "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} - - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response + }, ) - # Call the embedding function without encoding_format - response = embedding( - model="text-embedding-ada-002", - input="Hello world", - ) + client = openai.OpenAI( + api_key="sk-test", http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) - # Get the call arguments to verify what was sent to OpenAI SDK - call_args = mock_client_instance.embeddings.with_raw_response.create.call_args - assert ( - call_args is not None - ), "OpenAI SDK embeddings.create should have been called" + response = embedding( + model="text-embedding-ada-002", + input="Hello world", + api_key="sk-test", + client=client, + ) - call_kwargs = call_args[1] # Get kwargs - - assert "encoding_format" in call_kwargs - assert ( - call_kwargs["encoding_format"] == "float" - ), "encoding_format should default to float when not provided by user" - - print("✅ PASS: encoding_format='float' is correctly passed to OpenAI SDK") + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert "encoding_format" not in captured_bodies[0], ( + "encoding_format should be omitted from the upstream request when not provided by user" + ) def test_encoding_format_explicit_value_preserved(): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4c591be0f61..84225c3feb1 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3182,63 +3182,3 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None - -def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: - return respx_mock.post("https://api.openai.com/v1/embeddings").mock( - return_value=httpx.Response( - 200, - json={ - "object": "list", - "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], - "model": "text-embedding-3-small", - "usage": {"prompt_tokens": 2, "total_tokens": 2}, - }, - ) - ) - - -def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert "encoding_format" not in request_body - assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] - - -def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - litellm.embedding( - model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" - ) - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert request_body["encoding_format"] == "base64" - - -def test_embedding_openai_env_var_sets_default_encoding_format( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert request_body["encoding_format"] == "float" - - -@pytest.mark.asyncio -async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - mock_route: Final = _mock_openai_embedding_route(respx_mock) - - response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - - request_body: Final = json.loads(mock_route.calls.last.request.read()) - assert "encoding_format" not in request_body - assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py index 94e4e3c81e5..9842bf30585 100644 --- a/tests/test_litellm/test_openai_embedding_encoding_format_default.py +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -1,124 +1,102 @@ -from unittest.mock import MagicMock, patch +import json +from typing import Final +import httpx import pytest +import respx -from litellm import embedding +import litellm -@pytest.mark.parametrize( - "set_env, env_value, expected", - [ - (False, None, "float"), - (True, "base64", "base64"), - ], -) -def test_openai_embedding_encoding_format_default( - monkeypatch, set_env, env_value, expected -): - monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - if set_env: - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) - - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } +def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) ) - mock_response.headers = {} - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) +@pytest.fixture(autouse=True) +def clear_default_encoding_format_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == expected + +def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +def test_embedding_openai_explicit_encoding_format_wins_over_env_var( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +@pytest.mark.parametrize("env_value", ["float", "base64"]) +def test_embedding_openai_env_var_sets_default_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_value: str +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == env_value @pytest.mark.parametrize("env_none", ["none", "NONE", " none "]) -def test_openai_embedding_encoding_format_env_none_omits_param( - monkeypatch, env_none -): - """LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT=none omits encoding_format (provider default).""" +def test_embedding_openai_env_none_omits_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_none: str +) -> None: monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_none) + mock_route: Final = _mock_openai_embedding_route(respx_mock) - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert "encoding_format" not in call_kwargs + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body -def test_openai_embedding_encoding_format_explicit_overrides_env(monkeypatch): - """Request `encoding_format` wins over LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT.""" - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") +@pytest.mark.asyncio +async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + mock_route: Final = _mock_openai_embedding_route(respx_mock) - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} + response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - - embedding( - model="text-embedding-ada-002", - input="Hello world", - encoding_format="base64", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == "base64" + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] From b65592e623805749827d03b1fcacf03173a5d16d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:18:19 -0700 Subject: [PATCH 11/92] test: drop stray trailing blank line in test_main.py --- tests/test_litellm/test_main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 84225c3feb1..8cf878d05d9 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,4 +3181,3 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None - From 8e687f10047ce7f0454204b94ed08e7bac2151b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:04:44 -0700 Subject: [PATCH 12/92] test(embeddings): move legacy intercepts to the wire for the omitted-format path The omitted-format path deliberately no longer dispatches through embeddings.create, so four legacy tests now intercept at the transport or client.post instead. Also adds a bypass error-path unit test, rewords a stale comment and a README scope note, and ratchets the lint budgets down. --- basedpyright-code-budget.json | 2 +- litellm/llms/hosted_vllm/embedding/README.md | 2 +- litellm/utils.py | 8 +-- ruff-strict-budget.json | 8 +-- test-quality-budget.json | 4 +- .../test_litellm_proxy_provider.py | 71 +++++++++++-------- tests/llm_translation/test_nvidia_nim.py | 48 ++++++++----- tests/local_testing/test_exceptions.py | 17 +++-- tests/local_testing/test_router.py | 4 +- ...penai_embedding_encoding_format_default.py | 19 +++++ type-discipline-budget.json | 2 +- 11 files changed, 117 insertions(+), 68 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 962d1266fd7..ba57fd1278b 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38721 + "limit": 38720 }, "reportUnknownParameterType": { "limit": 19778 diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index 32b7ea5c560..50474aabdeb 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -4,7 +4,7 @@ VLLM is a superset of OpenAI's `embedding` endpoint. ## `encoding_format` -For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request: +For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request. `hosted_vllm/...` models use a separate handler that never adds the field on its own, so this resolution applies to the `openai/...`-style routes only: 1. Explicit value on the embedding call (`encoding_format=...`). 2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). diff --git a/litellm/utils.py b/litellm/utils.py index fa2226dbf2c..66a07d2e2db 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3542,10 +3542,10 @@ def get_optional_params_embeddings( non_default_params=non_default_params, optional_params={}, kwargs=kwargs ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": - # OpenAI SDKs (and litellm's own client) send encoding_format="float" - # by default; float lists are exactly what the vertex API returns, so - # the param is a no-op — don't reject the provider default. Other - # values (e.g. "base64") stay on the unsupported-param path below. + # OpenAI SDKs send encoding_format="float" by default; float lists are + # exactly what the vertex API returns, so the param is a no-op — don't + # reject the provider default. Other values (e.g. "base64") stay on + # the unsupported-param path below. if non_default_params.get("encoding_format") == "float": non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c60988eccc0..73c69a732fe 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2003 + "limit": 2001 }, "ANN202": { "limit": 845 @@ -240,13 +240,13 @@ "limit": 96 }, "TRY201": { - "limit": 405 + "limit": 403 }, "TRY203": { - "limit": 113 + "limit": 111 }, "TRY300": { - "limit": 857 + "limit": 855 }, "UP028": { "limit": 2 diff --git a/test-quality-budget.json b/test-quality-budget.json index ee33eb581d6..d834c581609 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 742 + "limit": 741 }, "TQ003": { "limit": 62 @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11139 + "limit": 11135 } } diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 1cb805bf9ba..8630259877d 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -5,6 +5,7 @@ from io import BytesIO from unittest.mock import AsyncMock +import httpx import litellm from litellm import completion, embedding import pytest @@ -92,44 +93,54 @@ async def test_litellm_gateway_from_sdk_embedding(is_async): litellm.set_verbose = True litellm._turn_on_debug() + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "my-vllm-model", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + if is_async: from openai import AsyncOpenAI - openai_client = AsyncOpenAI(api_key="fake-key") - mock_method = AsyncMock() - patch_target = openai_client.embeddings.create + openai_client = AsyncOpenAI( + api_key="fake-key", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + response = await litellm.aembedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) else: from openai import OpenAI - openai_client = OpenAI(api_key="fake-key") - mock_method = MagicMock() - patch_target = openai_client.embeddings.create + openai_client = OpenAI( + api_key="fake-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + response = litellm.embedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) - with patch.object(patch_target.__self__, patch_target.__name__, new=mock_method): - try: - if is_async: - await litellm.aembedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - else: - litellm.embedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - except Exception as e: - print(e) + request_body = captured_bodies[0] + print("Request body - {}".format(request_body)) - mock_method.assert_called_once() - - print("Call KWARGS - {}".format(mock_method.call_args.kwargs)) - - assert "Hello world" == mock_method.call_args.kwargs["input"] - assert "my-vllm-model" == mock_method.call_args.kwargs["model"] + assert "Hello world" == request_body["input"] + assert "my-vllm-model" == request_body["model"] + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] @pytest.mark.parametrize("is_async", [False, True]) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 7ee4f347f72..d5942e674d0 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -63,27 +63,39 @@ def test_embedding_nvidia_nim(): litellm.set_verbose = True from openai import OpenAI + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "nvidia/nv-embedqa-e5-v5", + "usage": {"prompt_tokens": 6, "total_tokens": 6}, + }, + ) + client = OpenAI( api_key="fake-api-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), ) - with patch.object(client.embeddings.with_raw_response, "create") as mock_client: - try: - litellm.embedding( - model="nvidia_nim/nvidia/nv-embedqa-e5-v5", - input="What is the meaning of life?", - input_type="passage", - dimensions=1024, - client=client, - ) - except Exception as e: - print(e) - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert request_body["input"] == "What is the meaning of life?" - assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" - assert request_body["extra_body"]["input_type"] == "passage" - assert request_body["dimensions"] == 1024 + response = litellm.embedding( + model="nvidia_nim/nvidia/nv-embedqa-e5-v5", + input="What is the meaning of life?", + input_type="passage", + dimensions=1024, + client=client, + ) + request_body = captured_bodies[0] + print("request_body: ", request_body) + assert request_body["input"] == "What is the meaning of life?" + assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" + assert request_body["input_type"] == "passage" + assert request_body["dimensions"] == 1024 + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] def test_chat_completion_nvidia_nim_with_tools(): diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 8370046446d..e6392cda406 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -5,7 +5,7 @@ import traceback from typing import Any import httpx -from openai import AsyncOpenAI, AuthenticationError, BadRequestError, OpenAIError, RateLimitError +from openai import AsyncAzureOpenAI, AsyncOpenAI, AuthenticationError, AzureOpenAI, BadRequestError, OpenAIError, RateLimitError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -895,7 +895,12 @@ def _pre_call_utils( ): if call_type == "embedding": data["input"] = "Hello world!" - mapped_target: Any = client.embeddings.with_raw_response + if isinstance(client, (AzureOpenAI, AsyncAzureOpenAI)): + mapped_target: Any = client.embeddings.with_raw_response + patched_attr = "create" + else: + mapped_target = client + patched_attr = "post" if sync_mode: original_function = litellm.embedding else: @@ -905,6 +910,7 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.chat.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.completion else: @@ -914,12 +920,13 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.text_completion else: original_function = litellm.atext_completion - return data, original_function, mapped_target + return data, original_function, mapped_target, patched_attr def _pre_call_utils_httpx( @@ -1003,7 +1010,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str ) data = {"model": model} - data, original_function, mapped_target = _pre_call_utils( + data, original_function, mapped_target, patched_attr = _pre_call_utils( call_type=call_type, data=data, client=openai_client, @@ -1049,7 +1056,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str with patch.object( mapped_target, - "create", + patched_attr, side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 370c43f8f44..c714bb4f9a7 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -2032,8 +2032,8 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): raise exception with patch.object( - openai_client.embeddings.with_raw_response, - "create", + openai_client, + "post", side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py index 9842bf30585..7a42eaf0f0a 100644 --- a/tests/test_litellm/test_openai_embedding_encoding_format_default.py +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -100,3 +100,22 @@ async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_omitted_encoding_format_maps_provider_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 429, + headers={"retry-after": "42", "x-should-retry": "false"}, + json={"error": {"message": "rate limited", "type": "rate_limit_error"}}, + ) + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", max_retries=0 + ) + + assert int(exc_info.value.litellm_response_headers["retry-after"]) == 42 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3f1a7defe7..67ebb2b3b17 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16564 + "limit": 16562 }, "LIT011": { "limit": 5577 From 2a88384e4ec2d2f183fb69f223ed376b2eb24991 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:15:32 -0700 Subject: [PATCH 13/92] style(utils): drop an em-dash from the vertex encoding_format comment --- litellm/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 66a07d2e2db..c4250dd0e4b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3543,9 +3543,9 @@ def get_optional_params_embeddings( ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": # OpenAI SDKs send encoding_format="float" by default; float lists are - # exactly what the vertex API returns, so the param is a no-op — don't - # reject the provider default. Other values (e.g. "base64") stay on - # the unsupported-param path below. + # exactly what the vertex API returns, so the param is a no-op and the + # provider default is not rejected. Other values (e.g. "base64") stay + # on the unsupported-param path below. if non_default_params.get("encoding_format") == "float": non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( From 855f56fa946674c3a25ee873fca6a09d4783edcd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:25:28 -0700 Subject: [PATCH 14/92] fix(openai): flatten top-level tool schema combinators on chat completions --- .../llms/openai/chat/gpt_transformation.py | 82 ++++++++++-- .../chat/test_openai_gpt_transformation.py | 118 ++++++++++++++++++ 2 files changed, 190 insertions(+), 10 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 5894658e5d2..89062f60fc7 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -4,7 +4,8 @@ Support for gpt model family import json import os -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload from urllib.parse import urlparse @@ -19,6 +20,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, + flatten_top_level_schema_combinators, get_tool_call_names, hoist_images_from_tool_messages, ) @@ -65,6 +67,22 @@ else: LiteLLMLoggingObj = Any +_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = tool.get("function") + if not isinstance(function, dict): + return tool + parameters: Final = function.get("parameters") + if not isinstance(parameters, dict): + return tool + flattened: Final = flatten_top_level_schema_combinators(parameters) + if flattened is parameters: + return tool + return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + + class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Reference: https://platform.openai.com/docs/api-reference/chat/create @@ -393,6 +411,26 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) return messages, tools + def _targets_openai_hosted_endpoint( + self, + custom_llm_provider: str | None, + api_base: str | None, + ) -> bool: + """ + True only for the generic `openai` provider actually pointed at + api.openai.com (no custom api_base, or an openai.com host): the one + backend enforcing OpenAI-only request strictness. + """ + if custom_llm_provider != "openai": + return False + resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") + if not resolved_api_base: + return True + hostname: Final = urlparse(resolved_api_base).hostname + if hostname is None: + return True + return hostname == "openai.com" or hostname.endswith(".openai.com") + def _should_preserve_cache_control_for_endpoint( self, custom_llm_provider: str | None, @@ -404,15 +442,37 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): api_base. Those can understand cache_control, so it must survive there. Real OpenAI cannot, so it is still stripped for an openai.com host. """ - if custom_llm_provider != "openai": - return False - resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") - if not resolved_api_base: - return False - hostname: Final = urlparse(resolved_api_base).hostname - if hostname is None: - return False - return hostname != "openai.com" and not hostname.endswith(".openai.com") + return custom_llm_provider == "openai" and not self._targets_openai_hosted_endpoint( + custom_llm_provider, api_base + ) + + def _flattened_tools_update_for_openai( + self, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> Mapping[str, object]: + """ + OpenAI's chat completions validator rejects tool `parameters` carrying + 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every + model family (unlike the Responses API, where GPT-5+ accepts them), so + tool schemas bound for api.openai.com get their top-level combinators + flattened; OpenAI-compatible backends on a custom api_base accept the + caller's schema as-is and keep it. + """ + tools: Final = optional_params.get("tools") + if not isinstance(tools, list): + return _NO_TOOLS_UPDATE + provider: Final = litellm_params.get("custom_llm_provider") + raw_api_base: Final = litellm_params.get("api_base") + if not self._targets_openai_hosted_endpoint( + provider if isinstance(provider, str) else None, + raw_api_base if isinstance(raw_api_base, str) else None, + ): + return _NO_TOOLS_UPDATE + flattened: Final = [ # mutable-ok: request tools are a JSON list + _tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + ] + return MappingProxyType({"tools": flattened}) def transform_request( self, @@ -444,6 +504,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": messages, **optional_params, + **self._flattened_tools_update_for_openai(optional_params, litellm_params), } async def async_transform_request( @@ -473,6 +534,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": transformed_messages, **optional_params, + **self._flattened_tools_update_for_openai(optional_params, litellm_params), } else: ## allow for any object specific behaviour to be handled diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 3f346b5e8e7..e65ef239068 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -975,3 +975,121 @@ class TestOpenAIPromptCacheBreakpointChatPath: assert request["messages"][1]["content"] == [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}] assert request["extra_body"] == {"prompt_cache_options": self.EXPLICIT} assert "prompt_cache_options" not in request + + +class TestToolSchemaCombinatorFlatteningForOpenAI: + """ + Regression tests for LIT-6488: OpenAI's chat completions validator rejects + tool parameters carrying a top-level anyOf/oneOf/allOf for every model + family (GPT-5 included, unlike the Responses API), so requests bound for + api.openai.com get those combinators flattened into one object schema, + while OpenAI-compatible backends on a custom api_base and other providers + keep the caller's schema untouched. + """ + + def setup_method(self): + self.config = OpenAIGPTConfig() + + @pytest.fixture(autouse=True) + def _clean_openai_base_env(self, monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None, raising=False) + + @staticmethod + def _anyof_tool(): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def _transform(self, config, model, litellm_params, tools): + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools}, + litellm_params=litellm_params, + headers={}, + ) + + def test_flattens_top_level_anyof_for_hosted_openai(self): + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [self._anyof_tool()] + ) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert request["tools"][0]["function"]["name"] == "automation_update" + + def test_gpt5_family_flattens_on_chat_completions(self): + request = self._transform( + OpenAIGPT5Config(), "gpt-5.6", {"custom_llm_provider": "openai", "api_base": None}, [self._anyof_tool()] + ) + assert "anyOf" not in request["tools"][0]["function"]["parameters"] + + def test_custom_api_base_keeps_union(self): + tool = self._anyof_tool() + request = self._transform( + self.config, + "gpt-4o", + {"custom_llm_provider": "openai", "api_base": "http://localhost:8000/v1"}, + [tool], + ) + assert request["tools"][0]["function"]["parameters"] == self._anyof_tool()["function"]["parameters"] + + def test_non_openai_provider_keeps_union(self): + request = self._transform( + self.config, "some-oss-model", {"custom_llm_provider": "groq", "api_base": None}, [self._anyof_tool()] + ) + assert request["tools"][0]["function"]["parameters"] == self._anyof_tool()["function"]["parameters"] + + def test_caller_tool_dict_is_not_mutated(self): + tool = self._anyof_tool() + self._transform(self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool]) + assert tool == self._anyof_tool() + + def test_clean_object_schema_passes_through_as_same_object(self): + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool] + ) + assert request["tools"][0] is tool + + @pytest.mark.asyncio + async def test_async_transform_request_flattens_for_hosted_openai(self): + request = await self.config.async_transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": [self._anyof_tool()]}, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} From 1c4674441c52afe22dbbe49402529fb67c78147c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:43:05 -0700 Subject: [PATCH 15/92] docs(openai): trim tool-flattening docstrings to upstream facts --- litellm/llms/openai/chat/gpt_transformation.py | 10 +--------- .../llms/openai/chat/test_openai_gpt_transformation.py | 5 +---- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 89062f60fc7..a81884702e9 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -416,11 +416,6 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): custom_llm_provider: str | None, api_base: str | None, ) -> bool: - """ - True only for the generic `openai` provider actually pointed at - api.openai.com (no custom api_base, or an openai.com host): the one - backend enforcing OpenAI-only request strictness. - """ if custom_llm_provider != "openai": return False resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") @@ -454,10 +449,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ OpenAI's chat completions validator rejects tool `parameters` carrying 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every - model family (unlike the Responses API, where GPT-5+ accepts them), so - tool schemas bound for api.openai.com get their top-level combinators - flattened; OpenAI-compatible backends on a custom api_base accept the - caller's schema as-is and keep it. + model family, unlike the Responses API, where GPT-5+ accepts them. """ tools: Final = optional_params.get("tools") if not isinstance(tools, list): diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index e65ef239068..19b245449a3 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -981,10 +981,7 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: """ Regression tests for LIT-6488: OpenAI's chat completions validator rejects tool parameters carrying a top-level anyOf/oneOf/allOf for every model - family (GPT-5 included, unlike the Responses API), so requests bound for - api.openai.com get those combinators flattened into one object schema, - while OpenAI-compatible backends on a custom api_base and other providers - keep the caller's schema untouched. + family, GPT-5 included, unlike the Responses API. """ def setup_method(self): From 5c7e6b80c9e4274b9582e7ffc898f91b224f5351 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:18:15 +0000 Subject: [PATCH 16/92] test: isolate global MCP registry and pin savings tests to bundled cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/conftest.py | 31 +++++++++++++++++++ .../proxy/spend_tracking/test_savings.py | 2 ++ 2 files changed, 33 insertions(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index b477bf3f406..c559e023c47 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -2,6 +2,37 @@ import os import pytest +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, +) + + +@pytest.fixture(autouse=True) +def _hermetic_mcp_server_registry(): + """Snapshot and restore the global manager's server-registry state around every test. + + ``global_mcp_server_manager`` is a module-global singleton, and many tests in this + package seed ``registry``/``config_mcp_servers`` (or clear them) without cleaning up. + In a shared CI shard the leaked entries poison later tests in the same worker, e.g. + the ``all_proxy_servers`` sentinel expansion in ``auth/`` suddenly sees a bridge + server registered by a discovery test, so the outcome depends on xdist scheduling. + Restoring the state here makes ordering irrelevant. + """ + saved_registry = dict(global_mcp_server_manager.registry) + saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) + saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) + saved_oauth_slots = global_mcp_server_manager._oauth_discovery_slots + try: + yield + finally: + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(saved_registry) + global_mcp_server_manager.config_mcp_servers.clear() + global_mcp_server_manager.config_mcp_servers.update(saved_config_servers) + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.update(saved_tool_mapping) + global_mcp_server_manager._oauth_discovery_slots = saved_oauth_slots + @pytest.fixture(autouse=True) def _hermetic_server_root_path(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index c3297ee6ae9..7dd18587df3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -14,6 +14,8 @@ from litellm.proxy.spend_tracking.savings import ( from litellm.router import Router from litellm.types.utils import Usage +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + def _anthropic_costs(model: str) -> tuple[float, float]: info = litellm.get_model_info(model=model, custom_llm_provider="anthropic") From dd031f1036eaff49c428da4a133e904b7b3a702f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:42:07 +0000 Subject: [PATCH 17/92] fix(ci): parse paginated gh api output without splitting on unicode line breaks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/close_duplicate_issues.py | 30 +++++++----- .../test_github_close_duplicate_issues.py | 46 +++++++++++++++++++ 2 files changed, 64 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/test_github_close_duplicate_issues.py diff --git a/.github/scripts/close_duplicate_issues.py b/.github/scripts/close_duplicate_issues.py index ec522af4f88..4c837c06418 100755 --- a/.github/scripts/close_duplicate_issues.py +++ b/.github/scripts/close_duplicate_issues.py @@ -39,6 +39,23 @@ def gh(*args: str) -> str: return result.stdout +def parse_concatenated_json(raw: str) -> list[dict]: + """Parse the concatenated JSON documents that `gh api --paginate` emits.""" + decoder = json.JSONDecoder() + issues: list[dict] = [] + idx = 0 + while idx < len(raw): + if raw[idx].isspace(): + idx += 1 + continue + parsed, idx = decoder.raw_decode(raw, idx) + if isinstance(parsed, list): + issues.extend(parsed) + else: + issues.append(parsed) + return issues + + def fetch_open_issues(repo: str | None) -> list[dict]: """Fetch all open issues (excluding PRs) via gh api --paginate.""" if repo: @@ -49,18 +66,7 @@ def fetch_open_issues(repo: str | None) -> list[dict]: endpoint = "repos/{owner}/{repo}/issues?state=open&per_page=100&sort=created&direction=asc" cmd = ["api", "--paginate", endpoint] - raw = gh(*cmd) - # gh --paginate concatenates JSON arrays, so we may get multiple arrays - issues = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - parsed = json.loads(line) - if isinstance(parsed, list): - issues.extend(parsed) - else: - issues.append(parsed) + issues = parse_concatenated_json(gh(*cmd)) # Filter out pull requests (they also appear in the issues endpoint) return [i for i in issues if "pull_request" not in i] diff --git a/tests/test_litellm/test_github_close_duplicate_issues.py b/tests/test_litellm/test_github_close_duplicate_issues.py new file mode 100644 index 00000000000..0ee9b3f096d --- /dev/null +++ b/tests/test_litellm/test_github_close_duplicate_issues.py @@ -0,0 +1,46 @@ +"""Unit tests for `.github/scripts/close_duplicate_issues.py`.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / ".github" + / "scripts" + / "close_duplicate_issues.py" +) + + +@pytest.fixture(scope="module") +def dedupe_module(): + spec = importlib.util.spec_from_file_location("close_duplicate_issues", SCRIPT_PATH) + assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" + module = importlib.util.module_from_spec(spec) + sys.modules["close_duplicate_issues"] = module + spec.loader.exec_module(module) + return module + + +def test_parse_concatenated_json_joins_paginated_arrays(dedupe_module): + page_one = json.dumps([{"number": 1, "title": "a"}, {"number": 2, "title": "b"}]) + page_two = json.dumps([{"number": 3, "title": "c"}]) + issues = dedupe_module.parse_concatenated_json(page_one + page_two) + assert [i["number"] for i in issues] == [1, 2, 3] + + +@pytest.mark.parametrize("separator", ["\u2028", "\u2029", "\x85"]) +def test_parse_concatenated_json_survives_unicode_line_breaks_in_bodies( + dedupe_module, separator +): + body = f"first{separator}second" + page_one = json.dumps([{"number": 1, "title": "a", "body": body}]) + page_two = json.dumps([{"number": 2, "title": "b", "body": "plain"}]) + issues = dedupe_module.parse_concatenated_json(page_one + page_two) + assert [i["number"] for i in issues] == [1, 2] + assert issues[0]["body"] == body From 9fc77f12227e36a6d8e86336e1995931659f1c25 Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Tue, 30 Jun 2026 12:10:48 -0500 Subject: [PATCH 18/92] feat(cost): support time-based off-peak pricing in cost calculation Some providers charge different per-token rates depending on the time of day. DeepSeek, for example, has historically discounted its chat and reasoner models during an off-peak window (16:30-00:30 UTC). LiteLLM's cost map only modeled static per-token pricing, so cost tracking could not stay accurate for these providers. This adds optional off-peak pricing to a model entry: input_cost_per_token_off_peak, output_cost_per_token_off_peak, cache_read_input_token_cost_off_peak, and an off_peak_hours_utc window expressed as "HH:MM-HH:MM" in UTC (the window may wrap past midnight). When the current UTC time falls inside the window, the cost calculator uses the off-peak rates and otherwise falls back to the standard rates, so existing models are unaffected. The fields are also accepted as custom pricing on a deployment, so they can be set from the proxy config or the SDK. The window check is a pure function that takes the current time as an argument, which keeps the regression tests deterministic without patching the clock. --- .../litellm_core_utils/llm_cost_calc/utils.py | 72 +++++++++ litellm/types/utils.py | 14 ++ litellm/utils.py | 1 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 152 ++++++++++++++++++ 4 files changed, 239 insertions(+) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 19e3f624268..576efb18bb0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -4,6 +4,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime, timezone from types import MappingProxyType from typing import Any, Final, Literal, TypedDict, cast @@ -290,10 +291,75 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, ) +def _is_within_off_peak_window(off_peak_hours_utc: str | list[str], current_time: datetime | None = None) -> bool: + """Return True if current_time (UTC, defaulting to now) falls inside any off-peak window. + + off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers + with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past + midnight. The start is inclusive and the end is exclusive; malformed windows are ignored. + """ + if current_time is None: + current_time = datetime.now(timezone.utc) + now = current_time.time() + windows = [off_peak_hours_utc] if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc + for window in windows: + try: + start_str, end_str = window.split("-") + start = datetime.strptime(start_str.strip(), "%H:%M").time() + end = datetime.strptime(end_str.strip(), "%H:%M").time() + except (ValueError, AttributeError): + continue + if start <= end: + if start <= now < end: + return True + elif now >= start or now < end: + return True + return False + + +def _coerce_off_peak_rate(value: object, default: float) -> float: + if isinstance(value, bool): + return default + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + return default + return default + + +def _apply_off_peak_pricing( + model_info: ModelInfo, + current_time: datetime | None, + prompt_base_cost: float, + completion_base_cost: float, + cache_read_cost: float, +) -> tuple[float, float, float]: + """Swap in off-peak per-token rates when the current UTC time is inside one of the model's + off_peak_pricing windows. Applied after threshold pricing so the discount is honored rather + than overwritten when a model combines off-peak and above-threshold rates. Any rate left + unset in off_peak_pricing falls back to the standard rate. + """ + off_peak = model_info.get("off_peak_pricing") + if not off_peak: + return prompt_base_cost, completion_base_cost, cache_read_cost + hours_utc = off_peak.get("hours_utc") + if not hours_utc or not _is_within_off_peak_window(hours_utc, current_time): + return prompt_base_cost, completion_base_cost, cache_read_cost + return ( + _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost), + _coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost), + _coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost), + ) + + def _get_token_base_cost( model_info: ModelInfo, usage: Usage, service_tier: str | None = None, + current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, ) -> tuple[float, float, float, float, float]: @@ -345,6 +411,9 @@ def _get_token_base_cost( k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] if not threshold_keys: + prompt_base_cost, completion_base_cost, cache_read_cost = _apply_off_peak_pricing( + model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost + ) return ( prompt_base_cost, completion_base_cost, @@ -451,6 +520,9 @@ def _get_token_base_cost( except Exception: continue + prompt_base_cost, completion_base_cost, cache_read_cost = _apply_off_peak_pricing( + model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost + ) return ( prompt_base_cost, completion_base_cost, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4bf8289d725..6583b125b62 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -193,6 +193,19 @@ class AgenticLoopParams(TypedDict, total=False): """The LLM provider name (e.g., 'bedrock', 'anthropic')""" +class OffPeakPricing(TypedDict, total=False): + """Time-windowed off-peak rates for providers that discount by time of day (e.g. DeepSeek). + + hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them for multiple daily windows; + a window may wrap past midnight. Any rate left unset falls back to the standard rate. + """ + + hours_utc: str | list[str] + input_cost_per_token: float + output_cost_per_token: float + cache_read_input_token_cost: float + + class ModelInfoBase(ProviderSpecificModelInfo, total=False): key: Required[str] # the key in litellm.model_cost which is returned @@ -225,6 +238,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. prompt_cache_min_tokens: int | None + off_peak_pricing: OffPeakPricing | None # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models diff --git a/litellm/utils.py b/litellm/utils.py index 5e9e115ed54..3389b8fcb78 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5842,6 +5842,7 @@ def _get_model_info_helper( cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), + off_peak_pricing=_model_info.get("off_peak_pricing", None), input_cost_per_character=_model_info.get("input_cost_per_character", None), input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), input_cost_per_token_above_200k_tokens=_model_info.get("input_cost_per_token_above_200k_tokens", None), diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3c4121977de..e226e255b05 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, + _is_within_off_peak_window, calculate_cache_writing_cost, generic_cost_per_token, get_token_type_cost_breakdown, @@ -3946,3 +3947,154 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r ) assert cost == expected_cost + + +def test_is_within_off_peak_window_same_day(): + from datetime import datetime, timezone + + window = "09:00-17:00" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 8, 59, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 17, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_wraps_midnight(): + from datetime import datetime, timezone + + window = "16:30-00:30" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 15, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 16, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 30, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_multiple_windows(): + from datetime import datetime, timezone + + # Providers like DeepSeek V4 have more than one daily peak/off-peak window. + windows = ["01:00-05:00", "13:00-16:00"] + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 3, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 14, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is False + # a malformed entry in the list is ignored, valid entries still match + assert _is_within_off_peak_window(["bad", "13:00-16:00"], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window([], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_malformed_returns_false(): + from datetime import datetime, timezone + + now = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + assert _is_within_off_peak_window("not-a-window", now) is False + assert _is_within_off_peak_window("16:30", now) is False + assert _is_within_off_peak_window("25:00-26:00", now) is False + + +def test_get_token_base_cost_applies_off_peak_pricing(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + }, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + assert off_peak[4] == 5e-8 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 1e-6 + assert peak[1] == 2e-6 + assert peak[4] == 1e-7 + + +def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert result[0] == 5e-7 + assert result[1] == 2e-6 + + +def test_get_token_base_cost_off_peak_wins_over_threshold(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_above_200k_tokens": 3e-6, + "output_cost_per_token_above_200k_tokens": 4e-6, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + }, + ) + usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 3e-6 + assert peak[1] == 4e-6 + + +def test_get_model_info_propagates_off_peak_fields(): + model_name = "test-off-peak-model" + off_peak_pricing = { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + } + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": off_peak_pricing, + } + } + ) + info = litellm.get_model_info(model=model_name) + assert info["off_peak_pricing"] == off_peak_pricing From f2c663515cd335482ab8bd29651e68cf102ab4cd Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Sun, 2 Aug 2026 21:34:09 -0500 Subject: [PATCH 19/92] fix(cost): evaluate off-peak windows in UTC for timezone-aware inputs _is_within_off_peak_window used current_time.time(), which drops tzinfo, so a caller passing a non-UTC aware datetime had the window compared against local wall-clock instead of UTC. That silently mispriced off-peak requests. Normalize aware datetimes to UTC before comparing; naive datetimes stay as-is per the documented UTC contract. Added a regression test with a UTC+8 datetime that fails without the fix --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 2 ++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 576efb18bb0..ffa39850a5e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -300,6 +300,8 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | list[str], current_time """ if current_time is None: current_time = datetime.now(timezone.utc) + elif current_time.tzinfo is not None: + current_time = current_time.astimezone(timezone.utc) now = current_time.time() windows = [off_peak_hours_utc] if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc for window in windows: diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index e226e255b05..463a871c6cc 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -3983,6 +3983,19 @@ def test_is_within_off_peak_window_multiple_windows(): assert _is_within_off_peak_window([], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is False +def test_is_within_off_peak_window_normalizes_timezone_aware_input(): + from datetime import datetime, timedelta, timezone + + # A caller may pass a non-UTC aware datetime; the window is UTC and must be + # evaluated in UTC, not against the caller's wall-clock. 09:00 at UTC+8 is + # 01:00 UTC, inside the 01:00-05:00 window. + tz_plus_8 = timezone(timedelta(hours=8)) + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 9, 0, tzinfo=tz_plus_8)) is True + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 12, 0, tzinfo=tz_plus_8)) is True + # 06:00 at UTC+8 is 22:00 UTC the previous day, outside the window + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 6, 0, tzinfo=tz_plus_8)) is False + + def test_is_within_off_peak_window_malformed_returns_false(): from datetime import datetime, timezone From c813386bb2a482a9fb4e5ece08267bfd5d1eb2b3 Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Mon, 10 Aug 2026 21:47:08 -0500 Subject: [PATCH 20/92] refactor(cost): conform off-peak pricing to current lint budgets Rebasing onto litellm_internal_staging picked up stricter ceilings than this branch was written against. Bind the off-peak results to fresh names instead of reassigning the base costs, mark the new locals Final, avoid rebinding the current_time parameter, and make the window parse explicit about UTC so DTZ007, LIT010 and LIT011 all stay within budget --- .../litellm_core_utils/llm_cost_calc/utils.py | 37 +++++++++---------- litellm/types/utils.py | 10 ++--- 2 files changed, 22 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index ffa39850a5e..66f47affb4d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -2,7 +2,7 @@ ## Helper utilities for cost_per_token() import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType @@ -291,24 +291,21 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, ) -def _is_within_off_peak_window(off_peak_hours_utc: str | list[str], current_time: datetime | None = None) -> bool: +def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool: """Return True if current_time (UTC, defaulting to now) falls inside any off-peak window. off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past midnight. The start is inclusive and the end is exclusive; malformed windows are ignored. """ - if current_time is None: - current_time = datetime.now(timezone.utc) - elif current_time.tzinfo is not None: - current_time = current_time.astimezone(timezone.utc) - now = current_time.time() - windows = [off_peak_hours_utc] if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc + reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() + windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc for window in windows: try: start_str, end_str = window.split("-") - start = datetime.strptime(start_str.strip(), "%H:%M").time() - end = datetime.strptime(end_str.strip(), "%H:%M").time() + start = datetime.strptime(start_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time() + end = datetime.strptime(end_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time() except (ValueError, AttributeError): continue if start <= end: @@ -344,10 +341,10 @@ def _apply_off_peak_pricing( than overwritten when a model combines off-peak and above-threshold rates. Any rate left unset in off_peak_pricing falls back to the standard rate. """ - off_peak = model_info.get("off_peak_pricing") + off_peak: Final = model_info.get("off_peak_pricing") if not off_peak: return prompt_base_cost, completion_base_cost, cache_read_cost - hours_utc = off_peak.get("hours_utc") + hours_utc: Final = off_peak.get("hours_utc") if not hours_utc or not _is_within_off_peak_window(hours_utc, current_time): return prompt_base_cost, completion_base_cost, cache_read_cost return ( @@ -413,15 +410,15 @@ def _get_token_base_cost( k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] if not threshold_keys: - prompt_base_cost, completion_base_cost, cache_read_cost = _apply_off_peak_pricing( + off_peak_prompt_cost, off_peak_completion_cost, off_peak_cache_read_cost = _apply_off_peak_pricing( model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost ) return ( - prompt_base_cost, - completion_base_cost, + off_peak_prompt_cost, + off_peak_completion_cost, cache_creation_cost, cache_creation_cost_above_1hr, - cache_read_cost, + off_peak_cache_read_cost, ) # Only sort the threshold keys (typically 1-2 keys instead of 66+) @@ -522,15 +519,15 @@ def _get_token_base_cost( except Exception: continue - prompt_base_cost, completion_base_cost, cache_read_cost = _apply_off_peak_pricing( + discounted_prompt_cost, discounted_completion_cost, discounted_cache_read_cost = _apply_off_peak_pricing( model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost ) return ( - prompt_base_cost, - completion_base_cost, + discounted_prompt_cost, + discounted_completion_cost, cache_creation_cost, cache_creation_cost_above_1hr, - cache_read_cost, + discounted_cache_read_cost, ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6583b125b62..3ab3a2382dc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -200,10 +200,10 @@ class OffPeakPricing(TypedDict, total=False): a window may wrap past midnight. Any rate left unset falls back to the standard rate. """ - hours_utc: str | list[str] - input_cost_per_token: float - output_cost_per_token: float - cache_read_input_token_cost: float + hours_utc: ReadOnly[str | Sequence[str]] + input_cost_per_token: ReadOnly[float] + output_cost_per_token: ReadOnly[float] + cache_read_input_token_cost: ReadOnly[float] class ModelInfoBase(ProviderSpecificModelInfo, total=False): @@ -238,7 +238,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. prompt_cache_min_tokens: int | None - off_peak_pricing: OffPeakPricing | None # time-windowed off-peak rates + off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models From d302301a4e6a250170ef96a621026af78760be4a Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Thu, 13 Aug 2026 17:28:00 -0500 Subject: [PATCH 21/92] test(cost): move off-peak tests beside the related cost tests They sat at the end of the file, which is where everyone else appends too, so this branch picked up a conflict there on nearly every rebase. Grouping them with the other _get_token_base_cost test keeps them clear of that churn and next to the code they cover. Pure move, no test changes --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 328 +++++++++--------- 1 file changed, 164 insertions(+), 164 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 463a871c6cc..f1340df8b69 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -410,6 +410,170 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): assert prompt_base_cost == 9e-6 +def test_is_within_off_peak_window_same_day(): + from datetime import datetime, timezone + + window = "09:00-17:00" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 8, 59, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 17, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_wraps_midnight(): + from datetime import datetime, timezone + + window = "16:30-00:30" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 15, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 16, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 30, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_multiple_windows(): + from datetime import datetime, timezone + + # Providers like DeepSeek V4 have more than one daily peak/off-peak window. + windows = ["01:00-05:00", "13:00-16:00"] + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 3, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 14, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is False + # a malformed entry in the list is ignored, valid entries still match + assert _is_within_off_peak_window(["bad", "13:00-16:00"], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window([], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_normalizes_timezone_aware_input(): + from datetime import datetime, timedelta, timezone + + # A caller may pass a non-UTC aware datetime; the window is UTC and must be + # evaluated in UTC, not against the caller's wall-clock. 09:00 at UTC+8 is + # 01:00 UTC, inside the 01:00-05:00 window. + tz_plus_8 = timezone(timedelta(hours=8)) + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 9, 0, tzinfo=tz_plus_8)) is True + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 12, 0, tzinfo=tz_plus_8)) is True + # 06:00 at UTC+8 is 22:00 UTC the previous day, outside the window + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 6, 0, tzinfo=tz_plus_8)) is False + + +def test_is_within_off_peak_window_malformed_returns_false(): + from datetime import datetime, timezone + + now = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + assert _is_within_off_peak_window("not-a-window", now) is False + assert _is_within_off_peak_window("16:30", now) is False + assert _is_within_off_peak_window("25:00-26:00", now) is False + + +def test_get_token_base_cost_applies_off_peak_pricing(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + }, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + assert off_peak[4] == 5e-8 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 1e-6 + assert peak[1] == 2e-6 + assert peak[4] == 1e-7 + + +def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert result[0] == 5e-7 + assert result[1] == 2e-6 + + +def test_get_token_base_cost_off_peak_wins_over_threshold(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_above_200k_tokens": 3e-6, + "output_cost_per_token_above_200k_tokens": 4e-6, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + }, + ) + usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 3e-6 + assert peak[1] == 4e-6 + + +def test_get_model_info_propagates_off_peak_fields(): + model_name = "test-off-peak-model" + off_peak_pricing = { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + } + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": off_peak_pricing, + } + } + ) + info = litellm.get_model_info(model=model_name) + assert info["off_peak_pricing"] == off_peak_pricing + + def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" @@ -3947,167 +4111,3 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r ) assert cost == expected_cost - - -def test_is_within_off_peak_window_same_day(): - from datetime import datetime, timezone - - window = "09:00-17:00" - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 8, 59, tzinfo=timezone.utc)) is False - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 17, 0, tzinfo=timezone.utc)) is False - - -def test_is_within_off_peak_window_wraps_midnight(): - from datetime import datetime, timezone - - window = "16:30-00:30" - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 15, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 16, 30, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 30, tzinfo=timezone.utc)) is False - assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is False - - -def test_is_within_off_peak_window_multiple_windows(): - from datetime import datetime, timezone - - # Providers like DeepSeek V4 have more than one daily peak/off-peak window. - windows = ["01:00-05:00", "13:00-16:00"] - assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 3, 0, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 14, 30, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is False - # a malformed entry in the list is ignored, valid entries still match - assert _is_within_off_peak_window(["bad", "13:00-16:00"], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is True - assert _is_within_off_peak_window([], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is False - - -def test_is_within_off_peak_window_normalizes_timezone_aware_input(): - from datetime import datetime, timedelta, timezone - - # A caller may pass a non-UTC aware datetime; the window is UTC and must be - # evaluated in UTC, not against the caller's wall-clock. 09:00 at UTC+8 is - # 01:00 UTC, inside the 01:00-05:00 window. - tz_plus_8 = timezone(timedelta(hours=8)) - assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 9, 0, tzinfo=tz_plus_8)) is True - assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 12, 0, tzinfo=tz_plus_8)) is True - # 06:00 at UTC+8 is 22:00 UTC the previous day, outside the window - assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 6, 0, tzinfo=tz_plus_8)) is False - - -def test_is_within_off_peak_window_malformed_returns_false(): - from datetime import datetime, timezone - - now = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) - assert _is_within_off_peak_window("not-a-window", now) is False - assert _is_within_off_peak_window("16:30", now) is False - assert _is_within_off_peak_window("25:00-26:00", now) is False - - -def test_get_token_base_cost_applies_off_peak_pricing(): - from datetime import datetime, timezone - from typing import cast - - from litellm.types.utils import ModelInfo - - model_info = cast( - ModelInfo, - { - "input_cost_per_token": 1e-6, - "output_cost_per_token": 2e-6, - "cache_read_input_token_cost": 1e-7, - "off_peak_pricing": { - "hours_utc": "16:30-00:30", - "input_cost_per_token": 5e-7, - "output_cost_per_token": 1e-6, - "cache_read_input_token_cost": 5e-8, - }, - }, - ) - usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) - - off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 - assert off_peak[4] == 5e-8 - - peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 1e-6 - assert peak[1] == 2e-6 - assert peak[4] == 1e-7 - - -def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): - from datetime import datetime, timezone - from typing import cast - - from litellm.types.utils import ModelInfo - - model_info = cast( - ModelInfo, - { - "input_cost_per_token": 1e-6, - "output_cost_per_token": 2e-6, - "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}, - }, - ) - usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) - - result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert result[0] == 5e-7 - assert result[1] == 2e-6 - - -def test_get_token_base_cost_off_peak_wins_over_threshold(): - from datetime import datetime, timezone - from typing import cast - - from litellm.types.utils import ModelInfo - - model_info = cast( - ModelInfo, - { - "input_cost_per_token": 1e-6, - "output_cost_per_token": 2e-6, - "input_cost_per_token_above_200k_tokens": 3e-6, - "output_cost_per_token_above_200k_tokens": 4e-6, - "off_peak_pricing": { - "hours_utc": "16:30-00:30", - "input_cost_per_token": 5e-7, - "output_cost_per_token": 1e-6, - }, - }, - ) - usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) - - off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 - - peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 3e-6 - assert peak[1] == 4e-6 - - -def test_get_model_info_propagates_off_peak_fields(): - model_name = "test-off-peak-model" - off_peak_pricing = { - "hours_utc": "16:30-00:30", - "input_cost_per_token": 5e-7, - "output_cost_per_token": 1e-6, - "cache_read_input_token_cost": 5e-8, - } - litellm.register_model( - { - model_name: { - "litellm_provider": "openai", - "mode": "chat", - "input_cost_per_token": 1e-6, - "output_cost_per_token": 2e-6, - "off_peak_pricing": off_peak_pricing, - } - } - ) - info = litellm.get_model_info(model=model_name) - assert info["off_peak_pricing"] == off_peak_pricing From 4f174ffdd1022a0ffa21c2ae6bd2154011a9368e Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Sat, 15 Aug 2026 10:49:51 -0500 Subject: [PATCH 22/92] fix(cost): apply off-peak rates on the tiered-pricing path Tiered pricing resolves its own base rates and returns early, before the off-peak swap ran, so a model carrying both tiered_pricing and off_peak_pricing billed the tier rate around the clock. Route every base-cost path through one helper so the window applies wherever the rates came from, and say plainly in the docstring that an off-peak rate replaces the rate it lands on rather than discounting it --- .../litellm_core_utils/llm_cost_calc/utils.py | 63 ++++++++++++------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 32 ++++++++++ 2 files changed, 73 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 66f47affb4d..be289f620cf 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -337,9 +337,10 @@ def _apply_off_peak_pricing( cache_read_cost: float, ) -> tuple[float, float, float]: """Swap in off-peak per-token rates when the current UTC time is inside one of the model's - off_peak_pricing windows. Applied after threshold pricing so the discount is honored rather - than overwritten when a model combines off-peak and above-threshold rates. Any rate left - unset in off_peak_pricing falls back to the standard rate. + off_peak_pricing windows. An off-peak rate replaces the rate that would otherwise apply + rather than discounting it, so a model that also has tiered or above-threshold pricing bills + the flat off-peak rate for the whole request while the window is open. Any rate left unset in + off_peak_pricing falls back to the standard rate. """ off_peak: Final = model_info.get("off_peak_pricing") if not off_peak: @@ -354,6 +355,22 @@ def _apply_off_peak_pricing( ) +def _apply_off_peak_to_base_costs( + model_info: ModelInfo, + current_time: datetime | None, + base_costs: tuple[float, float, float, float, float], +) -> tuple[float, float, float, float, float]: + """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path + produced them. Cache-creation rates are passed through untouched, since off_peak_pricing + has no field for them. + """ + prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs + off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing( + model_info, current_time, prompt, completion, cache_read + ) + return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) + + def _get_token_base_cost( model_info: ModelInfo, usage: Usage, @@ -376,7 +393,7 @@ def _get_token_base_cost( """ tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage) if tiered_base_costs is not None: - return tiered_base_costs + return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs) # Get service tier aware cost keys input_cost_key: Final = _get_service_tier_cost_key("input_cost_per_token", service_tier) @@ -410,15 +427,16 @@ def _get_token_base_cost( k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] if not threshold_keys: - off_peak_prompt_cost, off_peak_completion_cost, off_peak_cache_read_cost = _apply_off_peak_pricing( - model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost - ) - return ( - off_peak_prompt_cost, - off_peak_completion_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - off_peak_cache_read_cost, + return _apply_off_peak_to_base_costs( + model_info, + current_time, + ( + prompt_base_cost, + completion_base_cost, + cache_creation_cost, + cache_creation_cost_above_1hr, + cache_read_cost, + ), ) # Only sort the threshold keys (typically 1-2 keys instead of 66+) @@ -519,15 +537,16 @@ def _get_token_base_cost( except Exception: continue - discounted_prompt_cost, discounted_completion_cost, discounted_cache_read_cost = _apply_off_peak_pricing( - model_info, current_time, prompt_base_cost, completion_base_cost, cache_read_cost - ) - return ( - discounted_prompt_cost, - discounted_completion_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - discounted_cache_read_cost, + return _apply_off_peak_to_base_costs( + model_info, + current_time, + ( + prompt_base_cost, + completion_base_cost, + cache_creation_cost, + cache_creation_cost_above_1hr, + cache_read_cost, + ), ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index f1340df8b69..8f348b36d35 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -574,6 +574,38 @@ def test_get_model_info_propagates_off_peak_fields(): assert info["off_peak_pricing"] == off_peak_pricing +def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): + """Tiered pricing resolves base rates on its own path and returns early, so off-peak has to + be applied there too or a model carrying both would silently bill the tier rate all day.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-tiered" + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "tiered_pricing": [ + {"range": [0, 128000], "input_cost_per_token": 3e-6, "output_cost_per_token": 6e-6}, + ], + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + } + } + ) + info = litellm.get_model_info(model=model_name) + usage = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) + + inside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert inside[:2] == (5e-7, 1e-6) + + outside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert outside[:2] == (3e-6, 6e-6) + + def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" From 27aefade5fc9c9276cf752f949376a32ddedf3a7 Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Sat, 22 Aug 2026 23:11:57 -0500 Subject: [PATCH 23/92] fix(cost): treat an equal-ended off-peak window as the whole day A window whose start equals its end is the natural way to spell off-peak all day, and the docstring's promise that a window may wrap past midnight invites it. It took the non-wrap branch instead, where start <= now < end can never hold, so it matched nothing. It parses cleanly, so it never reached the branch that ignores malformed windows: no exception, no log, and the model billed at standard rates around the clock while the config said otherwise. Let equality fall through to the wrap branch, which covers every instant, and say so in the docstring. Reported by @xyzs996 in review. --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 5 +++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index be289f620cf..30c25253eea 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -296,7 +296,8 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past - midnight. The start is inclusive and the end is exclusive; malformed windows are ignored. + midnight, and a window whose start equals its end covers the whole day. The start is + inclusive and the end is exclusive; malformed windows are ignored. """ reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() @@ -308,7 +309,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ end = datetime.strptime(end_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time() except (ValueError, AttributeError): continue - if start <= end: + if start < end: if start <= now < end: return True elif now >= start or now < end: diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 8f348b36d35..c9ed5936f03 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -431,6 +431,19 @@ def test_is_within_off_peak_window_wraps_midnight(): assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is False +def test_is_within_off_peak_window_equal_start_and_end_covers_whole_day(): + """An equal start and end is the natural way to spell off-peak all day. It used to take the + non-wrap branch, where start <= now < end can never hold, so it matched nothing and billed at + standard rates around the clock without raising or logging anything.""" + from datetime import datetime, timezone + + for window in ("00:00-00:00", "10:00-10:00"): + for hour in range(24): + assert ( + _is_within_off_peak_window(window, datetime(2026, 1, 1, hour, 0, tzinfo=timezone.utc)) is True + ), f"{window} should cover {hour:02d}:00" + + def test_is_within_off_peak_window_multiple_windows(): from datetime import datetime, timezone From cc3ea1fb08387a511b7cc243613df29d41b97062 Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Sat, 22 Aug 2026 23:40:03 -0500 Subject: [PATCH 24/92] docs(cost): state that a naive off-peak current_time is read as UTC An aware value is converted, a naive one is taken to already be UTC rather than localised. Nothing signals the difference, so a caller passing datetime.now() instead of datetime.now(timezone.utc) shifts every window by the host's offset and bills silently wrong. Say so where a caller will read it. Reported by @xyzs996 in review. --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 30c25253eea..21680129ed4 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -298,6 +298,10 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past midnight, and a window whose start equals its end covers the whole day. The start is inclusive and the end is exclusive; malformed windows are ignored. + + An aware current_time is converted to UTC. A naive one is taken to already be UTC rather + than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(), + or every window shifts by the host's offset. """ reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() From 797848dd8243fe10c4bbb58126c97b373ecb87ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:18:59 -0700 Subject: [PATCH 25/92] fix(cost): bill OCR annotation pages via annotation_cost_per_page --- litellm/cost_calculator.py | 18 ++++--- litellm/llms/base_llm/ocr/transformation.py | 1 + .../llms/mistral/ocr/test_mistral_ocr_cost.py | 54 +++++++++++++++++++ 3 files changed, 66 insertions(+), 7 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3adc1c25dfd..d3ad6124751 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1910,12 +1910,15 @@ def ocr_cost( if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: float | None = None - if model_info is not None: - ocr_cost_per_page = model_info.get("ocr_cost_per_page") + ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None + annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed - if pages_processed is None: + annotation_pages: Final = response.usage_info.pages_processed_annotation or 0 + has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0 + + if pages_processed is None and not has_billable_annotation_pages: if cost_per_credit is not None or ocr_cost_per_page is None: # Surface missing usage data instead of silently under-reporting # cost. The previous behavior raised ValueError; we now return 0.0 @@ -1931,7 +1934,7 @@ def ocr_cost( return 0.0, 0.0 raise ValueError("OCR response pages_processed is None") - if ocr_cost_per_page is None: + if ocr_cost_per_page is None and not has_billable_annotation_pages: # No per-page pricing configured. Either the model is on credit-based # pricing (and credits weren't returned, so the credit branch above did # not match) or the model has no OCR pricing entry at all. Surface a @@ -1947,8 +1950,9 @@ def ocr_cost( ) return 0.0, 0.0 - total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed - return total_ocr_processing_cost, 0.0 + ocr_pages_cost: Final = (ocr_cost_per_page or 0.0) * (pages_processed or 0) + annotation_pages_cost: Final = (annotation_rate or 0.0) * annotation_pages + return ocr_pages_cost + annotation_pages_cost, 0.0 def vector_store_search_cost( diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index d1c77186ea8..3b302837032 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -75,6 +75,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" pages_processed: int | None = None + pages_processed_annotation: int | None = None credits: float | None = None doc_size_bytes: int | None = None diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 890df597933..a0e1616d4b2 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -24,6 +24,9 @@ OCR3_MODEL = "mistral/mistral-ocr-2512" OCR3_COST_PER_PAGE = 0.002 OCR3_ANNOTATION_COST_PER_PAGE = 0.003 +AZURE_DOC_AI_MODEL = "azure_ai/mistral-document-ai-2512" +AZURE_DOC_AI_COST_PER_PAGE = 0.003 + def _ocr_response(model: str, pages_processed: int) -> OCRResponse: return OCRResponse( @@ -33,6 +36,14 @@ def _ocr_response(model: str, pages_processed: int) -> OCRResponse: ) +def _annotated_ocr_response(model: str, pages_processed: int | None, annotation_pages: int) -> OCRResponse: + return OCRResponse( + pages=[], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, pages_processed_annotation=annotation_pages), + ) + + @pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) def test_model_info_ocr4_price(model: str) -> None: info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") @@ -79,3 +90,46 @@ def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) call_type="ocr", ) assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) + + +def test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 2, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(2 * OCR3_COST_PER_PAGE + 3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_only_response(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 0, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_pages_when_pages_processed_missing(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", None, 4), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(4 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate(local_model_cost_map) -> None: + info = litellm.get_model_info(model=AZURE_DOC_AI_MODEL, custom_llm_provider="azure_ai") + assert info.get("annotation_cost_per_page") is None + assert info["ocr_cost_per_page"] == AZURE_DOC_AI_COST_PER_PAGE + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-document-ai-2512", 0, 1), + model=AZURE_DOC_AI_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + ) + assert cost == pytest.approx(AZURE_DOC_AI_COST_PER_PAGE) From 7b4b92f54f87fd52835aeeeb90a087d111d3483a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:37:26 +0000 Subject: [PATCH 26/92] fix(registry): update veo 3.1 pricing with resolution tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 26 ++++++++++++------- model_prices_and_context_window.json | 26 ++++++++++++------- 2 files changed, 32 insertions(+), 20 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7071eaa0807..cc963a80d83 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23777,8 +23777,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23792,7 +23794,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23820,8 +23823,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23835,7 +23840,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -43369,8 +43375,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43399,8 +43405,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7071eaa0807..cc963a80d83 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23777,8 +23777,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23792,7 +23794,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23820,8 +23823,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23835,7 +23840,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -43369,8 +43375,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43399,8 +43405,8 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], From 9125a5b7a0af11892dda9260160efcfa52125619 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:39:34 +0000 Subject: [PATCH 27/92] fix(responses): json-encode object tool call arguments in the chat completions bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_tools.py | 17 +++++++ .../streaming_iterator.py | 9 ++-- .../transformation.py | 11 +++-- .../test_litellm_completion_responses.py | 49 +++++++++++++++++++ .../test_streaming_iterator_transformation.py | 34 +++++++++++++ 5 files changed, 111 insertions(+), 9 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index cccae06c74b..bd6abd3f45f 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -45,6 +45,23 @@ def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: return tool_name in custom_tool_names +def serialize_tool_call_arguments(raw_arguments: object, default: str = "") -> str: + """Render tool call arguments as the JSON string tool-call schemas require. + + Arguments normally arrive already JSON-encoded, but clients and providers + also send the decoded object. ``str()`` on a dict yields a Python repr with + single quotes, which every downstream JSON parser rejects with errors like + "Expecting ',' delimiter". + """ + if raw_arguments is None or raw_arguments == "": + return default + if isinstance(raw_arguments, str): + return raw_arguments + if isinstance(raw_arguments, (dict, list, tuple, bool, int, float)): + return json.dumps(raw_arguments) + return str(raw_arguments) + + def unwrap_custom_tool_arguments(arguments: str) -> str: """Extract the raw content string from JSON-wrapped arguments. diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..b2edf2bf9ed 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, + serialize_tool_call_arguments, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -213,10 +214,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args_delta = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args_delta = str(fn.get("arguments") or "") + fn_args_delta = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args_delta = str(getattr(fn, "arguments", "") or "") + fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) output_index = self._get_or_assign_tool_output_index(call_id) @@ -284,10 +285,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args = str(fn.get("arguments") or "") + fn_args = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args = str(getattr(fn, "arguments", "") or "") + fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) # Track if this is a new tool call that wasn't streamed diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f39df38d069..3ca01f2eb74 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -93,6 +93,7 @@ from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, is_custom_tool_call, + serialize_tool_call_arguments, unwrap_custom_tool_arguments, validated_allowed_callers, ) @@ -1010,7 +1011,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], tool_use_type), function=ChatCompletionToolCallFunctionChunk( name=str(function.get("name", "")), - arguments=str(function.get("arguments", "{}")), + arguments=serialize_tool_call_arguments(function.get("arguments"), "{}"), ), index=index, ) @@ -1539,7 +1540,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], _tool_use_definition.get("type") or "function"), function=ChatCompletionToolCallFunctionChunk( name=function.get("name") or "", - arguments=str(function.get("arguments") or ""), + arguments=serialize_tool_call_arguments(function.get("arguments")), ), index=0, ) @@ -1589,7 +1590,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=ChatCompletionToolCallFunctionChunk( name=f"{namespace}__{raw_name}" if qualify else raw_name, - arguments=str(raw_arguments or ""), + arguments=serialize_tool_call_arguments(raw_arguments), ), index=0, ) @@ -2022,7 +2023,7 @@ class LiteLLMCompletionResponsesConfig: function_definition = tool.function tool_name = function_definition.name or "" tool_id = tool.id or "" - tool_arguments = function_definition.get("arguments") or "" + tool_arguments = serialize_tool_call_arguments(function_definition.get("arguments")) # Check if this is a custom tool if is_custom_tool_call(tool_name, custom_tool_names): @@ -2557,7 +2558,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=Function( name=tool_call.get("name") or "", - arguments=tool_call.get("arguments") or "", + arguments=serialize_tool_call_arguments(tool_call.get("arguments")), ), ) 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 b96d2eb5322..16e099f0404 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 @@ -966,6 +966,55 @@ class TestFunctionCallTransformation: assert result[0]["tool_calls"][0]["function"]["arguments"] == "{}" + def test_function_call_transformation_json_encodes_object_arguments(self): + """A decoded arguments object must be JSON-encoded, not str()'d. + + Clients and providers sometimes send `arguments` as an object rather + than a JSON string; `str()` on a dict produces a Python repr with + single quotes, which downstream JSON parsers reject with errors like + "Expecting ',' delimiter". + """ + function_call_item = { + "type": "function_call", + "name": "shell", + "arguments": {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]}, + "call_id": "call_123", + "id": "call_123", + "status": "completed", + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( + function_call=function_call_item + ) + + arguments = result[0].get("tool_calls", [])[0].get("function", {}).get("arguments") + assert json.loads(arguments) == {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]} + assert "'" not in arguments + + def test_create_tool_call_chunk_json_encodes_object_arguments(self): + """Cached tool_call definitions with object arguments stay valid JSON.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={ + "id": "call_456", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls"}}, + }, + tool_call_id="call_456", + index=0, + ) + + assert json.loads(chunk["function"]["arguments"]) == {"command": "ls"} + + def test_create_tool_call_chunk_keeps_empty_arguments_default(self): + """Missing arguments still fall back to an empty JSON object.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={"id": "call_789", "type": "function", "function": {"name": "shell"}}, + tool_call_id="call_789", + index=0, + ) + + assert chunk["function"]["arguments"] == "{}" + def test_complete_input_transformation_with_function_calls(self): """Test the complete transformation with the exact input from the issue""" test_input = [ diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 823f656ddc5..01148f627f1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -10,6 +10,7 @@ before response.completed, and that every event of a bridged stream carries the spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -523,3 +524,36 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids assert len(set(response_ids)) == 1 assert response_ids[0].startswith("resp_") + + +def test_object_tool_call_arguments_stream_as_valid_json(): + """A provider that sends decoded object arguments must still stream valid JSON. + + `str()` on a dict yields a Python repr with single quotes, which clients + parsing function_call_arguments reject with errors like + "Expecting ',' delimiter". + """ + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_obj", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls", "flags": ["-l"]}}, + } + ] + ) + + streamed_arguments = "".join( + evt.delta + for evt in iterator._pending_tool_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ) + + assert json.loads(streamed_arguments) == {"command": "ls", "flags": ["-l"]} From ade21da9d075a9431178594163791eee5ab61e2d Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:54:05 +0000 Subject: [PATCH 28/92] refactor(responses): simplify tool call argument serializer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/custom_tools.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index bd6abd3f45f..90491739bb0 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -53,13 +53,11 @@ def serialize_tool_call_arguments(raw_arguments: object, default: str = "") -> s single quotes, which every downstream JSON parser rejects with errors like "Expecting ',' delimiter". """ - if raw_arguments is None or raw_arguments == "": - return default if isinstance(raw_arguments, str): - return raw_arguments - if isinstance(raw_arguments, (dict, list, tuple, bool, int, float)): - return json.dumps(raw_arguments) - return str(raw_arguments) + return raw_arguments or default + if raw_arguments is None: + return default + return json.dumps(raw_arguments, default=str) def unwrap_custom_tool_arguments(arguments: str) -> str: From c344c7a66beb3ab184182d227eae9fc4c346283d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 19:11:04 +0000 Subject: [PATCH 29/92] fix(registry): add zai/glm-5.2, together Qwen3.8-Flash, cerebras/gemma-4-31b, elevenlabs/scribe_v2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 50 +++++++++++++++++++ model_prices_and_context_window.json | 50 +++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc963a80d83..b7927767eeb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -55046,5 +55046,55 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc963a80d83..b7927767eeb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -55046,5 +55046,55 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } From be5997f3664f7359204fe58f3111a93f24471aee Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:48:16 -0700 Subject: [PATCH 30/92] feat: add Azure AI DeepSeek V4 Flash 0731 pricing --- model_prices_and_context_window.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7071eaa0807..6b4a9818554 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9959,6 +9959,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/embed-v-4-0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", From 1a80b7ae252e017b937dfb9cf665cc04152ce4ec Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:00:39 -0700 Subject: [PATCH 31/92] fix: sync Azure AI model backup registry --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7071eaa0807..6b4a9818554 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9959,6 +9959,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/embed-v-4-0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", From 92b46538e1f83455a038f5ac047e6ca522db8479 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:38:44 -0700 Subject: [PATCH 32/92] fix(policy_engine): restore request guardrails list after pipeline allow --- .../proxy/policy_engine/pipeline_executor.py | 52 +++++++-- .../policy_engine/test_pipeline_executor.py | 100 ++++++++++++++++++ 2 files changed, 142 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 50cb813c6fa..4be0f556ed7 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,6 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Sequence from typing import Any, Final, Literal import litellm @@ -114,11 +115,7 @@ class PipelineExecutor: # Handle terminal actions if action == "allow": - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": return PipelineExecutionResult( @@ -138,11 +135,7 @@ class PipelineExecutor: # action == "next" → continue to next step # Ran out of steps without a terminal action → default allow - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) @staticmethod async def _run_step( @@ -251,6 +244,45 @@ class PipelineExecutor: return None +def _allow_result( + step_results: Sequence[PipelineStepResult], + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> PipelineExecutionResult: + """Build the terminal-allow result, propagating pipeline modifications without the per-step guardrail override.""" + restored: Final = _restore_request_guardrails(working_data, request_data) + return PipelineExecutionResult( + terminal_action="allow", + step_results=list(step_results), # mutable-ok: PipelineExecutionResult field is a list + modified_data=restored if restored != request_data else None, + ) + + +def _restore_request_guardrails( + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> dict: # mutable-ok: merged back into the request dict, which downstream code mutates + """ + Restore the request's own metadata["guardrails"] activation list. + + _run_step overrides it to [step.guardrail] so should_run_guardrail() allows each + step; letting that override escape via modified_data permanently drops every + independently activated guardrail from later lifecycle stages (post_call, etc.). + """ + working_metadata: Final = working_data.get("metadata") + if not isinstance(working_metadata, dict): + return working_data + request_metadata: Final = request_data.get("metadata") + original_guardrails: Final = request_metadata.get("guardrails") if isinstance(request_metadata, dict) else None + stripped: Final = {k: v for k, v in working_metadata.items() if k != "guardrails"} # mutable-ok: request dict + if original_guardrails is not None: + restored: Final = {**stripped, "guardrails": original_guardrails} # mutable-ok: request dict + return {**working_data, "metadata": restored} # mutable-ok: request dict + if not stripped and not isinstance(request_metadata, dict): + return {k: v for k, v in working_data.items() if k != "metadata"} # mutable-ok: request dict + return {**working_data, "metadata": stripped} # mutable-ok: request dict + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. 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 054a5af4148..4fcb7d22588 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -749,6 +749,106 @@ async def test_single_step_pipeline_allow(monkeypatch): assert guard.calls == 1 +@pytest.mark.asyncio +async def test_allow_restores_independent_guardrails_list(monkeypatch): + """ + Request activates an independent guardrail; an unrelated pipeline runs and allows. + Expected: no modified_data escapes, so the request's guardrails list survives + and the independent guardrail still runs at later lifecycle stages (post_call). + Regression: LIT-6587 (pipeline clobbered the list with its last step's guardrail). + """ + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = { + "messages": [{"role": "user", "content": "clean content"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert pipeline_guard.calls == 1 + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert propagated["metadata"]["guardrails"] == ["independent-output-guard"] + assert data["metadata"]["guardrails"] == ["independent-output-guard"] + + +@pytest.mark.asyncio +async def test_allow_does_not_leak_guardrails_into_bare_request(monkeypatch): + """A request without metadata must not gain a metadata.guardrails list from the pipeline.""" + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = {"messages": [{"role": "user", "content": "clean content"}]} + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert "guardrails" not in propagated.get("metadata", {}) + assert "metadata" not in data + + +@pytest.mark.asyncio +async def test_data_forwarding_keeps_changes_and_restores_guardrails_list(monkeypatch): + """A pass_data pipeline's modifications propagate while the request's guardrails list is restored.""" + pii_guard = PiiMaskingGuardrail(guardrail_name="pii-masker") + content_guard = ContentCheckGuardrail(guardrail_name="content-check") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="pii-masker", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="content-check", on_fail="block", on_pass="allow"), + ], + ) + + monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard]) + + data = { + "messages": [{"role": "user", "content": "Hello John Smith"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="pii-then-safety", + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" + assert result.modified_data["metadata"]["guardrails"] == ["independent-output-guard"] + + @pytest.mark.asyncio async def test_step_results_include_duration(monkeypatch): """Step results should include timing information.""" From 191313e756fe606db5dceac80bed72618ff6679c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:22:05 -0700 Subject: [PATCH 33/92] test(websearch): register configured search tool in pre-request hook test PR #38113 made a configured search_tool_name fail fast when the router does not carry a matching search tool, which broke test_pre_request_hook_modifies_request_body: it names test-search-tool but never registers it. Stub the proxy router with that tool so the test exercises the conversion path again. --- .../test_websearch_interception_e2e.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py index 091ea106b91..cc7901b1710 100644 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py @@ -937,6 +937,14 @@ async def test_pre_request_hook_modifies_request_body(): print("✅ WebSearchInterceptionLogger initialized") + mock_router = MagicMock() + mock_router.search_tools = [ + { + "search_tool_name": "test-search-tool", + "litellm_params": {"search_provider": "tavily"}, + } + ] + # Track what actually gets sent to the API captured_request = {} @@ -987,7 +995,7 @@ async def test_pre_request_hook_modifies_request_body(): with patch( "litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler", side_effect=mock_anthropic_messages_handler, - ): + ), patch("litellm.proxy.proxy_server.llm_router", mock_router): print( "\n📝 Making request with native web_search_20250305 tool (stream=True)..." From f65bee6d74322804dea454ef0e6c904a26f89198 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:32:05 -0700 Subject: [PATCH 34/92] test(websearch): carry a reasoned test-quality suppression on the router patch --- .../test_websearch_interception_e2e.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py index cc7901b1710..fd95b7fa8f2 100644 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py @@ -995,7 +995,10 @@ async def test_pre_request_hook_modifies_request_body(): with patch( "litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler", side_effect=mock_anthropic_messages_handler, - ), patch("litellm.proxy.proxy_server.llm_router", mock_router): + ), patch( # test-quality-ok: the hook imports this process-global router at call time; no injection seam exists to register search_tools + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): print( "\n📝 Making request with native web_search_20250305 tool (stream=True)..." From a27e12367e2c3574586128a55e970fa5d17d5379 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 31 Aug 2026 21:50:40 -0700 Subject: [PATCH 35/92] fix(bedrock): forward native structured outputs on Invoke instead of silently inlining the schema --- litellm/llms/anthropic/chat/transformation.py | 24 +- .../anthropic_claude3_transformation.py | 44 +--- litellm/llms/bedrock/common_utils.py | 89 ++++++++ .../anthropic_claude3_transformation.py | 60 ++--- ...odel_prices_and_context_window_backup.json | 24 +- model_prices_and_context_window.json | 24 +- .../test_anthropic_chat_transformation.py | 42 ++++ ...ations_anthropic_claude3_transformation.py | 131 +++++++++-- .../test_anthropic_claude3_transformation.py | 206 +++++++++++++++--- .../llms/bedrock/test_bedrock_common_utils.py | 41 ++++ 10 files changed, 535 insertions(+), 150 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..a3c76d6a29b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1992,19 +1992,35 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return data def _apply_output_config(self, data: dict, model: str, optional_params: dict) -> None: - """Validate and apply output_config to the request data.""" + """Validate and apply output_config to the request data. + + The ``drop_params`` gate here is an effort gate: ``format`` is a + structured-output field, not an effort field, so it survives the drop + and is vetted where it is consumed (the map's + ``supports_native_structured_output`` flag on emission paths). + """ if "output_config" not in optional_params: return output_config: Final = optional_params.get("output_config") if not output_config or not isinstance(output_config, dict): return - if litellm.drop_params is True and not self._model_supports_effort_param(model, self._resolved_provider): + if ( + litellm.drop_params is True + and any(key != "format" for key in output_config) + and not self._model_supports_effort_param(model, self._resolved_provider) + ): litellm.verbose_logger.warning( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, model, ) - optional_params.pop("output_config", None) - data.pop("output_config", None) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + optional_params.pop("output_config", None) + data.pop("output_config", None) + return + format_only: Final = {"format": preserved_format} # mutable-ok: json body + optional_params["output_config"] = format_only # rebind-ok: out-param store + data["output_config"] = format_only # rebind-ok: out-param store return effort: Final = output_config.get("effort") valid_efforts: Final = ["high", "medium", "low", "xhigh", "max"] diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 40b90014f3b..8e709349400 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers -from litellm.litellm_core_utils.litellm_logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -16,17 +15,16 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse -from litellm.utils import _supports_factory if TYPE_CHECKING: import tiktoken @@ -212,36 +210,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("stream_chunk_size", None) - output_format: Final = anthropic_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_request, - ) - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_request, + ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_request, + ) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 72e3cc1b326..df65df642a2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -177,6 +177,95 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages +def _bedrock_model_supports(model: str, key: str) -> bool: + from litellm.utils import _supports_factory + + return _supports_factory(model=model, custom_llm_provider="bedrock", key=key) + + +def apply_bedrock_invoke_structured_output( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Route Anthropic structured-output params to what the Bedrock model supports. + + Consumes the legacy top-level ``output_format`` and the newer + ``output_config.format``, keeping the pre-existing precedence of the legacy + field when a request carries both. Models flagged + ``supports_native_structured_output`` in the model map get the schema + forwarded as ``output_config.format``, which Bedrock relays to the model for + enforced structured output. For every other model the schema is inlined into + the last user message as best-effort text, with a warning because nothing + enforces it. + """ + legacy_output_format: Final = request_body.pop("output_format", None) + output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body) + schema_format: Final = legacy_output_format if isinstance(legacy_output_format, dict) else output_config_format + if schema_format is None: + return + + if _bedrock_model_supports(model, "supports_native_structured_output"): + existing_output_config: Final = request_body.get("output_config") + if isinstance(existing_output_config, dict): + existing_output_config["format"] = schema_format + else: + request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param # mutable-ok: json + return + + verbose_logger.warning( + "Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` " + "in model_prices_and_context_window.json, so the JSON schema was inlined into " + "the last user message and is NOT enforced by the model.", + model, + ) + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=schema_format, + request_body=request_body, + ) + + +def strip_unsupported_bedrock_invoke_output_config_keys( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Drop ``output_config`` keys the Bedrock model does not accept. + + ``format`` survives unconditionally: it is only attached for models whose map + entry advertises ``supports_native_structured_output``. Effort-bearing keys + survive only when the map flags ``supports_output_config`` or a + ``supports_*_reasoning_effort`` tier; otherwise they are dropped with a + warning so Bedrock does not reject the request. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + output_config: Final = request_body.get("output_config") + if not isinstance(output_config, dict): + return + if all(key == "format" for key in output_config): + return + if _bedrock_model_supports(model, "supports_output_config") or AnthropicConfig._model_supports_effort_param( + model, "bedrock" + ): + return + + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` keys for " + "model=%s: neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + request_body.pop("output_config", None) + else: + request_body["output_config"] = {"format": preserved_format} # rebind-ok: out-param # mutable-ok: json + + def normalize_custom_field_on_tools(request_body: dict) -> None: """ Drop the ``custom`` field from each tool, first hoisting a boolean diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index f74a290d773..6ff9f0155f9 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,14 +29,14 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -51,7 +51,6 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk, ModelResponseStream from litellm.types.utils import GenericStreamingChunk as GChunk -from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -708,52 +707,25 @@ class AmazonAnthropicClaudeMessagesConfig( # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) self._remove_ttl_from_cache_control(anthropic_messages_request=anthropic_messages_request, model=model) - # 5. Convert structured-output params to inline schema. - # Bedrock Invoke doesn't support top-level `output_format`; its - # accepted `output_config` subset is also narrower than Anthropic's, so - # consume the newer `output_config.format` shape here instead of - # forwarding it as an unknown nested key. + # 5. Route structured-output params (`output_format` / + # `output_config.format`) to native enforcement or the inline-schema + # fallback, then strip `output_config` keys the model does not accept. + # Ref: https://github.com/BerriAI/litellm/issues/22797 existing_output_config: Final = anthropic_messages_request.get("output_config") if isinstance(existing_output_config, dict): anthropic_messages_request["output_config"] = dict(existing_output_config) - output_format: Final = anthropic_messages_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_messages_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_messages_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_messages_request, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_messages_request, + ) normalize_bedrock_opus_output_config_effort( model=model, output_config=anthropic_messages_request.get("output_config"), ) - - # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, - # but older models do not — strip it to avoid request rejection. - # Ref: https://github.com/BerriAI/litellm/issues/22797 - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_messages_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_messages_request, + ) # 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) # Ref: https://github.com/BerriAI/litellm/issues/22847 @@ -774,9 +746,11 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True - and "output_config" in anthropic_messages_request + and isinstance(remaining_output_config, dict) + and any(key != "format" for key in remaining_output_config) and not AnthropicConfig._model_supports_effort_param(model, "bedrock") ): verbose_logger.warning( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 718e6c489fd..80dc49a770b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1591,7 +1591,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1627,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1663,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1699,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1735,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1771,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2064,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2101,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2138,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2175,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2212,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2249,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 718e6c489fd..80dc49a770b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1591,7 +1591,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1627,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1663,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1699,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1735,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1771,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2064,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2101,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2138,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2175,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2212,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2249,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 25e2c3cda80..c4df46dea83 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6207,3 +6207,45 @@ def test_disabled_thinking_omitted_only_for_always_on_models( assert "thinking" not in request else: assert request["thinking"] == {"type": "disabled"} + + +def test_anthropic_drop_params_keeps_format_only_output_config(monkeypatch): + """``drop_params=True`` must not consume ``output_config.format``: the drop + gate is an effort gate and ``format`` is a structured-output field.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_anthropic_drop_params_reduces_mixed_output_config_to_format(monkeypatch): + """``drop_params=True`` drops the effort key on unsupported models but keeps + ``format`` so structured outputs still reach the provider.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"effort": "low", "format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index cea299280f8..a122d97a0f0 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -428,30 +428,58 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): def test_output_config_format_converted_for_bedrock_chat_invoke_request(): - """Bedrock Invoke chat path consumes ``output_config.format`` before forwarding.""" + """Bedrock Invoke chat path inlines ``output_config.format`` for models + without native structured-output support and keeps the effort key.""" config = AmazonAnthropicClaudeConfig() schema = { "type": "object", "properties": {"answer": {"type": "string"}}, } - result = config.transform_request( + with patch( # test-quality-ok: pin non-native path + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", + ): + result = config.transform_request( + model="anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "test"}], + optional_params={ + "max_tokens": 100, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"effort": "xhigh"} + last_content = result["messages"][0]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(): + """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort + for models with native structured-output support (Claude Opus 4.7).""" + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( model="anthropic.claude-opus-4-7", messages=[{"role": "user", "content": "test"}], optional_params={ "max_tokens": 100, - "output_config": { - "effort": "xhigh", - "format": {"type": "json_schema", "schema": schema}, - }, + "output_config": {"effort": "xhigh", "format": schema_format}, }, litellm_params={}, headers={}, ) - assert result.get("output_config") == {"effort": "xhigh"} - last_content = result["messages"][0]["content"] - assert json.loads(last_content[-1]["text"]) == schema + assert result.get("output_config") == {"effort": "xhigh", "format": schema_format} + assert "answer" not in json.dumps(result["messages"]) @pytest.mark.parametrize( @@ -488,7 +516,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( optional_params = {"max_tokens": 100, "output_config": {"effort": "high"}} with patch( - "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = config.transform_request( @@ -499,11 +527,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( headers={}, ) - mock_supports_factory.assert_called_once_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_once_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -542,3 +566,80 @@ def test_output_format_removed_from_bedrock_invoke_request(): assert ( "output_format" not in result ), f"output_format should be removed for Bedrock Invoke, got keys: {result.keys()}" + + +def test_bedrock_chat_invoke_forwards_output_config_format_natively(local_model_cost_map): + """Regression: ``output_config.format`` is forwarded verbatim on models Bedrock + enforces structured outputs for, instead of being inlined as prompt text.""" + import json + + config = AmazonAnthropicClaudeConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + "required": ["zebra_count"], + "additionalProperties": False, + }, + } + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_chat_invoke_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not eat ``output_config.format`` before the + native-forwarding router runs (Sonnet 4.5 has no effort flags).""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={"max_tokens": 100, "output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_model_cost_map, monkeypatch): + """``drop_params=True`` on a model without native structured-output support + still reaches the inline-schema fallback instead of losing the schema.""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema = {"type": "object", "properties": {"zebra_count": {"type": "integer"}}} + + result = AmazonAnthropicClaudeConfig().transform_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params={}, + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 8d07d38b1b6..09ebc1a3c95 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -935,7 +935,7 @@ def test_bedrock_messages_strips_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -970,7 +970,7 @@ def test_bedrock_messages_preserves_output_config_for_claude_4_6(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1003,7 +1003,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = cfg.transform_anthropic_messages_request( @@ -1014,11 +1014,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): headers={}, ) - mock_supports_factory.assert_called_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -1038,7 +1034,7 @@ def test_bedrock_messages_forwards_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1054,27 +1050,29 @@ def test_bedrock_messages_forwards_output_config(): def test_bedrock_messages_forwards_output_config_with_output_format(): - """``output_config`` is forwarded; ``output_format`` is converted to inline schema.""" + """Legacy ``output_format`` is forwarded as ``output_config.format`` on models + that support native structured outputs, alongside the effort key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } optional_params = { "max_tokens": 4096, "output_config": {"effort": "low"}, - "output_format": { - "type": "json_schema", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - }, + "output_format": schema_format, } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1085,12 +1083,14 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): headers={}, ) - assert result.get("output_config") == {"effort": "low"} + assert result.get("output_config") == {"effort": "low", "format": schema_format} assert "output_format" not in result + assert "answer" not in json.dumps(result["messages"]) def test_bedrock_messages_converts_output_config_format_to_inline_schema(): - """``output_config.format`` is consumed so Bedrock does not see an unknown nested key.""" + """Without native structured-output support, ``output_config.format`` falls back + to the inline schema so Bedrock does not see an unknown nested key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams @@ -1110,8 +1110,8 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1146,7 +1146,7 @@ def test_bedrock_messages_normalizes_output_config_effort_for_opus( cfg = AmazonAnthropicClaudeMessagesConfig() with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1184,8 +1184,8 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1229,7 +1229,7 @@ def test_bedrock_messages_does_not_mutate_callers_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): cfg.transform_anthropic_messages_request( @@ -1271,7 +1271,7 @@ def test_bedrock_messages_strips_output_config_with_output_format(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -1332,7 +1332,7 @@ def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): litellm.drop_params = True try: with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1375,7 +1375,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1482,7 +1482,7 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -3104,3 +3104,149 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): break await asyncio.sleep(0.01) assert logging_obj.completion_start_time is not None + + +def test_bedrock_messages_forwards_output_config_format_natively(local_model_cost_map): + """Regression: on a model Bedrock enforces structured outputs for (Claude + Sonnet 4.5), ``output_config.format`` must be forwarded verbatim, not + silently rewritten into inline prompt text.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "zebra_count": {"type": "integer"}, + "is_tuesday": {"type": "boolean"}, + }, + "required": ["zebra_count", "is_tuesday"], + "additionalProperties": False, + }, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): + """Bedrock rejects ``output_config.format`` for the Claude 5 family, so the + schema falls back to the inline-text path instead of a deterministic 400.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema = { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(local_model_cost_map): + """When a request carries both schema forms, the legacy top-level + ``output_format`` keeps winning, matching the pre-existing precedence.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + legacy_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"legacy_field": {"type": "string"}}}, + } + newer_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"newer_field": {"type": "string"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_format": legacy_format, + "output_config": {"format": newer_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": legacy_format} + assert "output_format" not in result + assert "newer_field" not in json.dumps(result) + + +def test_bedrock_messages_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not strip a natively forwarded + ``output_config.format`` on models without effort support (Sonnet 4.5).""" + import litellm + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setattr(litellm, "drop_params", True) + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_model_cost_map): + """Sonnet 4.5 has native structured-output support but no effort support, so + a mixed ``output_config`` keeps ``format`` and drops ``effort``.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 4096, + "output_config": {"format": schema_format, "effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 389bf4a8e40..609b5c75801 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -520,3 +520,44 @@ def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_stati assert merged["aws_secret_access_key"] == "caller-secret" assert merged["aws_session_token"] == "caller-token" assert merged["aws_region_name"] == "us-west-2" + + +def test_strip_unsupported_output_config_keeps_format_drops_effort(local_model_cost_map): + """On a model with neither effort flag, only the ``format`` key survives.""" + from litellm.llms.bedrock.common_utils import ( + strip_unsupported_bedrock_invoke_output_config_keys, + ) + + schema_format = {"type": "json_schema", "schema": {"type": "object"}} + body = {"output_config": {"effort": "high", "format": schema_format}} + + strip_unsupported_bedrock_invoke_output_config_keys( + model="anthropic.claude-3-haiku-20240307-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": schema_format} + + +def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_map): + """The legacy ``output_format`` wins over ``output_config.format`` when a + request carries both, matching the pre-existing precedence.""" + from litellm.llms.bedrock.common_utils import ( + apply_bedrock_invoke_structured_output, + ) + + legacy = {"type": "json_schema", "schema": {"type": "object", "properties": {"a": {"type": "string"}}}} + newer = {"type": "json_schema", "schema": {"type": "object", "properties": {"b": {"type": "string"}}}} + body = { + "messages": [{"role": "user", "content": "hi"}], + "output_format": legacy, + "output_config": {"format": newer}, + } + + apply_bedrock_invoke_structured_output( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": legacy} + assert "output_format" not in body From 847d737b8e93dd3abcff819e0537e84550b8c436 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:13:05 -0700 Subject: [PATCH 36/92] test(ui): budget DOM-structure assertions in dashboard tests Turn on testing-library/no-node-access, no-container and prefer-screen-queries as warnings and baseline them in eslint-budgets.json so the counts can only go down. These three rules catch tests that assert on DOM structure rather than on what a user can observe: reaching through parentElement chains, querying the container by CSS selector, and destructuring queries off render instead of going through screen. Those assertions break on refactors that change nothing a user sees, and stay green when the behaviour underneath is broken. Baselines are the current counts, so nothing fails today. --- ui/litellm-dashboard/eslint-budgets.json | 5 ++++- ui/litellm-dashboard/eslint.config.mjs | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index bbf69c4a77a..d986722c3b3 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -4,5 +4,8 @@ "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, - "local/no-long-condition-chain": { "max": 265, "target": 120 } + "local/no-long-condition-chain": { "max": 265, "target": 120 }, + "testing-library/no-container": { "max": 150, "target": 50 }, + "testing-library/no-node-access": { "max": 760, "target": 500 }, + "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } } diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 23cc5096bb1..f5e3b23b3ec 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -104,10 +104,13 @@ const eslintConfig = [ plugins: { "testing-library": testingLibrary, "jest-dom": jestDom }, rules: { "testing-library/await-async-queries": "error", + "testing-library/no-container": "warn", + "testing-library/no-node-access": "warn", "testing-library/no-wait-for-multiple-assertions": "error", "testing-library/no-wait-for-side-effects": "error", "testing-library/prefer-find-by": "error", "testing-library/prefer-presence-queries": "error", + "testing-library/prefer-screen-queries": "warn", "jest-dom/prefer-checked": "error", "jest-dom/prefer-empty": "error", "jest-dom/prefer-enabled-disabled": "error", From 2fbea77afa31f488c442d347ff96da06edf30c9d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:19:27 -0700 Subject: [PATCH 37/92] test(ui): assert DataTable behavior instead of DOM structure The shared DataTable test reached for elements by CSS selector and by walking parentElement chains, then asserted on Tailwind class strings. It had no role queries at all, so a wrapper div anywhere in the render tree broke it while changing nothing a user sees. Columns, rows and headers are now found the way a user finds them: by role and by the text on screen. The compact skeleton row is compared against the loaded row's height rather than a hard-coded h-8, so renaming the class no longer breaks the test but shrinking the row still does. The fillHeight and maxBodyHeight cases stay class assertions. jsdom has no layout engine, so there is nothing behavioural to assert there. What they no longer do is derive their elements from incidental nesting: the three layout wrappers and the header now publish a stable test id, which is also why the resizer's write-only data-resizer attribute became one. Budgets drop with the counts: no-container 150 to 133, no-node-access 760 to 723. --- ui/litellm-dashboard/eslint-budgets.json | 4 +- .../shared/DataTable/DataTable.test.tsx | 114 +++++++++--------- .../components/shared/DataTable/DataTable.tsx | 13 +- 3 files changed, 66 insertions(+), 65 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index d986722c3b3..3bfe56724eb 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -5,7 +5,7 @@ "max-depth": { "max": 70, "target": 30 }, "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, - "testing-library/no-container": { "max": 150, "target": 50 }, - "testing-library/no-node-access": { "max": 760, "target": 500 }, + "testing-library/no-container": { "max": 133, "target": 50 }, + "testing-library/no-node-access": { "max": 723, "target": 500 }, "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } } diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 58a0cd94997..7ead9bb64d4 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -1,5 +1,5 @@ import type { ColumnDef, ExpandedState } from "@tanstack/react-table"; -import { render, screen, waitFor } from "@testing-library/react"; +import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -21,6 +21,12 @@ function person(id: string, name: string, flagged = false): Person { const names = (): (string | null)[] => screen.getAllByTestId("name-cell").map((el) => el.textContent); +const heightClassesOf = (el: HTMLElement | undefined): string[] => + (el?.className ?? "") + .split(/\s+/) + .filter((cls) => cls.startsWith("h-")) + .sort(); + const nameCellColumns: ColumnDef[] = [ { accessorKey: "name", @@ -229,12 +235,10 @@ describe("DataTable sorting", () => { describe("DataTable layout", () => { it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { - const { container } = render(); + render(); - const table = container.querySelector("table"); - expect(table).not.toBeNull(); // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow. - expect(table?.style.minWidth).toBe("100%"); + expect(screen.getByRole("table")).toHaveStyle({ minWidth: "100%" }); }); }); @@ -354,17 +358,21 @@ describe("DataTable loading", () => { const { rerender } = render( , ); - const skeletonRow = screen.getAllByTestId("skeleton-row").at(0); - const loadedRowHeight = "h-8"; - expect(skeletonRow?.className).toContain(loadedRowHeight); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); rerender(); - expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight); + const loadedHeight = heightClassesOf(screen.getByRole("row", { name: /Charlie/ })); + + expect(loadedHeight).not.toEqual([]); + expect(skeletonHeight).toEqual(loadedHeight); }); it("does not force the compact height on default-size skeleton rows", () => { - render(); - expect(screen.getAllByTestId("skeleton-row").at(0)?.className).not.toContain("h-8"); + const { rerender } = render(); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); + + rerender(); + expect(heightClassesOf(screen.getAllByRole("row").at(-1))).not.toEqual(skeletonHeight); }); it("varies skeleton shape and width per column instead of one fixed bar", () => { @@ -420,7 +428,7 @@ describe("DataTable loading", () => { describe("DataTable column visibility", () => { it("hides a column when toggled off in the view-options menu", async () => { const user = userEvent.setup(); - const { container } = render( + render( { />, ); - expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull(); + expect(screen.getByRole("columnheader", { name: "Email" })).toBeInTheDocument(); await user.click(screen.getByTestId("view-options-trigger")); await user.click(await screen.findByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).toBeNull()); + await waitFor(() => expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument()); await user.click(screen.getByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull()); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); }); it("omits columns that opt out of hiding from the menu", async () => { @@ -468,14 +476,10 @@ describe("DataTable column visibility", () => { describe("DataTable pinned columns", () => { it("applies sticky positioning to a pinned column only", () => { - const { container } = render(); + render(); - const pinnedHead = container.querySelector('th[data-header-id="name"]'); - const normalHead = container.querySelector('th[data-header-id="email"]'); - - expect(pinnedHead?.style.position).toBe("sticky"); - expect(pinnedHead?.style.left).toBe("0px"); - expect(normalHead?.style.position).toBe(""); + expect(screen.getByRole("columnheader", { name: "Name" })).toHaveStyle({ position: "sticky", left: "0px" }); + expect(screen.getByRole("columnheader", { name: "Email" })).not.toHaveStyle({ position: "sticky" }); }); }); @@ -570,7 +574,7 @@ describe("DataTable expansion", () => { describe("DataTable row styling and footer", () => { it("applies rowClassName to the matching row only", () => { const data = [person("a", "Alice", true), person("b", "Bob", false)]; - const { container } = render( + render( { />, ); - expect(container.querySelector('tr[data-row-id="a"]')?.className).toContain("flagged-row"); - expect(container.querySelector('tr[data-row-id="b"]')?.className).not.toContain("flagged-row"); + expect(screen.getByRole("row", { name: /Alice/ })).toHaveClass("flagged-row"); + expect(screen.getByRole("row", { name: /Bob/ })).not.toHaveClass("flagged-row"); }); it("renders the footer slot inside a tfoot element", () => { @@ -596,63 +600,57 @@ describe("DataTable row styling and footer", () => { />, ); - expect(screen.getByTestId("footer-row").closest("tfoot")).not.toBeNull(); + const rowGroups = screen.getAllByRole("rowgroup"); + expect(within(rowGroups.at(-1) as HTMLElement).getByText("Total: 3")).toBeInTheDocument(); }); }); describe("DataTable layout", () => { it("exposes resize handles with stable selectors only when resizing is enabled", () => { - const { container, rerender } = render( - , - ); - expect(container.querySelectorAll("[data-resizer][data-header-id]").length).toBe(2); + const { rerender } = render(); + expect(screen.getByTestId("column-resizer-name")).toBeInTheDocument(); + expect(screen.getByTestId("column-resizer-email")).toBeInTheDocument(); rerender(); - expect(container.querySelectorAll("[data-resizer]").length).toBe(0); + expect(screen.queryByTestId("column-resizer-name")).not.toBeInTheDocument(); }); it("makes the header sticky and constrains body height when maxBodyHeight is set", () => { - const { container } = render(); - expect(container.querySelector("thead")?.className).toContain("sticky"); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - expect(scroller).toHaveStyle({ maxHeight: "240px" }); + render(); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky"); + expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" }); }); it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - const frame = scroller.parentElement as HTMLElement; - const outer = frame.parentElement as HTMLElement; + render(); + const outer = screen.getByTestId("data-table-root"); + const frame = screen.getByTestId("data-table-frame"); + const scroller = screen.getByTestId("data-table-scroller"); // A ceiling, not a stretch: flex-1 here would hold the footer at the bottom on a two-row table. - expect(outer.className).toContain("max-h-full"); - expect(outer.className).not.toContain("flex-1"); - expect(frame.className).not.toContain("flex-1"); - expect(scroller.className).not.toContain("flex-1"); + expect(outer).toHaveClass("max-h-full", "flex-col"); + expect(outer).not.toHaveClass("flex-1"); + expect(frame).toHaveClass("flex-col"); + expect(frame).not.toHaveClass("flex-1"); + expect(scroller).not.toHaveClass("flex-1"); - expect(outer.className).toContain("flex-col"); - expect(frame.className).toContain("flex-col"); - expect(scroller.className).toContain("min-h-0"); - expect(scroller.className).toContain("overflow-auto"); + expect(scroller).toHaveClass("min-h-0", "overflow-auto"); expect(scroller).toHaveStyle({ maxHeight: "" }); // Without this the Table primitive's own overflow container captures the sticky header. - expect(scroller.className).toContain("[&_[data-slot=table-container]]:overflow-visible"); + expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible"); - const thead = container.querySelector("thead") as HTMLElement; - expect(thead.className).toContain("sticky"); // Rows pass under the header, so the semi-transparent row tint alone would let them show through. - expect(thead.className).toContain("bg-background"); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background"); }); it("leaves the default layout untouched when neither height mode is set", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; + render(); + const scroller = screen.getByTestId("data-table-scroller"); - expect(scroller.className).toContain("overflow-x-auto"); - expect(scroller.className).not.toContain("min-h-0"); + expect(scroller).toHaveClass("overflow-x-auto"); + expect(scroller).not.toHaveClass("min-h-0"); expect(scroller).toHaveStyle({ maxHeight: "" }); - expect((scroller.parentElement as HTMLElement).className).not.toContain("flex-col"); - expect(container.querySelector("thead")?.className).not.toContain("sticky"); - expect(container.querySelector("thead")?.className).not.toContain("bg-background"); + expect(screen.getByTestId("data-table-frame")).not.toHaveClass("flex-col"); + expect(screen.getByTestId("data-table-head")).not.toHaveClass("sticky", "bg-background"); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index e8357b37715..60267606951 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -195,8 +195,7 @@ function DataTableHeadCell({ header, size, stickyHeader, enableColumnResi )} {canResize && (
column.resetSize()} @@ -589,15 +588,19 @@ export function DataTable(props: DataTableProps -
+
+
{toolbar !== undefined &&
{toolbar(table)}
}
- + {table.getHeaderGroups().map((headerGroup) => ( {headerGroup.headers.map((header) => ( From c5ba2b5fcf6f698b0e0f06a190ebd9d3b553c51b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 00:26:49 -0700 Subject: [PATCH 38/92] test(ui): query the screen instead of the render result Two changes, both about finding elements the way a user finds them. Twenty-six test files destructured queries off render and called them bare. Those queries are scoped to the render container, so they quietly miss anything portalled into the body, and they read as if they were free functions. They now go through screen. ChatMessageBubble and the key info panel derived elements by walking closest/parentElement/firstElementChild and then asserted on the classes they found. A wrapper element anywhere in between broke them. The bubble surface, the avatar and the budget reset value now publish a test id, so the assertions survive markup changes and still fail when the styling they check actually regresses. Budgets drop with the counts: prefer-screen-queries 221 to 21, no-node-access 723 to 716. The 21 remaining prefer-screen-queries are not all fixable: 18 of them are within(dialog) results in MCPToolsetsTab, which the rule cannot tell apart from a render result. Target is 18, not 0. --- ui/litellm-dashboard/eslint-budgets.json | 4 +- .../_components/APIReferenceView.test.tsx | 12 +- .../cache_settings/RedisTypeSelector.test.tsx | 8 +- .../_components/CacheLeakageCard.test.tsx | 50 ++++---- .../CostOptimizationView.activity.test.tsx | 16 +-- .../_components/CostOptimizationView.test.tsx | 50 ++++---- .../_components/PromptCachingTab.test.tsx | 8 +- .../_components/UsageTab.test.tsx | 118 +++++++++--------- .../_components/guardrail_info.test.tsx | 78 +++++------- .../_components/pii_components.test.tsx | 18 ++- .../_components/pii_configuration.test.tsx | 6 +- .../_components/mcp_servers.test.tsx | 28 ++--- .../PriceDataManagementTab.test.tsx | 6 +- .../models-and-endpoints/page.test.tsx | 50 ++++---- .../chat_ui/ChatMessageBubble.test.tsx | 8 +- .../components/chat_ui/ChatMessageBubble.tsx | 2 + .../components/compareUI/CompareUI.test.tsx | 30 ++--- .../components/MessageDisplay.test.tsx | 30 ++--- .../EntityUsageExportModal.test.tsx | 13 +- .../add_model/advanced_settings.test.tsx | 34 ++--- .../add_model/litellm_model_name.test.tsx | 12 +- .../bulk_create_users_button.test.tsx | 4 +- ...cost_optimization_feedback_banner.test.tsx | 18 +-- .../organization/organization_view.test.tsx | 4 +- .../src/components/settings.test.tsx | 20 +-- .../shared/PaginationStatusAlerts.test.tsx | 16 +-- .../shared/charts/area_chart.test.tsx | 6 +- .../shared/charts/bar_chart.test.tsx | 4 +- .../KeyInfoView.handleKeyUpdate.test.tsx | 4 +- .../key_info_view.budget_display.test.tsx | 2 +- .../components/templates/key_info_view.tsx | 2 +- 31 files changed, 323 insertions(+), 338 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 3bfe56724eb..e7545eb2383 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -6,6 +6,6 @@ "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, - "testing-library/no-node-access": { "max": 723, "target": 500 }, - "testing-library/prefer-screen-queries": { "max": 221, "target": 0 } + "testing-library/no-node-access": { "max": 716, "target": 500 }, + "testing-library/prefer-screen-queries": { "max": 21, "target": 18 } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx index b2e0a42eecb..dae19032e20 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx @@ -13,17 +13,17 @@ describe("APIReferenceView", () => { it("uses the API doc base url when provided", () => { const apiDocUrl = "https://docs.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(apiDocUrl)); }); it("falls back to the proxy base url when the docs url is missing", () => { const proxyUrl = "https://proxy.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(proxyUrl)); }); @@ -31,7 +31,7 @@ describe("APIReferenceView", () => { const apiDocUrl = "https://docs-preferred.litellm.test"; const proxyUrl = "https://proxy-backup.litellm.test"; - const { getAllByTestId } = render( + render( { />, ); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); const renderedCode = codeBlocks[0].textContent ?? ""; expect(renderedCode).toContain(apiDocUrl); expect(renderedCode).not.toContain(proxyUrl); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx index 9d4d5a6d425..372f27e2be1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx @@ -1,12 +1,10 @@ import { describe, expect, it } from "vitest"; import RedisTypeSelector from "./RedisTypeSelector"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; describe("RedisTypeSelector", () => { it("should render the component", () => { - const { getAllByText } = render( - {}} />, - ); - expect(getAllByText(/Redis/i).length).toBeGreaterThan(0); + render( {}} />); + expect(screen.getAllByText(/Redis/i).length).toBeGreaterThan(0); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index 8c36b934789..f320d8e0f97 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { DailyData, KeyMetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; @@ -79,25 +79,25 @@ const renderWith = (results: DailyData[], overrides: Partial describe("CacheLeakageCard", () => { it("ranks leaking keys by uncached prompt tokens and shows cache hit ratio", () => { - const { getByText, getByLabelText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-caching": key("caching-key", { prompt_tokens: 1000, cache_read_input_tokens: 900 }), "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }), ]); - expect(getByText("leaky-key")).toBeInTheDocument(); - expect(getByText("0.0%")).toBeInTheDocument(); - expect(getByText("90.0%")).toBeInTheDocument(); + expect(screen.getByText("leaky-key")).toBeInTheDocument(); + expect(screen.getByText("0.0%")).toBeInTheDocument(); + expect(screen.getByText("90.0%")).toBeInTheDocument(); [ "Input tokens you sent in this range that weren't served from or written to the cache", "Share of your input tokens that were served from the cache", "About how much you'd save if this uncached input used prompt caching. Estimated as uncached input tokens times what your cached traffic already nets per cached token (realized cache savings, after write premiums, ÷ cache read and write tokens). Blank when caching is not currently saving anything overall.", - ].forEach((info) => expect(getByLabelText(info)).toBeInTheDocument()); + ].forEach((info) => expect(screen.getByLabelText(info)).toBeInTheDocument()); }); it("sorts by the clicked column, worst cache hit rate first", () => { - const { getAllByRole, getByText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-a": key("alpha", { prompt_tokens: 10000, @@ -111,48 +111,48 @@ describe("CacheLeakageCard", () => { }), }), ]); - const firstDataRow = () => getAllByRole("row")[1]; + const firstDataRow = () => screen.getAllByRole("row")[1]; expect(firstDataRow()).toHaveTextContent("alpha"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("bravo"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("alpha"); }); it("switches to the model view and lists only Anthropic models", () => { - const { getByText, queryByText } = renderWith([ + renderWith([ dayWithModels("2026-07-12", { "claude-sonnet-5": { prompt_tokens: 5000, cache_read_input_tokens: 0 }, "gpt-4o": { prompt_tokens: 8000, cache_read_input_tokens: 0 }, }), ]); - fireEvent.click(getByText("By model")); + fireEvent.click(screen.getByText("By model")); - expect(getByText("Cache leakage by model")).toBeInTheDocument(); - expect(getByText("claude-sonnet-5")).toBeInTheDocument(); - expect(queryByText("gpt-4o")).not.toBeInTheDocument(); + expect(screen.getByText("Cache leakage by model")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet-5")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o")).not.toBeInTheDocument(); }); it("shows an empty state when no key used tokens in the range", () => { - const { getByText, queryByRole } = renderWith([dayWithKeys("2026-07-12", {})]); + renderWith([dayWithKeys("2026-07-12", {})]); - expect(getByText("No key usage in this range.")).toBeInTheDocument(); - expect(queryByRole("table")).not.toBeInTheDocument(); + expect(screen.getByText("No key usage in this range.")).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); }); it("tells the user the table is still filling in while fallback pages stream", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { getByText, getByRole } = renderWith([day], { isFetchingMore: true }); + renderWith([day], { isFetchingMore: true }); - expect(getByRole("table")).toBeInTheDocument(); + expect(screen.getByRole("table")).toBeInTheDocument(); expect( - getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).toBeInTheDocument(); }); @@ -160,10 +160,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day], { loading: true }); + renderWith([day], { loading: true }); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); @@ -171,10 +171,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day]); + renderWith([day]); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 1cc7bec13d1..03250e3e53b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, waitFor } from "@testing-library/react"; +import { fireEvent, render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -56,7 +56,7 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { getByRole, getByTestId, findByTestId, queryByText } = render( + render( , @@ -64,12 +64,12 @@ describe("CostOptimizationView daily activity", () => { await waitFor(() => expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1)); - fireEvent.click(getByRole("tab", { name: "Prompt Caching" })); - await findByTestId("caching-settings"); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Caching" })); + await screen.findByTestId("caching-settings"); expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1); expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); - expect(queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); }); it("shows the fetch-progress banner while the paginated fallback streams pages in", async () => { @@ -84,13 +84,13 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { findByText, getByRole } = render( + render( , ); - expect(await findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); - expect(getByRole("button", { name: "Stop" })).toBeInTheDocument(); + expect(await screen.findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Stop" })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index d5df5aa75da..028367555a1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -45,32 +45,32 @@ describe("CostOptimizationView", () => { }); it("renders the standard page header with the sidebar's Cost Optimization icon", () => { - const { container, getByRole, getByText } = renderView(); + const { container } = renderView(); - expect(getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); - expect(getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); + expect(screen.getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); + expect(screen.getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); expect(container.querySelector(".lucide-piggy-bank")).not.toBeNull(); }); it("renders the four cost-optimization tabs", () => { - const { getByText } = renderView(); + renderView(); - expect(getByText("Overall")).toBeInTheDocument(); - expect(getByText("Prompt Compression")).toBeInTheDocument(); - expect(getByText("Prompt Caching")).toBeInTheDocument(); - expect(getByText("Auto-Router")).toBeInTheDocument(); + expect(screen.getByText("Overall")).toBeInTheDocument(); + expect(screen.getByText("Prompt Compression")).toBeInTheDocument(); + expect(screen.getByText("Prompt Caching")).toBeInTheDocument(); + expect(screen.getByText("Auto-Router")).toBeInTheDocument(); }); it("defaults to the Overall tab and switches the active tab on click", () => { - const { getByRole } = renderView(); + renderView(); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); - fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Compression" })); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); // Unlike the other three pages in this cleanup, Cost Optimization keeps its @@ -80,21 +80,21 @@ describe("CostOptimizationView", () => { // are proxy-admin-only, so those are what disappear. describe("proxy-admin-only tabs", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])("shows %s the Overall tab only", (userRole) => { - const { getByRole, queryByRole } = renderView(userRole); + renderView(userRole); - expect(getByRole("tab", { name: "Overall" })).toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Overall" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); }); it("never mounts the panels behind the admin-only endpoints for an internal user", () => { - const { getByTestId, queryByTestId } = renderView("Internal User"); + renderView("Internal User"); - expect(getByTestId("usage-tab")).toBeInTheDocument(); - expect(queryByTestId("compression-tab")).not.toBeInTheDocument(); - expect(queryByTestId("caching-tab")).not.toBeInTheDocument(); - expect(queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); + expect(screen.getByTestId("usage-tab")).toBeInTheDocument(); + expect(screen.queryByTestId("compression-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("caching-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx index 38517dab0ab..2c602033171 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx @@ -1,4 +1,4 @@ -import { render, waitFor } from "@testing-library/react"; +import { render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; const mockGetGeneralSettingsCall = vi.fn(); @@ -37,10 +37,10 @@ describe("PromptCachingTab", () => { cancelled: false, cancel: vi.fn(), }; - const { getByTestId } = render(); + render(); - expect(getByTestId("caching-settings")).toBeInTheDocument(); - expect(getByTestId("cache-leakage-card")).toBeInTheDocument(); + expect(screen.getByTestId("caching-settings")).toBeInTheDocument(); + expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument(); await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity }))); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index df23e5509bf..f85a667a074 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ToolSpendResponse } from "@/components/networking"; @@ -152,13 +152,13 @@ describe("UsageTab", () => { gateway_injected_caching_savings_spend: 0.006, compression_saved_tokens: 100000, }; - const { getByText } = renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); + renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); - expect(getByText("$0.1500")).toBeInTheDocument(); - expect(getByText("$0.1400")).toBeInTheDocument(); - expect(getByText("$0.0100")).toBeInTheDocument(); - expect(getByText("$0.0160")).toBeInTheDocument(); - expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + expect(screen.getByText("$0.1500")).toBeInTheDocument(); + expect(screen.getByText("$0.1400")).toBeInTheDocument(); + expect(screen.getByText("$0.0100")).toBeInTheDocument(); + expect(screen.getByText("$0.0160")).toBeInTheDocument(); + expect(screen.getByText("140,000 tokens compressed")).toBeInTheDocument(); }); const twoDays = () => [ @@ -167,11 +167,11 @@ describe("UsageTab", () => { ]; it("opens on a running total anchored at $0 at the start of the range", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative prepends a synthetic $0 point at the range start (Jul 1) so the // line rises from zero rather than floating; the daily running totals follow. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(3); expect(series[0]).toMatchObject({ date: "Jul 1", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); @@ -183,12 +183,12 @@ describe("UsageTab", () => { // The original complaint: a one-day range plotted a single floating dot. The // synthetic start anchor gives the line a zero origin to climb from. const oneDay = new Date(2026, 6, 24); - const { getByTestId } = renderWith( - [day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], - { from: oneDay, to: oneDay }, - ); + renderWith([day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], { + from: oneDay, + to: oneDay, + }); - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ date: "Jul 24", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ date: "Jul 24", Compression: 0.2, "Prompt caching": 0.05 }); @@ -202,49 +202,49 @@ describe("UsageTab", () => { day("2026-07-13", { gateway_injected_caching_savings_spend: 0.1 }), day("2026-07-12", { gateway_injected_caching_savings_spend: 0.04 }), ]; - const { getByTestId, getByRole } = renderWith(newestFirst); + renderWith(newestFirst); // The $0 anchor leads, then the days climb oldest to newest. - const cumulative = readSeries(getByTestId("area-chart")); + const cumulative = readSeries(screen.getByTestId("area-chart")); expect(cumulative.map((p: { date: string }) => p.date)).toEqual(["Jul 1", "Jul 12", "Jul 13"]); expect(cumulative[1]["Prompt caching"]).toBeCloseTo(0.04, 5); expect(cumulative[2]["Prompt caching"]).toBeCloseTo(0.14, 5); expect(cumulative[2]["Prompt caching"]).toBeGreaterThan(cumulative[1]["Prompt caching"]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const perDay = readSeries(getByTestId("bar-chart")); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const perDay = readSeries(screen.getByTestId("bar-chart")); expect(perDay.map((p: { date: string }) => p.date)).toEqual(["Jul 12", "Jul 13"]); }); it("draws bars of the raw per-interval readings on the other tab", async () => { - const { getByRole, getByTestId, queryByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative opens on the area line. - expect(getByTestId("area-chart")).toBeInTheDocument(); + expect(screen.getByTestId("area-chart")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); // Per day switches to a bar chart of the unaccumulated daily savings, with no // synthetic anchor prepended. - expect(queryByTestId("area-chart")).not.toBeInTheDocument(); - const series = readSeries(getByTestId("bar-chart")); + expect(screen.queryByTestId("area-chart")).not.toBeInTheDocument(); + const series = readSeries(screen.getByTestId("bar-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); }); it("says what the line means and over what range", async () => { - const { getByText, getByRole } = renderWith(twoDays()); + renderWith(twoDays()); - expect(getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); - expect(getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + expect(screen.getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + expect(screen.getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); }); it("builds the per-driver donut from the range totals, not the running total", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -252,9 +252,9 @@ describe("UsageTab", () => { }); it("omits a driver slice when that driver has no savings", () => { - const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([{ driver: "Compression", color: "emerald", usd: expect.closeTo(0.04, 5) }]); }); @@ -262,7 +262,7 @@ describe("UsageTab", () => { // Stacking sums the series into one bar. Auto-router savings go negative when a // model switch pays for a cold cache, and that segment would be drawn below the // axis while the rest of the bar still read as the day's total. - const { getByRole, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -270,8 +270,8 @@ describe("UsageTab", () => { }), ]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const bars = getByTestId("bar-chart"); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const bars = screen.getByTestId("bar-chart"); expect(bars).toHaveAttribute("data-stack", "false"); expect(readSeries(bars)[0]).toMatchObject({ "Auto-router": -0.05 }); }); @@ -281,10 +281,10 @@ describe("UsageTab", () => { // per day"). Hand-rolled rows made it compete with the legend and the toggle for // width, so the header grew a line on one tab and the chart moved with it. CardHeader // sizes the action column to its content and gives the rest to the title column. - const { getByRole, getByTestId, container } = renderWith(twoDays()); + const { container } = renderWith(twoDays()); const header = () => { - const legend = getByTestId("chart-legend"); + const legend = screen.getByTestId("chart-legend"); const action = legend.closest('[data-slot="card-action"]') as HTMLElement; const cardHeader = action.parentElement as HTMLElement; const description = cardHeader.querySelector('[data-slot="card-description"]') as HTMLElement; @@ -295,12 +295,12 @@ describe("UsageTab", () => { expect(before.action).toBeTruthy(); expect(before.description).toBeTruthy(); // the toggle rides in the same action slot as the legend, so neither moves alone - expect(before.action.contains(getByRole("tablist"))).toBe(true); + expect(before.action.contains(screen.getByRole("tablist"))).toBe(true); // the subtitle lives outside that slot, so its length cannot reposition the controls expect(before.action.contains(before.description)).toBe(false); expect(before.description).toHaveTextContent(/Running total saved/); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); const after = header(); expect(after.action).toBe(before.action); @@ -314,7 +314,7 @@ describe("UsageTab", () => { // Switching models leaves the new one with a cold cache, so a route can cost more // than the baseline would have. A negative slice is meaningless in a donut, but the // total has to keep the loss or the page can only ever report good news. - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -322,16 +322,16 @@ describe("UsageTab", () => { }), ]); - expect(getByText("$0.0700")).toBeInTheDocument(); - expect(getByText("-$0.0500")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("-$0.0500")).toBeInTheDocument(); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices.map((d: { driver: string }) => d.driver)).toEqual(["Compression", "Prompt caching"]); - expect(getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); + expect(screen.getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); }); it("carries auto-router savings into the summary card, donut slice, and cumulative series", () => { - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.04, gateway_injected_caching_savings_spend: 0.006, @@ -345,11 +345,11 @@ describe("UsageTab", () => { ]); // Total saved now sums three drivers, and the auto-router card carries its own total. - expect(getByText("$0.2260")).toBeInTheDocument(); - expect(getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("$0.2260")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); // The driver donut gains a third slice priced from the range totals. - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -357,7 +357,7 @@ describe("UsageTab", () => { ]); // And the cumulative line accumulates the auto-router series alongside the others. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series[2]["Auto-router"]).toBeCloseTo(0.07, 5); }); @@ -371,9 +371,9 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); // The 64px bar cap is this card's opt-in; the shared BarChart must not cap @@ -391,14 +391,16 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const [totalByTool, dailyByTool] = bars.slice(-2); expect(dailyByTool).toHaveAttribute("data-show-legend", "false"); expect(totalByTool).toHaveAttribute("data-colors", dailyByTool.getAttribute("data-colors")); - const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + const toolLegends = screen + .getAllByTestId("chart-legend") + .filter((legend) => legend.textContent === "search,read_file"); expect(toolLegends).toHaveLength(1); }); @@ -415,23 +417,23 @@ describe("UsageTab", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])( "hides the card and never calls the endpoint for %s", async (userRole) => { - const { queryByText, getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend, userRole, }); // Liveness gate: the daily-activity charts still render for this role, // so the absence below is the gate, not an empty tab. - expect(getByTestId("donut-chart")).toBeInTheDocument(); - expect(queryByText("Spend by tool")).not.toBeInTheDocument(); + expect(screen.getByTestId("donut-chart")).toBeInTheDocument(); + expect(screen.queryByText("Spend by tool")).not.toBeInTheDocument(); await vi.waitFor(() => expect(mockGetToolSpend).not.toHaveBeenCalled()); }, ); it("keeps the card and the endpoint call for an admin", async () => { - const { findByText } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); - expect(await findByText("Spend by tool")).toBeInTheDocument(); + expect(await screen.findByText("Spend by tool")).toBeInTheDocument(); expect(mockGetToolSpend).toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 2b90a1d8cbc..fcffc2122e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -1,5 +1,5 @@ import * as networking from "@/components/networking"; -import { fireEvent, render, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, waitFor, within, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import GuardrailInfoView from "./guardrail_info"; @@ -65,21 +65,19 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getAllByText, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); // Wait for the loading to complete and data to be rendered await waitFor(() => { // The guardrail name appears in multiple places (title and settings tab) - const elements = getAllByText("Test Guardrail"); + const elements = screen.getAllByText("Test Guardrail"); expect(elements.length).toBeGreaterThan(0); }); // Verify other key elements are present - expect(getByText("Back to Guardrails")).toBeInTheDocument(); - expect(getByText("Overview")).toBeInTheDocument(); - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Back to Guardrails")).toBeInTheDocument(); + expect(screen.getByText("Overview")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); it("should render a tag-based mode object rather than crashing the detail view", async () => { @@ -105,11 +103,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findAllByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); it("should render the provider logo from the bundled guardrail logo map", async () => { @@ -135,11 +131,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByAltText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - const logo = await findByAltText("Presidio PII logo"); + const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); }); @@ -167,25 +161,27 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText, findByText, container } = render( + const { container } = render( {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Click the Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); // Wait for the Settings panel to render await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); await userEvent.hover(within(container).getByRole("img", { name: "Config guardrail details" })); - expect(await findByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); + expect( + await screen.findByText("Guardrail is defined in the config file and cannot be edited."), + ).toBeInTheDocument(); }); it("should render the guardrail info", async () => { @@ -216,12 +212,10 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("PII Entity Configuration")).toBeInTheDocument(); + expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); }); }); it("should handle content filter updates correctly", async () => { @@ -251,30 +245,28 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - const { getByText, getByLabelText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Go to Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); // Enter Edit Mode - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Modify Guardrail Name to force an update - const nameInput = getByLabelText("Guardrail Name"); + const nameInput = screen.getByLabelText("Guardrail Name"); fireEvent.change(nameInput, { target: { value: "Updated Name" } }); // Save with only name change - const saveButton = getByText("Save Changes"); + const saveButton = screen.getByText("Save Changes"); fireEvent.click(saveButton); await waitFor(() => { @@ -300,16 +292,16 @@ describe("Guardrail Info", () => { // Enter Edit Mode again to make changes await waitFor(() => { - expect(getByText("Edit Settings")).toBeInTheDocument(); + expect(screen.getByText("Edit Settings")).toBeInTheDocument(); }); - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Now modify the values using the mock button - const simulateChangeButton = getByText("Simulate Change"); + const simulateChangeButton = screen.getByText("Simulate Change"); fireEvent.click(simulateChangeButton); // Save again - fireEvent.click(getByText("Save Changes")); + fireEvent.click(screen.getByText("Save Changes")); await waitFor(() => { expect(networking.updateGuardrailCall).toHaveBeenCalled(); @@ -339,12 +331,10 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByRole, getByRole, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx index 2f839487111..30ff2cf7c5f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import { CategoryFilter, QuickActions, PiiEntityList } from "./pii_components"; import type { PiiEntityCategory } from "@/components/guardrails/types"; @@ -6,25 +6,21 @@ import type { PiiEntityCategory } from "@/components/guardrails/types"; describe("CategoryFilter", () => { it("should render", () => { const emptyCategories: PiiEntityCategory[] = []; - const { getByText } = render( - {}} />, - ); - expect(getByText("Filter by category")).toBeInTheDocument(); + render( {}} />); + expect(screen.getByText("Filter by category")).toBeInTheDocument(); }); }); describe("QuickActions", () => { it("should render", () => { - const { getByText } = render( - {}} onUnselectAll={() => {}} hasSelectedEntities={false} />, - ); - expect(getByText("Quick Actions")).toBeInTheDocument(); + render( {}} onUnselectAll={() => {}} hasSelectedEntities={false} />); + expect(screen.getByText("Quick Actions")).toBeInTheDocument(); }); }); describe("PiiEntityList", () => { it("should render", () => { - const { getByText } = render( + render( { entityToCategoryMap={new Map()} />, ); - expect(getByText("No PII types match your filter criteria")).toBeInTheDocument(); + expect(screen.getByText("No PII types match your filter criteria")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx index 00c568ef35b..4f822578fe8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx @@ -1,10 +1,10 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import PiiConfiguration from "./pii_configuration"; describe("PiiConfiguration", () => { it("should render", () => { - const { getByText } = render( + render( { entityCategories={[]} />, ); - expect(getByText("Configure PII Protection")).toBeInTheDocument(); + expect(screen.getByText("Configure PII Protection")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 8dca87e24ef..8abb8855e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -45,7 +45,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -53,11 +53,11 @@ describe("MCPServers", () => { // Wait for the component to load and check if title renders await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the title is rendered - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); it("should render mocked MCP servers data in the table", async () => { @@ -96,7 +96,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); const queryClient = createQueryClient(); - const { getByText, getAllByText } = render( + render( , @@ -104,19 +104,19 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Wait for the mocked data to render in the table await waitFor(() => { - expect(getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); }); // Verify the mocked server data is rendered in the table - expect(getByText("Test Server 1")).toBeInTheDocument(); - expect(getByText("Test Server 2")).toBeInTheDocument(); - expect(getAllByText("test-server-1").length).toBeGreaterThan(0); - expect(getAllByText("test-server-2").length).toBeGreaterThan(0); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 2")).toBeInTheDocument(); + expect(screen.getAllByText("test-server-1").length).toBeGreaterThan(0); + expect(screen.getAllByText("test-server-2").length).toBeGreaterThan(0); // Verify the API was called // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock @@ -168,7 +168,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -176,7 +176,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the health check API was called (without a server ID filter — the hook always @@ -211,7 +211,7 @@ describe("MCPServers", () => { ); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -219,7 +219,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify that health check was initiated diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index 8b34d61ebad..7cd418bc176 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -1,5 +1,5 @@ /* @vitest-environment jsdom */ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PriceDataManagementTab from "./PriceDataManagementTab"; @@ -11,7 +11,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ describe("PriceDataManagementTab", () => { it("renders its content standalone, without a tab-panel ancestor", () => { - const { getByText } = render(); - expect(getByText("Price Data Management")).toBeInTheDocument(); + render(); + expect(screen.getByText("Price Data Management")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 521f89a39f2..a504d75bb63 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsPage from "./page"; @@ -64,48 +64,48 @@ describe("ModelsAndEndpointsPage", () => { }); it("renders the admin tab bar and the All Models panel by default", () => { - const { getByRole, getByTestId } = renderPage(); - expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); - expect(getByTestId("panel-all-models")).toBeInTheDocument(); + renderPage(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); + expect(screen.getByTestId("panel-all-models")).toBeInTheDocument(); }); it("switches tabs in-memory, mounting only the active panel", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId, queryByTestId } = renderPage(); - await user.click(getByRole("tab", { name: "Health Status" })); - expect(getByTestId("panel-health")).toBeInTheDocument(); - expect(queryByTestId("panel-all-models")).not.toBeInTheDocument(); + renderPage(); + await user.click(screen.getByRole("tab", { name: "Health Status" })); + expect(screen.getByTestId("panel-health")).toBeInTheDocument(); + expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); }); it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; - const { getByTestId, queryByRole } = renderPage(); - expect(getByTestId("model-info")).toHaveTextContent("model:abc-123"); - expect(queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.getByTestId("model-info")).toHaveTextContent("model:abc-123"); + expect(screen.queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); }); it("renders the team detail overlay from the ?team drill-in", () => { detailState.teamId = "team-9"; - const { getByTestId } = renderPage(); - expect(getByTestId("team-info")).toHaveTextContent("team:team-9"); + renderPage(); + expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); }); it("hides admin-only tabs for a non-admin user", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); - expect(queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); }); // Auto-routers are excluded from the All Models table, so this tab is their home: the only // place in the product to list, create, edit or delete one. describe("Auto-Routers tab", () => { it("sits third, after All Models and Add Model", () => { - const { getAllByRole } = renderPage(); + renderPage(); - const tabs = getAllByRole("tab").map((tab) => tab.textContent); + const tabs = screen.getAllByRole("tab").map((tab) => tab.textContent); expect(tabs[0]).toContain("All Models"); expect(tabs[1]).toBe("Add Model"); expect(tabs[2]).toContain("Auto-Routers"); @@ -115,17 +115,17 @@ describe("ModelsAndEndpointsPage", () => { it("renders its panel when selected", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId } = renderPage(); + renderPage(); - await user.click(getByRole("tab", { name: /Auto-Routers/ })); - expect(getByTestId("panel-auto-routers")).toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: /Auto-Routers/ })); + expect(screen.getByTestId("panel-auto-routers")).toBeInTheDocument(); }); it("is hidden from non-admins, who cannot write models", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); + renderPage(); - expect(queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index a83c11d1444..ed02e7e16cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -93,13 +93,13 @@ describe("ChatMessageBubble", () => { ])("should paint the $role surface from theme tokens, not fixed colours", ({ role, bubble, avatar }) => { render(); - const header = screen.getByText(role).closest("div") as HTMLElement; - const surface = header.parentElement as HTMLElement; + const surface = screen.getByTestId("message-surface"); + const avatarEl = screen.getByTestId("message-avatar"); expect(surface).toHaveClass(...bubble); expect(surface).not.toHaveAttribute("style"); - expect(header.firstElementChild).toHaveClass(avatar); - expect(header.firstElementChild).not.toHaveAttribute("style"); + expect(avatarEl).toHaveClass(avatar); + expect(avatarEl).not.toHaveAttribute("style"); }); it("should show model badge for assistant messages when model is provided", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index 8c54d9e89fa..bd6a4bc49a4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -46,6 +46,7 @@ function ChatMessageBubble({ return (
{ describe("CompareUI", () => { it("should render", () => { - const { getByTestId } = render(); - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); - expect(getByTestId("message-input")).toBeInTheDocument(); + render(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("message-input")).toBeInTheDocument(); }); it("adds a comparison when Add Comparison button is clicked", async () => { const user = userEvent.setup(); - const { container, getByTestId } = render( - , - ); + const { container } = render(); // Verify initial state: 2 comparison panels - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); let comparisonPanels = container.querySelectorAll('[data-testid^="comparison-panel-"]'); expect(comparisonPanels).toHaveLength(2); @@ -117,15 +115,13 @@ describe("CompareUI", () => { }); // Verify the original 2 panels are still there - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); }); it("should handle image upload and send message with attachment", async () => { const user = userEvent.setup(); - const { getByTestId, queryByTestId } = render( - , - ); + render(); const file = new File(["test content"], "test-image.png", { type: "image/png" }); @@ -138,13 +134,13 @@ describe("CompareUI", () => { } await waitFor(() => { - expect(getByTestId("has-attachment")).toBeInTheDocument(); + expect(screen.getByTestId("has-attachment")).toBeInTheDocument(); }); - const textarea = getByTestId("message-textarea"); + const textarea = screen.getByTestId("message-textarea"); fireEvent.change(textarea, { target: { value: "Describe this image" } }); - const sendButton = getByTestId("send-button"); + const sendButton = screen.getByTestId("send-button"); expect(sendButton).toBeEnabled(); await user.click(sendButton); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx index 72a1d41f9fe..69b50d4088f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { MessageType } from "@/components/chat_ui/types"; import { MessageDisplay } from "./MessageDisplay"; @@ -39,9 +39,9 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByText } = render(); - expect(getByText("Hello")).toBeInTheDocument(); - expect(getByText("Hi there!")).toBeInTheDocument(); + render(); + expect(screen.getByText("Hello")).toBeInTheDocument(); + expect(screen.getByText("Hi there!")).toBeInTheDocument(); }); it("displays user and assistant messages with proper grouping and shows loading state", () => { @@ -64,13 +64,13 @@ describe("MessageDisplay", () => { }, }, ]; - const { getByText, getByTestId } = render(); - expect(getByText("You")).toBeInTheDocument(); - expect(getByText("What is 2+2?")).toBeInTheDocument(); - expect(getByText("gpt-4")).toBeInTheDocument(); - expect(getByText("calculator")).toBeInTheDocument(); - expect(getByText("2+2 equals 4")).toBeInTheDocument(); - expect(getByTestId("response-metrics")).toBeInTheDocument(); + render(); + expect(screen.getByText("You")).toBeInTheDocument(); + expect(screen.getByText("What is 2+2?")).toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("calculator")).toBeInTheDocument(); + expect(screen.getByText("2+2 equals 4")).toBeInTheDocument(); + expect(screen.getByTestId("response-metrics")).toBeInTheDocument(); }); it("should display image attachment in user message", () => { @@ -86,10 +86,10 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByTestId, getByText } = render(); - expect(getByText("What is in this image? [Image attached]")).toBeInTheDocument(); - expect(getByTestId("chat-image-renderer")).toBeInTheDocument(); - const image = getByTestId("chat-image-renderer").querySelector("img"); + render(); + expect(screen.getByText("What is in this image? [Image attached]")).toBeInTheDocument(); + expect(screen.getByTestId("chat-image-renderer")).toBeInTheDocument(); + const image = screen.getByTestId("chat-image-renderer").querySelector("img"); expect(image).toHaveAttribute("src", "blob:test-image-url"); }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx index 8edf2174eee..cb04323e4c9 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx @@ -10,6 +10,7 @@ */ import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import EntityUsageExportModal from "./EntityUsageExportModal"; @@ -73,13 +74,13 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByRole } = renderWithProviders(); + renderWithProviders(); // Default primary action reflects CSV export - expect(getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); // Click export - await user.click(getByRole("button", { name: /Export CSV/i })); + await user.click(screen.getByRole("button", { name: /Export CSV/i })); // Verifies export function was invoked with correct parameters expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily", "Tag", "tag", {}); @@ -97,14 +98,14 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByText, getByRole } = renderWithProviders(); + renderWithProviders(); // Choose the alternate export type - click the label to trigger radio - const dailyModelLabel = getByText(/Day-by-day by tag and model/i); + const dailyModelLabel = screen.getByText(/Day-by-day by tag and model/i); await user.click(dailyModelLabel); // Export with default CSV format - const exportBtn = getByRole("button", { name: /Export CSV/i }); + const exportBtn = screen.getByRole("button", { name: /Export CSV/i }); await user.click(exportBtn); // Ensure the selected scope flowed through diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 01e00d903aa..ceb9d1ca713 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -1,4 +1,4 @@ -import { act, fireEvent, render, waitFor } from "@testing-library/react"; +import { act, fireEvent, render, waitFor, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { MountedFormHost } from "../../../tests/mounted-form-host"; import AdvancedSettings from "./advanced_settings"; @@ -35,51 +35,51 @@ describe("AdvancedSettings", () => { }); it("should render tags list", async () => { - const { getByText } = renderAdvancedSettings(); - fireEvent.click(getByText("Advanced Settings")); + renderAdvancedSettings(); + fireEvent.click(screen.getByText("Advanced Settings")); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); }); it("should render the litellm params", async () => { - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("LiteLLM Params")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Params")).toBeInTheDocument(); }); }); it("hides every PTU field when PTU cost attribution is disabled", async () => { - const { getByText, queryByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(queryByText(label)).not.toBeInTheDocument(); + expect(screen.queryByText(label)).not.toBeInTheDocument(); } - expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); }); it("shows every PTU field when PTU cost attribution is enabled", async () => { mockUsePtuCostAttributionEnabled.mockReturnValue(true); - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("PTU Count")).toBeInTheDocument(); + expect(screen.getByText("PTU Count")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(getByText(label)).toBeInTheDocument(); + expect(screen.getByText(label)).toBeInTheDocument(); } - expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx index 64074481603..d6476089cd7 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; import { getPlaceholder, Providers } from "../provider_info_helpers"; import { MountedFormHost } from "../../../tests/mounted-form-host"; @@ -6,7 +6,7 @@ import LiteLLMModelNameField from "./litellm_model_name"; describe("LitellmModelNameField", () => { it("should render", () => { - const { getByText } = render( + render( { /> , ); - expect(getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); }); it("should show Azure placeholder as 'my-deployment'", () => { - const { getByPlaceholderText, queryByPlaceholderText } = render( + render( , ); - expect(getByPlaceholderText("my-deployment")).toBeInTheDocument(); - expect(queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("my-deployment")).toBeInTheDocument(); + expect(screen.queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx index 24685f9c129..e1faf8be1df 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx @@ -26,8 +26,8 @@ const openUploadStep = async () => { describe("BulkCreateUsersButton", () => { it("should render", () => { - const { getByText } = render(); - expect(getByText("+ Bulk Invite Users")).toBeInTheDocument(); + render(); + expect(screen.getByText("+ Bulk Invite Users")).toBeInTheDocument(); }); it("parses a CSV chosen through the file input", async () => { diff --git a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx index bad92555bd5..3fc1d41dfdd 100644 --- a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it } from "vitest"; import CostOptimizationFeedbackBanner from "./cost_optimization_feedback_banner"; @@ -10,24 +10,24 @@ describe("CostOptimizationFeedbackBanner", () => { }); it("renders with a link to the feedback discussion", () => { - const { getByText } = render(); - const link = getByText("Share Feedback").closest("a"); + render(); + const link = screen.getByText("Share Feedback").closest("a"); expect(link).toHaveAttribute("href", "https://github.com/BerriAI/litellm/discussions/32172"); }); it("hides itself and persists the dismissal when the dismiss button is clicked", () => { - const { getByText, queryByText, getByLabelText } = render(); - expect(getByText("Help shape cost optimization")).toBeInTheDocument(); + render(); + expect(screen.getByText("Help shape cost optimization")).toBeInTheDocument(); - fireEvent.click(getByLabelText("Dismiss banner")); + fireEvent.click(screen.getByLabelText("Dismiss banner")); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); expect(localStorage.getItem(STORAGE_KEY)).toBe("true"); }); it("stays dismissed on remount once persisted", () => { localStorage.setItem(STORAGE_KEY, "true"); - const { queryByText } = render(); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + render(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index b8d8e3ba9c8..700d19eb13c 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -108,7 +108,7 @@ beforeEach(() => { test("renders organization view after loading data", async () => { mockUseOrganization.mockReturnValue({ data: mockOrg, isLoading: false } as any); - const { findAllByText } = renderWithProviders( + renderWithProviders( {}} @@ -120,7 +120,7 @@ test("renders organization view after loading data", async () => { />, ); - const [orgName] = await findAllByText("Acme Corp"); + const [orgName] = await screen.findAllByText("Acme Corp"); expect(orgName).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index b97fef32402..762b23f413e 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -77,21 +77,21 @@ describe("Settings", () => { }); it("should render the logging callbacks tab when access token is provided", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); }); it("should display additional settings tabs", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); - expect(getByText("Alerting Types")).toBeInTheDocument(); - expect(getByText("Alerting Settings")).toBeInTheDocument(); - expect(getByText("Email Alerts")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("Alerting Types")).toBeInTheDocument(); + expect(screen.getByText("Alerting Settings")).toBeInTheDocument(); + expect(screen.getByText("Email Alerts")).toBeInTheDocument(); }); }); @@ -279,13 +279,13 @@ describe("Settings", () => { }); it("should display CloudZero Cost Tracking tab", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx index 3698b68155e..5bd48b1aa4c 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PaginationStatusAlerts from "./PaginationStatusAlerts"; @@ -6,7 +6,7 @@ import PaginationStatusAlerts from "./PaginationStatusAlerts"; describe("PaginationStatusAlerts", () => { it("shows page progress and wires the Stop button while fetching", () => { const cancel = vi.fn(); - const { getByRole, getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); - fireEvent.click(getByRole("button", { name: "Stop" })); + expect(screen.getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Stop" })); expect(cancel).toHaveBeenCalledTimes(1); }); it("shows the partial-data notice after a cancel, frozen at the last fetched page", () => { - const { getByText } = render( + render( { />, ); - expect(getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); + expect(screen.getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); }); it("names the subject it is fetching", () => { - const { getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); }); it("renders nothing when idle", () => { diff --git a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx index 6e4ab14be33..c81ade526dd 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import React from "react"; import { describe, expect, it } from "vitest"; import { AreaChart } from "./area_chart"; @@ -21,9 +21,9 @@ describe("AreaChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx index a322eeb3ad0..3cd7bd2fd08 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx @@ -21,9 +21,9 @@ describe("BarChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 6592ff8c357..705679fe2e6 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -445,7 +445,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { ); fireEvent.click(screen.getByText("Settings")); - expect(screen.getByText("Budget Reset").parentElement?.textContent).toContain("Every 30d"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Every 30d"); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { @@ -456,7 +456,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { fireEvent.click(screen.getByText("Mock Submit")); await waitFor(() => { - expect(screen.getByText("Budget Reset").parentElement?.textContent).toBe("Budget ResetNever"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx index f506c0e51d7..a02f80f7800 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx @@ -381,6 +381,6 @@ describe("KeyInfoView budget reset visibility", () => { await waitFor(() => { expect(screen.getByText("Budget Reset")).toBeInTheDocument(); }); - expect(screen.getByText("Budget Reset").parentElement).toHaveTextContent("Never"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0be80c3e173..0dd0dd6d6af 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -895,7 +895,7 @@ export default function KeyInfoView({

Budget Reset

-

+

{currentKeyData.budget_reset_at ? `${currentKeyData.budget_duration ? `Every ${currentKeyData.budget_duration}, next ` : ""}${formatTimestamp(currentKeyData.budget_reset_at)}` : "Never"} From 95c7ca8801b2fdde31763f8d5869355646bb93e4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:59:08 +0000 Subject: [PATCH 39/92] fix(tests): restore module attributes after reload in mcp identity env tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/test_mcp_server_identity_env.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py index ac7082c2668..934810c305f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py @@ -23,6 +23,14 @@ MGMT_MODULE = "litellm.proxy.management_endpoints.mcp_management_endpoints" @contextlib.contextmanager def _env_and_reload(**env): saved = {key: os.environ.get(key) for key in env} + utils_module = importlib.import_module(UTILS_MODULE) + mgmt_module = importlib.import_module(MGMT_MODULE) + # Restore the pre-reload module attributes afterwards instead of reloading + # a third time: a reload re-creates every class in the module, so modules + # that imported names like MCPMissingUserEnvVarsError before this test + # would keep raising the old class while pytest.raises in later tests + # matches the new one + snapshots = {module: dict(vars(module)) for module in (utils_module, mgmt_module)} def _apply_env(values): for key, value in values.items(): @@ -32,8 +40,8 @@ def _env_and_reload(**env): os.environ[key] = value def _reload(): - utils = importlib.reload(importlib.import_module(UTILS_MODULE)) - mgmt = importlib.reload(importlib.import_module(MGMT_MODULE)) + utils = importlib.reload(utils_module) + mgmt = importlib.reload(mgmt_module) return utils, mgmt try: @@ -41,7 +49,10 @@ def _env_and_reload(**env): yield _reload() finally: _apply_env(saved) - _reload() + for module, snapshot in snapshots.items(): + for key in [key for key in vars(module) if key not in snapshot]: + delattr(module, key) + vars(module).update(snapshot) def test_defaults_used_when_env_unset(): From 215bf03373617c6de89aa47c19dc3be7d5094634 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:59:38 +0000 Subject: [PATCH 40/92] refactor(types): replace Any with precise types across 73 modules Narrows reportAny / reportExplicitAny hot spots in provider transformations, proxy endpoints, integrations and secret managers by introducing TypedDicts, Protocols and object-typed boundaries instead of Any, then ratchets the budget ceilings down to match. reportAny 14765 -> 14076, reportExplicitAny 4493 -> 4128, ANN401 387 -> 307 --- basedpyright-code-budget.json | 16 +-- litellm/caching/caching.py | 9 +- litellm/caching/qdrant_semantic_cache.py | 18 ++- .../handler.py | 14 +-- litellm/google_genai/main.py | 30 ++--- litellm/images/main.py | 24 ++-- .../SlackAlerting/slack_alerting.py | 24 +++- .../bitbucket/bitbucket_prompt_manager.py | 37 +++--- litellm/integrations/cloudzero/transform.py | 27 ++++- litellm/integrations/custom_guardrail.py | 13 ++- litellm/integrations/datadog/datadog.py | 44 ++++--- .../integrations/datadog/datadog_llm_obs.py | 19 ++-- .../integrations/dotprompt/prompt_manager.py | 25 ++-- litellm/integrations/galileo.py | 35 +++++- litellm/integrations/gitlab/gitlab_client.py | 88 ++++++++++++-- litellm/integrations/langfuse/langfuse.py | 39 ++++--- litellm/integrations/opik/opik.py | 27 ++++- .../opik/opik_payload_builder/extractors.py | 21 ++-- litellm/integrations/otel/plumbing/metrics.py | 45 ++++++-- .../vector_store_pre_call_hook.py | 6 +- .../litellm_core_utils/realtime_streaming.py | 2 +- .../streaming_chunk_builder_utils.py | 14 ++- .../a2a/chat/guardrail_translation/handler.py | 27 +++-- litellm/llms/anthropic/chat/transformation.py | 51 ++++++--- litellm/llms/anthropic/files/handler.py | 4 +- litellm/llms/azure/azure.py | 16 +-- litellm/llms/azure_ai/agents/handler.py | 20 +--- .../llms/bedrock/realtime/transformation.py | 6 +- .../black_forest_labs/image_edit/handler.py | 56 +++++++-- .../image_generation/handler.py | 35 ++++-- litellm/llms/codestral/completion/handler.py | 52 ++++++++- .../llms/deepinfra/rerank/transformation.py | 39 ++++++- .../gemini/interactions/transformation.py | 62 ++++++++-- litellm/llms/gemini/videos/transformation.py | 36 +++--- .../huggingface/embedding/transformation.py | 22 +++- .../llms/openai/chat/gpt_transformation.py | 14 ++- .../chat/guardrail_translation/handler.py | 19 ++-- .../llms/openai/responses/transformation.py | 45 ++++++-- litellm/llms/openai_like/chat/handler.py | 28 ++++- .../image_generation/transformation.py | 24 +++- litellm/llms/sap/credentials.py | 55 ++++++--- .../llms/vertex_ai/files/transformation.py | 49 +++++--- .../llms/vertex_ai/gemini/transformation.py | 18 +-- litellm/llms/vertex_ai/vertex_llm_base.py | 53 ++++++--- litellm/passthrough/main.py | 28 ++--- .../mcp_server/semantic_tool_filter.py | 19 ++-- .../proxy/agent_endpoints/a2a_endpoints.py | 19 +++- litellm/proxy/auth/handle_jwt.py | 56 +++++++-- litellm/proxy/common_utils/debug_utils.py | 107 +++++++++++++----- litellm/proxy/db/db_spend_update_writer.py | 10 +- .../guardrails/guardrail_hooks/akto/akto.py | 8 +- .../guardrail_hooks/grayswan/grayswan.py | 50 ++++++-- .../guardrails/guardrail_hooks/lasso/lasso.py | 11 +- .../guardrail_hooks/pillar/pillar.py | 54 +++++++-- .../semantic_guard/semantic_guard.py | 15 ++- .../guardrail_hooks/tool_permission.py | 56 ++++++--- .../vigil_guard/vigil_guard.py | 2 +- litellm/proxy/hooks/litellm_skills/main.py | 23 +++- .../hooks/parallel_request_limiter_v3.py | 24 +++- .../model_management_endpoints.py | 4 +- .../organization_endpoints.py | 15 ++- litellm/proxy/management_endpoints/ui_sso.py | 6 +- .../vertex_passthrough_logging_handler.py | 4 +- .../proxy/response_api_endpoints/endpoints.py | 12 +- litellm/proxy/route_llm_request.py | 4 +- litellm/proxy/video_endpoints/endpoints.py | 10 +- litellm/rag/main.py | 11 +- .../mcp/litellm_proxy_mcp_handler.py | 4 +- litellm/router_strategy/budget_limiter.py | 13 ++- .../complexity_router/complexity_router.py | 4 +- .../hashicorp_secret_manager.py | 105 ++++++++++++++--- litellm/types/llms/openai.py | 18 +-- litellm/types/router.py | 9 +- .../vector_stores/vector_store_registry.py | 6 +- ruff-strict-budget.json | 14 +-- type-discipline-budget.json | 8 +- 76 files changed, 1458 insertions(+), 579 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a07b9352659..df52069e71f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 14765 + "limit": 14076 }, "reportArgumentType": { "limit": 2216 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4493 + "limit": 4128 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5607 + "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15310 + "limit": 15306 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38368 + "limit": 38350 }, "reportUnknownParameterType": { - "limit": 19633 + "limit": 19626 }, "reportUnknownVariableType": { - "limit": 29908 + "limit": 29890 }, "reportUnnecessaryCast": { "limit": 111 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 828 + "limit": 826 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index cefe6aae9ed..754815fce47 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -12,6 +12,7 @@ import hashlib import json import time import traceback +from collections.abc import Mapping from enum import Enum from typing import Any, Final @@ -506,7 +507,7 @@ class Cache: def _get_cache_logic( self, - cached_result: Any | None, + cached_result: object | None, max_age: float | None, ): """ @@ -538,8 +539,8 @@ class Cache: return cached_result @staticmethod - def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - cache_lookup_kwargs: Final[dict[str, Any]] = {} + def _get_safe_cache_lookup_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: + cache_lookup_kwargs: Final[dict[str, object]] = {} for prompt_kwarg in ("messages", "input"): if prompt_kwarg in kwargs: cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] @@ -552,7 +553,7 @@ class Cache: @staticmethod def _update_metadata_from_cache_lookup_kwargs( - original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any] + original_kwargs: Mapping[str, object], cache_lookup_kwargs: Mapping[str, object] ) -> None: original_metadata: Final = original_kwargs.get("metadata") cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata") diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 4898700c403..c5876e993d3 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -12,7 +12,7 @@ import ast import asyncio import json import os -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import print_verbose @@ -39,6 +39,12 @@ if TYPE_CHECKING: from litellm.router import Router +class _QdrantCollectionDetailsResponse(Protocol): + """The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object.""" + + def json(self) -> dict[str, object]: ... + + class QdrantSemanticCache(BaseCache): CACHE_KEY_FIELD_NAME = "litellm_cache_key" embedding_max_input_tokens: int | None = None @@ -115,15 +121,15 @@ class QdrantSemanticCache(BaseCache): raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}") if collection_exists.json()["result"]["exists"]: - collection_details = self.sync_client.get( + collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get( url=f"{self.qdrant_api_base}/collections/{self.collection_name}", headers=self.headers, ) - self.collection_info = collection_details.json() + self.collection_info: dict[str, object] = collection_details.json() print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}") self._ensure_cache_key_payload_index() else: - quantization_params: dict[str, Any] + quantization_params: dict[str, dict[str, object]] if quantization_config is None or quantization_config == "binary": quantization_params = { "binary": { @@ -214,7 +220,7 @@ class QdrantSemanticCache(BaseCache): resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router), ) - def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: + def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse: """Embed via the proxy Router when it serves the model, else direct.""" try: from litellm.proxy.proxy_server import llm_model_list, llm_router @@ -241,7 +247,7 @@ class QdrantSemanticCache(BaseCache): num_retries=0, ) - async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: + async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse: try: from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 727c39c16ec..f494d6610a1 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -45,14 +45,14 @@ class ResponsesToCompletionBridgeHandler: return bool(stream) @staticmethod - def _is_preformatted_cached_chat_stream(result: Any) -> bool: + def _is_preformatted_cached_chat_stream(result: object) -> bool: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response" @staticmethod def _coerce_response_object( - response_obj: Any, + response_obj: object, hidden_params: dict | None, ) -> "ResponsesAPIResponse": if isinstance(response_obj, ResponsesAPIResponse): @@ -78,8 +78,8 @@ class ResponsesToCompletionBridgeHandler: for _ in stream_iter: pass - completed: Final = getattr(stream_iter, "completed_response", None) - response_obj: Final = getattr(completed, "response", None) if completed else None + completed: Final[object] = getattr(stream_iter, "completed_response", None) + response_obj: Final[object] = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") @@ -93,8 +93,8 @@ class ResponsesToCompletionBridgeHandler: async for _ in stream_iter: pass - completed: Final = getattr(stream_iter, "completed_response", None) - response_obj: Final = getattr(completed, "response", None) if completed else None + completed: Final[object] = getattr(stream_iter, "completed_response", None) + response_obj: Final[object] = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") @@ -157,7 +157,7 @@ class ResponsesToCompletionBridgeHandler: def completion( self, *args, **kwargs ) -> Union[ - Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", ]: diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index b5815bd3f7c..c1822e4720d 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -52,10 +52,10 @@ class GenerateContentSetupResult(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) model: str - request_body: dict[str, Any] + request_body: dict[str, object] custom_llm_provider: str generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None - generate_content_config_dict: dict[str, Any] + generate_content_config_dict: dict[str, object] native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj @@ -68,7 +68,7 @@ class GenerateContentHelper: @staticmethod def mock_generate_content_response( mock_response: str = "This is a mock response from Google GenAI generate_content.", - ) -> dict[str, Any]: + ) -> dict[str, object]: """Mock response for generate_content for testing purposes""" return { "text": mock_response, @@ -239,9 +239,9 @@ async def agenerate_content( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -307,9 +307,9 @@ def generate_content( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -397,9 +397,9 @@ async def agenerate_content_stream( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -492,9 +492,9 @@ def generate_content_stream( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, diff --git a/litellm/images/main.py b/litellm/images/main.py index 1688087c2da..617f8e08ab6 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -3,7 +3,7 @@ import contextvars import importlib from collections.abc import Coroutine from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload if TYPE_CHECKING: from litellm.images.utils import ImageEditRequestUtils @@ -151,7 +151,7 @@ def image_generation( *, aimg_generation: Literal[True], **kwargs, -) -> Coroutine[Any, Any, ImageResponse]: +) -> Coroutine[object, object, ImageResponse]: ... @@ -197,7 +197,7 @@ def image_generation( api_version: str | None = None, custom_llm_provider=None, **kwargs, -) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: +) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Maps the https://api.openai.com/v1/images/generations endpoint. @@ -723,14 +723,14 @@ def image_edit( user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: +) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Maps the image edit functionality, similar to OpenAI's images/edits endpoint. """ @@ -769,7 +769,7 @@ def image_edit( images: Final = image if isinstance(image, list) else ([image] if image is not None else []) headers_from_kwargs: Final = kwargs.get("headers") - merged_extra_headers: Final[dict[str, Any]] = {} + merged_extra_headers: Final[dict[str, object]] = {} if isinstance(headers_from_kwargs, dict): merged_extra_headers.update(headers_from_kwargs) if isinstance(extra_headers, dict): @@ -974,9 +974,9 @@ async def aimage_edit( user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1044,7 +1044,7 @@ async def aimage_edit( ) -def __getattr__(name: str) -> Any: +def __getattr__(name: str) -> type["ImageEditRequestUtils"]: """Lazy import handler for images.main module""" if name == "ImageEditRequestUtils": # Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94d734546be..c137164ecdb 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -545,7 +545,6 @@ class SlackAlerting(CustomBatchLogger): # Get the appropriate budget alert type handler budget_alert_class: Final = get_budget_alert_type(type) _id: Final = budget_alert_class.get_id(user_info) - user_info_json: Final = user_info.model_dump(exclude_none=True) user_info_str: Final = self._get_user_info_str(user_info) event_message = budget_alert_class.get_event_message() @@ -575,7 +574,22 @@ class SlackAlerting(CustomBatchLogger): webhook_event = WebhookEvent( event=event, event_message=event_message, - **user_info_json, + spend=user_info.spend, + max_budget=user_info.max_budget, + soft_budget=user_info.soft_budget, + token=user_info.token, + customer_id=user_info.customer_id, + user_id=user_info.user_id, + team_id=user_info.team_id, + team_alias=user_info.team_alias, + organization_id=user_info.organization_id, + user_email=user_info.user_email, + key_alias=user_info.key_alias, + projected_exceeded_date=user_info.projected_exceeded_date, + projected_spend=user_info.projected_spend, + event_group=user_info.event_group, + alert_emails=user_info.alert_emails, + max_budget_alert_emails=user_info.max_budget_alert_emails, ) await self.send_alert( message=event_message + "\n\n" + user_info_str, @@ -657,7 +671,7 @@ class SlackAlerting(CustomBatchLogger): """ Create a standard message for a budget alert """ - _all_fields_as_dict: Final = user_info.model_dump(exclude_none=True) + _all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True) _all_fields_as_dict.pop("token") msg = "" for k, v in _all_fields_as_dict.items(): @@ -1006,7 +1020,7 @@ class SlackAlerting(CustomBatchLogger): except Exception: pass - async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any): + async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: object): base_model_from_user: Final = getattr(passed_model_info, "base_model", None) model_info = {} base_model = "" @@ -1973,7 +1987,7 @@ Model Info: try: message = f"`{event_name}`\n" - key_event_dict: Final = key_event.model_dump() + key_event_dict: Final[dict[str, object]] = key_event.model_dump() # Add Created by information first message += "*Action Done by:*\n" diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index 6a03e3ee93c..ff34bd91e31 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -3,6 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system Fetches .prompt files from BitBucket repositories and provides team-based access control. """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from jinja2 import DictLoader, select_autoescape @@ -65,7 +66,7 @@ class BitBucketTemplateManager: def __init__( self, - bitbucket_config: dict[str, Any], + bitbucket_config: Mapping[str, object], prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config @@ -123,7 +124,7 @@ class BitBucketTemplateManager: template_content = content # Parse YAML frontmatter - metadata: dict[str, Any] = {} + metadata: dict[str, object] = {} if frontmatter_str: try: import yaml @@ -141,9 +142,9 @@ class BitBucketTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, object]: """Basic YAML parser for simple cases when PyYAML is not available.""" - result: Final[dict[str, Any]] = {} + result: Final[dict[str, object]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): @@ -162,7 +163,7 @@ class BitBucketTemplateManager: result[key] = value.strip("\"'") return result - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str: + def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str: """Render a template with the given variables.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -209,7 +210,7 @@ class BitBucketPromptManager(CustomPromptManagement): def __init__( self, - bitbucket_config: dict[str, Any], + bitbucket_config: Mapping[str, object], prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config @@ -234,7 +235,7 @@ class BitBucketPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, ) -> tuple[str, dict[str, Any]]: """ Get a prompt template and render it with variables. @@ -267,12 +268,12 @@ class BitBucketPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -316,9 +317,9 @@ class BitBucketPromptManager(CustomPromptManagement): except Exception as e: # Log error but don't fail the call - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e) return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]: @@ -384,14 +385,14 @@ class BitBucketPromptManager(CustomPromptManagement): def post_call_hook( self, user_id: str | None, - response: Any, + response: object, input_messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: Mapping[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> Any: + ) -> object: """ Post-call hook for any post-processing after the LLM call. """ diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index f0d4d67fc22..ffc8fe1c1f5 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -19,14 +19,29 @@ """Transform LiteLLM data to CloudZero AnyCost CBF format.""" from datetime import datetime -from typing import Any, Final +from typing import Final, SupportsFloat, SupportsIndex, SupportsInt import polars as pl +from typing_extensions import Buffer from ...types.integrations.cloudzero import CBFRecord from .cz_resource_names import CZEntityType, CZRNGenerator +def _as_int(value: object) -> int: + """The integer form of a spend table cell, computed the way :func:`int` computes it.""" + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex)): + return int(value) + raise TypeError(f"int() argument must be a string or a number, not {type(value).__name__!r}") + + +def _as_float(value: object) -> float: + """The floating point form of a spend table cell, computed the way :func:`float` computes it.""" + if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(value) + raise TypeError(f"float() argument must be a string or a number, not {type(value).__name__!r}") + + class CBFTransformer: """Transform LiteLLM usage data to CloudZero Billing Format (CBF).""" @@ -82,15 +97,15 @@ class CBFTransformer: return pl.DataFrame(cbf_data) - def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord: + def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord: """Create a single CBF record from LiteLLM daily spend row.""" # Parse date (daily spend tables use date strings like '2025-04-19') usage_date: Final = self._parse_date(row.get("date")) # Calculate total tokens - prompt_tokens: Final = int(row.get("prompt_tokens", 0)) - completion_tokens: Final = int(row.get("completion_tokens", 0)) + prompt_tokens: Final = _as_int(row.get("prompt_tokens", 0)) + completion_tokens: Final = _as_int(row.get("completion_tokens", 0)) total_tokens: Final = prompt_tokens + completion_tokens # Create CloudZero Resource Name (CZRN) as resource_id @@ -154,7 +169,7 @@ class CBFTransformer: "time/usage_start": ( usage_date.isoformat() if usage_date else None ), # Required: ISO-formatted UTC datetime - "cost/cost": float(row.get("spend", 0.0)), # Required: billed cost + "cost/cost": _as_float(row.get("spend", 0.0)), # Required: billed cost "resource/id": resource_id, # CZRN (CloudZero Resource Name) # Usage metrics for token consumption "usage/amount": total_tokens, # Numeric value of tokens consumed @@ -187,7 +202,7 @@ class CBFTransformer: return CBFRecord(cbf_record) - def _parse_date(self, date_str) -> datetime | None: + def _parse_date(self, date_str: object) -> datetime | None: """Parse date string from daily spend tables (e.g., '2025-04-19').""" if date_str is None: return None diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8dc6881d23e..e87ac9521ae 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -2,6 +2,7 @@ import contextvars import hashlib import os import secrets +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args @@ -227,13 +228,13 @@ class CustomGuardrail(CustomLogger): ) super().__init__(**kwargs) - def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str: + def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str: """Return a custom violation message if template is configured.""" if not self.violation_message_template: return default - format_context: Final[dict[str, Any]] = {"default_message": default} + format_context: Final[dict[str, object]] = {"default_message": default} if context: format_context.update(context) try: @@ -661,7 +662,7 @@ class CustomGuardrail(CustomLogger): value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails") return value if isinstance(value, list) else [] - def _is_valid_response_type(self, result: Any) -> bool: + def _is_valid_response_type(self, result: object) -> bool: """ Check if result is a valid LLMResponseTypes instance. @@ -722,7 +723,7 @@ class CustomGuardrail(CustomLogger): return None return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None: + def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -747,7 +748,7 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool: + def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: marker: Final = self._pre_call_marker() if marker is None: return False @@ -1170,7 +1171,7 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ # Convert None to empty dict to satisfy type requirements - guardrail_response: dict[str, Any] | str = {} if response is None else response + guardrail_response: dict[str, object] | str = {} if response is None else response # For apply_guardrail functions in custom_code_guardrail scenario, # simplify the logged response to "allow", "deny", or "mask" diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 04f1c6dff15..866076a3c49 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -20,10 +20,11 @@ import time import traceback from collections.abc import Sequence from datetime import datetime as datetimeObj -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx from httpx import Response +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -62,6 +63,18 @@ from litellm.types.utils import StandardLoggingPayload from ..additional_logging_utils import AdditionalLoggingUtils +if TYPE_CHECKING: + from fastapi import HTTPException + + from litellm.proxy._types import UserAPIKeyAuth + + +class _DatadogLoggingKwargs(TypedDict, total=False): + """The subset of logging ``kwargs`` that the Datadog payload builder reads.""" + + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + + # max number of logs DD API can accept @@ -87,6 +100,11 @@ def _resolve_dd_batch_size() -> int: return max(1, min(value, DD_MAX_BATCH_SIZE)) +def _span_attribute(span: object, name: str) -> object: + """Read an optional attribute off whatever span object the active tracer hands back.""" + return getattr(span, name, None) + + class DataDogLogger( CustomBatchLogger, AdditionalLoggingUtils, @@ -271,9 +289,9 @@ class DataDogLogger( self, request_data: dict, original_exception: Exception, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", traceback_str: str | None = None, - ) -> Any | None: + ) -> "HTTPException | None": """ Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog. @@ -297,7 +315,7 @@ class DataDogLogger( status_code = int(_code) # Use project-standard sanitized user context when running in proxy - user_context: dict[str, Any] = {} + user_context: dict[str, object] = {} try: from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, @@ -553,8 +571,8 @@ class DataDogLogger( def create_datadog_logging_payload( self, - kwargs: dict | Any, - response_obj: Any, + kwargs: _DatadogLoggingKwargs, + response_obj: object, start_time: datetime.datetime, end_time: datetime.datetime, ) -> DatadogPayload: @@ -562,8 +580,8 @@ class DataDogLogger( Helper function to create a datadog payload for logging Args: - kwargs (Union[dict, Any]): request kwargs - response_obj (Any): llm api response + kwargs: request kwargs, read for its standard logging object + response_obj: llm api response start_time (datetime.datetime): start time of request end_time (datetime.datetime): end time of request @@ -625,7 +643,7 @@ class DataDogLogger( self, payload: ServiceLoggerPayload, error: str | None = "", - parent_otel_span: Any | None = None, + parent_otel_span: object = None, start_time: datetimeObj | float | None = None, end_time: float | datetimeObj | None = None, event_metadata: dict | None = None, @@ -659,7 +677,7 @@ class DataDogLogger( self, payload: ServiceLoggerPayload, error: str | None = "", - parent_otel_span: Any | None = None, + parent_otel_span: object = None, start_time: datetimeObj | float | None = None, end_time: float | datetimeObj | None = None, event_metadata: dict | None = None, @@ -696,7 +714,7 @@ class DataDogLogger( def _create_v0_logging_payload( self, - kwargs: dict | Any, + kwargs: dict, response_obj: Any, start_time: datetime.datetime, end_time: datetime.datetime, @@ -810,11 +828,11 @@ class DataDogLogger( if current_span is None: return None - trace_id: Final = getattr(current_span, "trace_id", None) + trace_id: Final = _span_attribute(current_span, "trace_id") if trace_id is None: return None - span_id: Final = getattr(current_span, "span_id", None) + span_id: Final = _span_attribute(current_span, "span_id") trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)} if span_id is not None: trace_context["span_id"] = str(span_id) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 704f0323e95..e5789965c6e 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -9,6 +9,7 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp import asyncio import json import os +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Any, Final, Literal @@ -334,7 +335,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): def _get_response_messages( self, standard_logging_payload: StandardLoggingPayload, call_type: str | None - ) -> list[Any]: + ) -> list[object]: """ Get the messages from the response object @@ -484,7 +485,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]: + def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]: if messages is None: return [] if isinstance(messages, str): @@ -495,11 +496,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): return [str(messages.get("content", ""))] return [] - def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: + def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload """ - _metadata: Final[dict[str, Any]] = { + _metadata: Final[dict[str, object]] = { "model_name": standard_logging_payload.get("model", "unknown"), "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"), "id": standard_logging_payload.get("id", "unknown"), @@ -647,7 +648,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): return spend_metrics - def _process_input_messages_preserving_tool_calls(self, messages: list[Any]) -> list[dict[str, Any]]: + def _process_input_messages_preserving_tool_calls(self, messages: Sequence[object]) -> list[dict[str, object]]: """ Process input messages while preserving tool_calls and tool message types. @@ -671,13 +672,13 @@ class DataDogLLMObsLogger(CustomBatchLogger): return processed @staticmethod - def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, Any]: + def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, object]: """ Extract tool call information into key-value pairs for Datadog metadata. Similar to OpenTelemetry's implementation but adapted for Datadog's format. """ - kv_pairs: Final[dict[str, Any]] = {} + kv_pairs: Final[dict[str, object]] = {} for idx, tool_call in enumerate(tool_calls): try: # Extract tool call ID @@ -712,11 +713,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): return kv_pairs - def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: + def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: """ Extract tool call information from both input messages and response for Datadog metadata. """ - tool_call_metadata: Final[dict[str, Any]] = {} + tool_call_metadata: Final[dict[str, object]] = {} try: # Extract tool calls from input messages diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index fd0b17ba746..9c82ff7c5ba 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -3,12 +3,21 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d """ import re +from collections.abc import Mapping from pathlib import Path from typing import Any, Final import yaml from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import NotRequired, ReadOnly, TypedDict + + +class _PromptFileJson(TypedDict): + """JSON form of a .prompt file: rendered template text plus its frontmatter.""" + + content: ReadOnly[NotRequired[str]] + metadata: ReadOnly[NotRequired[dict[str, object]]] def strip_version_suffix(prompt_id: str) -> str | None: @@ -167,7 +176,7 @@ class PromptManager: template_id=prompt_id, ) - def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]: + def _parse_frontmatter(self, content: str) -> tuple[dict[str, object], str]: """Parse YAML frontmatter from prompt content.""" # Match YAML frontmatter between --- delimiters frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$" @@ -178,7 +187,7 @@ class PromptManager: template_content = match.group(2) try: - frontmatter = yaml.safe_load(frontmatter_yaml) or {} + frontmatter: dict[str, object] = yaml.safe_load(frontmatter_yaml) or {} except yaml.YAMLError as e: raise ValueError(f"Invalid YAML frontmatter: {e}") else: @@ -191,7 +200,7 @@ class PromptManager: def render( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, version: int | None = None, ) -> str: """ @@ -231,7 +240,7 @@ class PromptManager: except Exception as e: raise ValueError(f"Error rendering template '{prompt_id}': {e}") - def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None: + def _validate_input(self, variables: Mapping[str, object], schema: Mapping[str, str]) -> None: """Basic validation of input variables against schema.""" for field_name, field_type in schema.items(): if field_name in variables: @@ -291,7 +300,7 @@ class PromptManager: """Get a list of all available prompt IDs.""" return list(self.prompts.keys()) - def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None: + def get_prompt_metadata(self, prompt_id: str) -> dict[str, object] | None: """Get metadata for a specific prompt.""" template: Final = self.prompts.get(prompt_id) return template.metadata if template else None @@ -302,12 +311,12 @@ class PromptManager: if self.prompt_directory: self._load_prompts() - def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None: + def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, object] | None = None) -> None: """Add a prompt template programmatically.""" template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id) self.prompts[prompt_id] = template - def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]: + def prompt_file_to_json(self, file_path: str | Path) -> _PromptFileJson: """Convert a .prompt file to JSON format. Args: @@ -324,7 +333,7 @@ class PromptManager: return {"content": template_content.strip(), "metadata": frontmatter} - def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str: + def json_to_prompt_file(self, prompt_data: _PromptFileJson) -> str: """Convert JSON prompt data to .prompt file format. Args: diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 23727801a6f..b27618993a3 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -6,10 +6,11 @@ import re import uuid from collections.abc import Mapping, Sequence from datetime import datetime, timezone, tzinfo -from typing import Any, Final, TypedDict, cast +from typing import Any, Final, Protocol, cast import httpx from pydantic import BaseModel, Field +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -35,6 +36,34 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai" GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000 +class _GalileoLoginBody(TypedDict): + """Decoded body of the Galileo login response.""" + + access_token: ReadOnly[str] + + +class _GalileoLoginResponse(Protocol): + """The login call's HTTP response, read for the access token it carries.""" + + def json(self) -> _GalileoLoginBody: ... + + +class _JsonResponse(Protocol): + """An HTTP response read only for whatever JSON body it decodes to.""" + + def json(self) -> object: ... + + +def _login_access_token(response: _GalileoLoginResponse) -> str: + """Read the bearer token out of a Galileo login response body.""" + return response.json()["access_token"] + + +def _decoded_body(response: _JsonResponse) -> object: + """Decode a response body without asserting anything about its shape.""" + return response.json() + + class GalileoStandardLoggingFields(TypedDict, total=False): call_type: str model: str @@ -156,7 +185,7 @@ class GalileoObserve(CustomLogger): }, ) galileo_login_response.raise_for_status() - access_token: Final = galileo_login_response.json()["access_token"] + access_token: Final = _login_access_token(galileo_login_response) self.headers = { "accept": "application/json", "Content-Type": "application/json", @@ -421,7 +450,7 @@ class GalileoObserve(CustomLogger): try: verbose_logger.debug( "Galileo Logger HTTP error response json: %s", - response.json(), + _decoded_body(response), ) except Exception: pass diff --git a/litellm/integrations/gitlab/gitlab_client.py b/litellm/integrations/gitlab/gitlab_client.py index 0690ccc8c15..813a2ef2821 100644 --- a/litellm/integrations/gitlab/gitlab_client.py +++ b/litellm/integrations/gitlab/gitlab_client.py @@ -4,12 +4,80 @@ Now supports selecting a tag via `config["tag"]`; falls back to branch ("main"). """ import base64 -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, TypedDict from urllib.parse import quote +from typing_extensions import ReadOnly + from litellm.llms.custom_httpx.http_handler import HTTPHandler +class GitLabFilePayload(TypedDict, total=False): + """A repository-files API entry.""" + + content: ReadOnly[str] + encoding: ReadOnly[str] + + +class GitLabTreeEntry(TypedDict, total=False): + """A repository-tree API entry.""" + + path: ReadOnly[str] + type: ReadOnly[str] + + +class GitLabBranch(TypedDict, total=False): + """A repository-branches API entry.""" + + name: ReadOnly[str] + type: ReadOnly[str] + + +class GitLabFileMetadata(TypedDict): + """The response headers a raw file request exposes as metadata.""" + + content_type: ReadOnly[str | None] + content_length: ReadOnly[str | None] + last_modified: ReadOnly[str | None] + + +class _FileJsonResponse(Protocol): + def json(self) -> GitLabFilePayload: ... + + +class _TreeJsonResponse(Protocol): + def json(self) -> Sequence[GitLabTreeEntry] | None: ... + + +class _ProjectJsonResponse(Protocol): + def json(self) -> Mapping[str, object]: ... + + +class _BranchesJsonResponse(Protocol): + def json(self) -> Sequence[GitLabBranch] | None: ... + + +def _file_payload(resp: _FileJsonResponse) -> GitLabFilePayload: + """The JSON body of a repository-files response.""" + return resp.json() + + +def _tree_entries(resp: _TreeJsonResponse) -> Sequence[GitLabTreeEntry]: + """The entries of a repository-tree response.""" + return resp.json() or [] + + +def _project_info(resp: _ProjectJsonResponse) -> Mapping[str, object]: + """The JSON body of a project response.""" + return resp.json() + + +def _branch_entries(resp: _BranchesJsonResponse) -> Sequence[GitLabBranch] | None: + """The JSON body of a repository-branches response.""" + return resp.json() + + class GitLabClient: """ Client for interacting with the GitLab API to fetch files. @@ -42,12 +110,12 @@ class GitLabClient: self.project: str | int = project self.access_token: str = str(access_token) - self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth' + self.auth_method: str = config.get("auth_method", "token") # 'token' or 'oauth' self.branch = config.get("branch", None) if not self.branch: self.branch = "main" self.tag = config.get("tag") - self.base_url = config.get("base_url", "https://gitlab.com/api/v4") + self.base_url: str = config.get("base_url", "https://gitlab.com/api/v4") if not all([self.project, self.access_token]): raise ValueError("project and access_token are required") @@ -159,7 +227,7 @@ class GitLabClient: if resp.status_code == 404: return None resp.raise_for_status() - data: Final = resp.json() + data: Final = _file_payload(resp) content: Final = data.get("content") encoding: Final = data.get("encoding", "") if content and encoding == "base64": @@ -208,7 +276,7 @@ class GitLabClient: return [] resp.raise_for_status() - data: Final = resp.json() or [] + data: Final = _tree_entries(resp) files: Final[list[str]] = [] for item in data: if item.get("type") == "blob": @@ -229,13 +297,13 @@ class GitLabClient: raise Exception("Authentication failed. Check your GitLab token and auth_method.") raise Exception(f"Failed to list files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """Get information about the project/repository.""" url: Final = f"{self.base_url}/projects/{self._project_enc}" try: resp: Final = self.http_handler.get(url, headers=self.headers) resp.raise_for_status() - return resp.json() + return _project_info(resp) except Exception as e: raise Exception(f"Failed to get repository info: {e}") @@ -247,18 +315,18 @@ class GitLabClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> list[GitLabBranch]: """Get list of branches in the repository.""" url: Final = f"{self.base_url}/projects/{self._project_enc}/repository/branches" try: resp: Final = self.http_handler.get(url, headers=self.headers) resp.raise_for_status() - data: Final = resp.json() + data: Final = _branch_entries(resp) return data if isinstance(data, list) else [] except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> GitLabFileMetadata | None: """ Get minimal metadata about a file via RAW endpoint headers at a given ref. diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 296c2b5714e..9576eabaa34 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -89,7 +89,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): - prompt_tokens_details: Final = getattr(usage_obj, "prompt_tokens_details", None) + prompt_tokens_details: Final[object] = getattr(usage_obj, "prompt_tokens_details", None) if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"): cached_tokens: Final = getattr(prompt_tokens_details, "cached_tokens", None) if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0: @@ -623,9 +623,16 @@ class LangFuseLogger: ) # Apply custom masking function if provided - if masking_function is not None and callable(masking_function): - input = self._apply_masking_function(input, masking_function) - output = self._apply_masking_function(output, masking_function) + masked_input: Final[object] = ( + self._apply_masking_function(input, masking_function) + if masking_function is not None and callable(masking_function) + else input + ) + masked_output: Final[object] = ( + self._apply_masking_function(output, masking_function) + if masking_function is not None and callable(masking_function) + else output + ) clean_metadata = redact_user_api_key_info(metadata=clean_metadata) @@ -651,15 +658,15 @@ class LangFuseLogger: # Special keys that are found in the function arguments and not the metadata if "input" in update_trace_keys: - trace_params["input"] = input if not mask_input else "redacted-by-litellm" + trace_params["input"] = masked_input if not mask_input else "redacted-by-litellm" if "output" in update_trace_keys: - trace_params["output"] = output if not mask_output else "redacted-by-litellm" + trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { "id": trace_id, "name": trace_name, "session_id": session_id, - "input": input if not mask_input else "redacted-by-litellm", + "input": masked_input if not mask_input else "redacted-by-litellm", "version": clean_metadata.pop( "trace_version", clean_metadata.get("version", None) ), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence @@ -669,9 +676,9 @@ class LangFuseLogger: trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None) if level == "ERROR": - trace_params["status_message"] = output + trace_params["status_message"] = masked_output else: - trace_params["output"] = output if not mask_output else "redacted-by-litellm" + trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" if debug is True or (isinstance(debug, str) and debug.lower() == "true"): debug_metadata: Final = { @@ -708,7 +715,7 @@ class LangFuseLogger: ("aws_region_name", aws_region_name, bool(aws_region_name)), ("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs), ) - enrichments: Final[Mapping[str, Any]] = { + enrichments: Final[Mapping[str, object]] = { key: value for key, value, include in candidate_enrichments if include } @@ -802,8 +809,8 @@ class LangFuseLogger: "end_time": end_time, "model": model_name, "model_parameters": optional_params, - "input": input if not mask_input else "redacted-by-litellm", - "output": output if not mask_output else "redacted-by-litellm", + "input": masked_input if not mask_input else "redacted-by-litellm", + "output": masked_output if not mask_output else "redacted-by-litellm", "usage": usage, "usage_details": usage_details, "metadata": { @@ -825,8 +832,8 @@ class LangFuseLogger: prompt_management_metadata=prompt_management_metadata, langfuse_client=self.Langfuse, ) - if output is not None and isinstance(output, str) and level == "ERROR": - generation_params["status_message"] = output + if masked_output is not None and isinstance(masked_output, str) and level == "ERROR": + generation_params["status_message"] = masked_output if self._supports_completion_start_time(): generation_params["completion_start_time"] = kwargs.get("completion_start_time", None) @@ -935,7 +942,7 @@ class LangFuseLogger: return Version(self.langfuse_sdk_version) >= Version("2.7.3") @staticmethod - def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any: + def _apply_masking_function(data: object, masking_function: Callable[[object], object]) -> object: """ Apply a masking function to data, handling different data types. @@ -1049,7 +1056,7 @@ def _add_prompt_to_generation_params( generation_params: dict, clean_metadata: dict, prompt_management_metadata: StandardLoggingPromptManagementMetadata | None, - langfuse_client: Any, + langfuse_client: object, ) -> dict: from langfuse import Langfuse from langfuse.model import ( diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index fae93f03d1e..ce47d7fe27a 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -4,9 +4,12 @@ Opik Logger that logs LLM events to an Opik server import asyncio import traceback +from collections.abc import Mapping from datetime import datetime from typing import Any, Final +from typing_extensions import ReadOnly, TypedDict, Unpack + from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.llms.custom_httpx.http_handler import ( @@ -23,7 +26,7 @@ except Exception: opik_client = None -def _should_skip_event(kwargs: dict[str, Any]) -> bool: +def _should_skip_event(kwargs: Mapping[str, object]) -> bool: """Check if event should be skipped due to missing standard_logging_object.""" if kwargs.get("standard_logging_object") is None: verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found") @@ -31,12 +34,24 @@ def _should_skip_event(kwargs: dict[str, Any]) -> bool: return False +class _OpikLoggerKwargs(TypedDict, total=False): + """Constructor options accepted by ``OpikLogger``.""" + + project_name: ReadOnly[str | None] + url: ReadOnly[str | None] + api_key: ReadOnly[str | None] + workspace: ReadOnly[str | None] + batch_size: ReadOnly[int | None] + flush_interval: ReadOnly[int | None] + max_queue_size: ReadOnly[int | None] + + class OpikLogger(CustomBatchLogger): """ Opik Logger for logging events to an Opik Server """ - def __init__(self, **kwargs: Any) -> None: + def __init__(self, **kwargs: Unpack[_OpikLoggerKwargs]) -> None: self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_httpx_client = _get_httpx_client() @@ -95,7 +110,7 @@ class OpikLogger(CustomBatchLogger): async def async_log_success_event( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], response_obj: Any, start_time: datetime, end_time: datetime, @@ -163,7 +178,7 @@ class OpikLogger(CustomBatchLogger): except Exception as e: verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc()) - def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: + def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None: try: response: Final = self.sync_httpx_client.post( url=url, @@ -178,7 +193,7 @@ class OpikLogger(CustomBatchLogger): def log_success_event( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], response_obj: Any, start_time: datetime, end_time: datetime, @@ -247,7 +262,7 @@ class OpikLogger(CustomBatchLogger): except Exception as e: verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc()) - async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: + async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None: try: response: Final = await self.async_httpx_client.post( url=url, diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index 92a7eca7f3e..4dd3d40fae3 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -1,6 +1,7 @@ """Data extraction functions for Opik payload building.""" import json +from collections.abc import Mapping from typing import Any, Final from litellm import _logging @@ -35,8 +36,8 @@ def normalize_provider_name(provider: str | None) -> str | None: def extract_opik_metadata( - litellm_metadata: dict[str, Any], - standard_logging_metadata: dict[str, Any], + litellm_metadata: Mapping[str, Any], + standard_logging_metadata: Mapping[str, Any], ) -> dict[str, Any]: """ Merge Opik metadata from three sources in increasing priority order: @@ -97,7 +98,7 @@ def extract_span_identifiers( def extract_tags( - opik_metadata: dict[str, Any], + opik_metadata: Mapping[str, Any], custom_llm_provider: str | None, ) -> list[str]: """ @@ -122,7 +123,7 @@ def apply_proxy_header_overrides( project_name: str, tags: list[str], thread_id: str | None, - proxy_headers: dict[str, Any], + proxy_headers: Mapping[str, str], ) -> tuple[str, list[str], str | None]: """ Apply overrides from proxy request headers (opik_* prefix). @@ -148,7 +149,7 @@ def apply_proxy_header_overrides( thread_id = value elif param_key == "tags": try: - parsed_tags = json.loads(value) + parsed_tags: object = json.loads(value) if isinstance(parsed_tags, list): tags.extend(parsed_tags) except (json.JSONDecodeError, TypeError): @@ -158,11 +159,11 @@ def apply_proxy_header_overrides( def extract_and_build_metadata( - opik_metadata: dict[str, Any], - standard_logging_metadata: dict[str, Any], - standard_logging_object: dict[str, Any], - litellm_kwargs: dict[str, Any], -) -> dict[str, Any]: + opik_metadata: Mapping[str, object], + standard_logging_metadata: Mapping[str, object], + standard_logging_object: Mapping[str, object], + litellm_kwargs: Mapping[str, object], +) -> dict[str, object]: """ Build the complete metadata dictionary from all available sources. diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index c7e491c002a..e1623f4697f 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -11,9 +11,10 @@ identical metrics. The attribute cardinality filter is reused from v1 by import from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime -from typing import Any, Final, TypeAlias +from typing import Any, Final, Literal, Protocol, TypeAlias from opentelemetry.metrics import Histogram, Meter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -151,6 +152,29 @@ METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset( BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",) +class _TokenUsage(TypedDict, total=False): + """The token counts a response's ``usage`` carries, as the recorder reads them.""" + + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + + +class _ResponseView(Protocol): + """The one read the recorder makes on a litellm response object.""" + + def get(self, key: Literal["usage"], /) -> _TokenUsage | None: ... + + +class _MetricKwargs(TypedDict, total=False): + """The logging kwargs the recorder reads directly.""" + + call_type: ReadOnly[str | None] + litellm_params: ReadOnly[Mapping[str, object] | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | float | str | None] + api_call_start_time: ReadOnly[datetime | float | str | None] + + def resolve_error_type(kwargs: Mapping[str, Any]) -> str: """The ``error.type`` value for a failed request. @@ -192,8 +216,8 @@ class GenAIMetricRecorder: def record( self, - kwargs: Mapping[str, Any], - response_obj: Any, + kwargs: _MetricKwargs, + response_obj: _ResponseView | None, start_time: datetime, end_time: datetime, ) -> None: @@ -218,7 +242,7 @@ class GenAIMetricRecorder: def record_failure( self, - kwargs: Mapping[str, Any], + kwargs: _MetricKwargs, start_time: datetime, end_time: datetime, ) -> None: @@ -342,7 +366,7 @@ class GenAIMetricRecorder: # Per-metric recording # ------------------------------------------------------------------ # - def _record_token_usage(self, response_obj: Any, common_attrs: dict) -> None: + def _record_token_usage(self, response_obj: _ResponseView | None, common_attrs: dict) -> None: if not response_obj: return usage: Final = response_obj.get("usage") @@ -353,7 +377,7 @@ class GenAIMetricRecorder: self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs) self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs) - def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None: + def _record_time_to_first_token(self, kwargs: _MetricKwargs, common_attrs: dict) -> None: time_to_first_chunk: Final = time_to_first_chunk_seconds(kwargs) if time_to_first_chunk is None: return @@ -361,15 +385,14 @@ class GenAIMetricRecorder: def _record_time_per_output_token( self, - kwargs: Mapping[str, Any], - response_obj: Any, + kwargs: _MetricKwargs, + response_obj: _ResponseView | None, end_time: datetime, duration_s: float, common_attrs: dict, ) -> None: - completion_tokens = None - if response_obj and (usage := response_obj.get("usage")): - completion_tokens = usage.get("completion_tokens") + usage: Final = response_obj.get("usage") if response_obj else None + completion_tokens: Final = usage.get("completion_tokens") if usage else None if completion_tokens is None or completion_tokens <= 0: return 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 aa29162ba1f..07d4f959489 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 @@ -13,7 +13,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.utils import StandardCallbackDynamicParams +from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, VectorStoreResultContent, @@ -226,7 +226,7 @@ class VectorStorePreCallHook(CustomLogger): self, request_data: dict, response: Any, - call_type: Any | None, + call_type: CallTypes | None, ) -> Any | None: """ Add search results to the response after successful LLM call. @@ -283,7 +283,7 @@ class VectorStorePreCallHook(CustomLogger): self, request_data: dict, response_chunk: Any, - call_type: Any | None, + call_type: CallTypes | None, ) -> Any | None: """ Add search results to the final streaming chunk. diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 9125ed6e70a..8479e108d17 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1500,6 +1500,6 @@ class RealTimeStreaming: pass -def client_sent_openai_beta_realtime_header(websocket: Any) -> bool: +def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e2139d688b..3978a01a5db 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -73,6 +73,18 @@ class _ContentChunk(TypedDict): choices: Sequence[_ContentChoice] +class _FunctionCallDelta(TypedDict): + function_call: ReadOnly[FunctionCall] + + +class _FunctionCallChoice(TypedDict): + delta: ReadOnly[_FunctionCallDelta] + + +class _FunctionCallChunk(TypedDict): + choices: ReadOnly[Sequence[_FunctionCallChoice]] + + class _AudioDelta(TypedDict, total=False): audio: ChatCompletionAudioDelta | None @@ -588,7 +600,7 @@ class ChunkProcessor: return tool_calls_list - def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall: + def get_combined_function_call_content(self, function_call_chunks: Sequence["_FunctionCallChunk"]) -> FunctionCall: argument_list: Final = [] delta = function_call_chunks[0]["choices"][0]["delta"] function_call = delta.get("function_call", "") diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 1c5ba951942..f1c7451796d 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -11,8 +11,11 @@ A2A Protocol Format: """ import json +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Optional +from typing_extensions import ReadOnly, TypedDict + from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.types.utils import GenericGuardrailAPIInputs @@ -23,6 +26,13 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth +class _A2ATextPart(TypedDict, total=False): + """The subset of an A2A message part this handler reads text from.""" + + kind: ReadOnly[str] + text: ReadOnly[str] + + class A2AGuardrailHandler(BaseTranslation): """ Handler for processing A2A Protocol messages with guardrails. @@ -41,7 +51,7 @@ class A2AGuardrailHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, - ) -> Any: + ) -> dict: """ Process A2A input messages by applying guardrails to text content. @@ -214,12 +224,12 @@ class A2AGuardrailHandler(BaseTranslation): async def process_output_streaming_response( self, - responses_so_far: list[Any], + responses_so_far: list[object], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, - ) -> list[Any]: + ) -> list[object]: """ Process A2A streaming output by applying guardrails to accumulated text. @@ -305,11 +315,12 @@ class A2AGuardrailHandler(BaseTranslation): def _parse_streaming_responses( self, - responses_so_far: list[Any], - ) -> tuple[list[dict[str, Any] | None], list[tuple[int, dict[str, Any]]]]: + responses_so_far: list[object], + ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" - parsed: Final[list[dict[str, Any] | None]] = [None] * len(responses_so_far) + parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) for i, item in enumerate(responses_so_far): + obj: dict[str, object] if isinstance(item, dict): obj = item elif isinstance(item, str): @@ -326,7 +337,7 @@ class A2AGuardrailHandler(BaseTranslation): def _collect_text_from_parsed_chunks( self, - valid_parsed: list[tuple[int, dict[str, Any]]], + valid_parsed: list[tuple[int, dict[str, object]]], ) -> tuple[str, list[int]]: """Collect text from parsed chunks, returning combined text and indices.""" from litellm.llms.a2a.common_utils import extract_text_from_a2a_response @@ -411,7 +422,7 @@ class A2AGuardrailHandler(BaseTranslation): def _extract_texts_from_parts( self, - parts: list[dict[str, Any]], + parts: Sequence[_A2ATextPart], path: tuple[str, ...], texts_to_check: list[str], task_mappings: list[tuple[tuple[str, ...], int]], diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..057a96ebd49 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, cast import httpx from pydantic import ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -125,7 +126,25 @@ else: _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 -_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[Any], bool]]] = MappingProxyType( + +class _AnthropicUsageIteration(TypedDict, total=False): + """One entry of the ``usage.iterations`` array on an Anthropic response.""" + + input_tokens: ReadOnly[int | None] + output_tokens: ReadOnly[int | None] + cache_creation_input_tokens: ReadOnly[int | None] + cache_read_input_tokens: ReadOnly[int | None] + + +class _AnthropicToolResultBlock(TypedDict, total=False): + """A ``*_tool_result`` content block on an Anthropic response.""" + + type: ReadOnly[str] + tool_use_id: ReadOnly[str] + content: ReadOnly[object] + + +_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[object], bool]]] = MappingProxyType( { "null": lambda v: v is None, "boolean": lambda v: isinstance(v, bool), @@ -440,7 +459,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("speed", None) @staticmethod - def _raise_invalid_reasoning_effort(model: str, value: Any, llm_provider: str) -> NoReturn: + def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn: """Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``. Args: @@ -2059,22 +2078,22 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, completion_response: dict ) -> tuple[ str, - list[Any] | None, + list[object] | None, list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, str | None, list[ChatCompletionToolCallChunk], - list[Any] | None, - list[Any] | None, - list[Any] | None, + list[object] | None, + list[_AnthropicToolResultBlock] | None, + list[object] | None, ]: text_content = "" - citations: list[Any] | None = None + citations: list[object] | None = None thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None reasoning_content: str | None = None tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] - web_search_results: list[Any] | None = None - tool_results: list[Any] | None = None - compaction_blocks: list[Any] | None = None + web_search_results: list[object] | None = None + tool_results: list[_AnthropicToolResultBlock] | None = None + compaction_blocks: list[object] | None = None for idx, content in enumerate(completion_response["content"]): if content["type"] == "text": text_content += content["text"] @@ -2284,7 +2303,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): raw_speed: Final = _usage.get("speed") resolved_speed: Final = raw_speed if isinstance(raw_speed, str) else speed - iterations: Final[list[Any] | None] = _usage.get("iterations") + iterations: Final[Sequence[_AnthropicUsageIteration] | None] = _usage.get("iterations") if iterations: prompt_tokens = sum(it.get("input_tokens", 0) or 0 for it in iterations) completion_tokens = sum(it.get("output_tokens", 0) or 0 for it in iterations) @@ -2377,7 +2396,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _build_code_interpreter_results( self, - tool_results: list[Any], + tool_results: Sequence[_AnthropicToolResultBlock], code_by_id: dict[str, str], container_id: str | None, ) -> list[OutputCodeInterpreterCall]: @@ -2403,11 +2422,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _build_provider_specific_fields( self, completion_response: dict, - citations: list[Any] | None, + citations: Sequence[object] | None, thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, - web_search_results: list[Any] | None, - tool_results: list[Any] | None, - compaction_blocks: list[Any] | None, + web_search_results: Sequence[object] | None, + tool_results: Sequence[_AnthropicToolResultBlock] | None, + compaction_blocks: Sequence[object] | None, tool_calls: list[ChatCompletionToolCallChunk], ) -> dict[str, Any]: provider_specific_fields: Final[dict[str, Any]] = { diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 5fdf2ceff7f..dfd62ca575b 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Coroutine -from typing import Any, Final +from typing import Final import httpx @@ -116,7 +116,7 @@ class AnthropicFilesHandler: api_key: str | None = None, timeout: float | httpx.Timeout = 600.0, max_retries: int | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Retrieve file content from Anthropic. diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 2bcc830851a..46a9dd1a531 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Any, Final +from typing import Final import httpx from openai import ( @@ -374,7 +374,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) error_body: Final = getattr(e, "body", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) @@ -392,7 +392,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): model: str, api_base: str, data: dict, - timeout: Any, + timeout: float | httpx.Timeout, dynamic_params: bool, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, @@ -502,7 +502,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): dynamic_params: bool, data: dict[str, object], model: str, - timeout: Any, + timeout: float | httpx.Timeout, max_retries: int, azure_ad_token: str | None = None, azure_ad_token_provider: Callable | None = None, @@ -578,7 +578,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): dynamic_params: bool, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout, max_retries: int, azure_ad_token: str | None = None, azure_ad_token_provider: Callable | None = None, @@ -634,7 +634,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) message: Final = getattr(e, "message", str(e)) error_body: Final = getattr(e, "body", None) if error_headers is None and error_response: @@ -754,7 +754,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): aembedding=None, headers: dict | None = None, litellm_params: dict | None = None, - ) -> EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse]: + ) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]: if headers: optional_params["extra_headers"] = headers if self._client_session is None: @@ -1268,7 +1268,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers["Authorization"] = f"Bearer {azure_ad_token}" # init AzureOpenAI Client - azure_client_params: Final[dict[str, Any]] = self.initialize_azure_sdk_client( + azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client( litellm_params=litellm_params or {}, api_key=api_key, model_name=model or "", diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index a13b1300e55..f7382190fca 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -51,15 +51,13 @@ else: AsyncHTTPHandler = Any -class _AzureRawAnnotation(TypedDict, total=False): - type: ReadOnly[str] +class _AzureRawAnnotation(ChatCompletionAnnotation, total=False): text: ReadOnly[str] start_index: ReadOnly[int] end_index: ReadOnly[int] - url_citation: ReadOnly[ChatCompletionAnnotationURLCitation] -_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation | _AzureRawAnnotation +_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation class _AzureText(TypedDict, total=False): @@ -223,18 +221,11 @@ class AzureAIAgentsHandler: """Build the ModelResponse from agent output.""" from litellm.types.utils import Choices, Message, Usage - message_kwargs: Final[dict[str, Any]] = { - "content": content, - "role": "assistant", - } - if annotations: - message_kwargs["annotations"] = annotations - model_response.choices = [ Choices( finish_reason="stop", index=0, - message=Message(**message_kwargs), + message=Message(content=content, role="assistant", annotations=annotations or None), ) ] model_response.model = model @@ -655,9 +646,6 @@ class AzureAIAgentsHandler: if data_str == "[DONE]": # Send final chunk with finish_reason - final_delta_kwargs: dict[str, Any] = {"content": None} - if collected_annotations: - final_delta_kwargs["annotations"] = collected_annotations final_chunk = ModelResponseStream( id=response_id, created=created, @@ -667,7 +655,7 @@ class AzureAIAgentsHandler: StreamingChoices( finish_reason="stop", index=0, - delta=Delta(**final_delta_kwargs), + delta=Delta(content=None, annotations=collected_annotations or None), ) ], ) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 28c2e446d10..1f4c81d6491 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -7,7 +7,7 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. import base64 import json import uuid as uuid_lib -from typing import Any, Final, cast +from typing import Final, cast from pydantic import BaseModel @@ -633,7 +633,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final = json.loads(message) + json_message: Final[dict[str, object]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -1182,7 +1182,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Create a function call arguments done event # This is a custom event format that matches what clients expect - function_call_event: Final[dict[str, Any]] = { + function_call_event: Final[dict[str, object]] = { "type": "response.function_call_arguments.done", "event_id": f"event_{uuid.uuid4()}", "response_id": current_response_id, diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index 1ff02a6f8d9..178acb0de0d 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -8,9 +8,11 @@ then we poll until the result is ready. import asyncio import time -from typing import Any, Final +from collections.abc import Coroutine, Mapping +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -33,6 +35,42 @@ from ..common_utils import ( from .transformation import BlackForestLabsImageEditConfig +class _BFLSubmitBody(TypedDict, total=False): + """Decoded body of the BFL submit response, which hands back a polling URL.""" + + errors: ReadOnly[object] + polling_url: ReadOnly[str] + + +class _BFLPollBody(TypedDict, total=False): + """Decoded body of a BFL polling response.""" + + status: ReadOnly[str] + + +class _BFLSubmitResponse(Protocol): + """The submit call's HTTP response, read for its status, body text and decoded body.""" + + @property + def status_code(self) -> int: ... + + @property + def text(self) -> str: ... + + def json(self) -> _BFLSubmitBody: ... + + +class _BFLPollResponse(Protocol): + """A polling call's HTTP response, read only for the task status it carries.""" + + def json(self) -> _BFLPollBody: ... + + +def _poll_status(response: _BFLPollResponse) -> str | None: + """Read the task status out of a BFL polling response body.""" + return response.json().get("status") + + class BlackForestLabsImageEdit: """ Black Forest Labs Image Edit handler. @@ -53,10 +91,10 @@ class BlackForestLabsImageEdit: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, aimage_edit: bool = False, - ) -> ImageResponse | Any: + ) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Main entry point for image edit requests. @@ -185,7 +223,7 @@ class BlackForestLabsImageEdit: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: """ @@ -281,7 +319,7 @@ class BlackForestLabsImageEdit: def _poll_for_result_sync( self, - initial_response: httpx.Response, + initial_response: _BFLSubmitResponse, headers: dict, sync_client: HTTPHandler, max_wait: float = DEFAULT_MAX_POLLING_TIME, @@ -356,8 +394,7 @@ class BlackForestLabsImageEdit: message=f"Polling failed: {response.text}", ) - data = response.json() - status = data.get("status") + status = _poll_status(response) verbose_logger.debug("BFL poll status: %s", status) @@ -383,7 +420,7 @@ class BlackForestLabsImageEdit: async def _poll_for_result_async( self, - initial_response: httpx.Response, + initial_response: _BFLSubmitResponse, headers: dict, async_client: AsyncHTTPHandler, max_wait: float = DEFAULT_MAX_POLLING_TIME, @@ -447,8 +484,7 @@ class BlackForestLabsImageEdit: message=f"Polling failed: {response.text}", ) - data = response.json() - status = data.get("status") + status = _poll_status(response) verbose_logger.debug("BFL poll status: %s", status) diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index 03e4999c5aa..879bef37b58 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -8,9 +8,11 @@ then we poll until the result is ready. import asyncio import time -from typing import Any, Final +from collections.abc import Coroutine, Mapping +from typing import Final, Protocol, TypedDict import httpx +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger @@ -33,6 +35,23 @@ from ..common_utils import ( from .transformation import BlackForestLabsImageGenerationConfig +class _BFLTaskPayload(TypedDict, total=False): + """The body BFL returns for a submitted or polled generation task.""" + + errors: ReadOnly[object] + polling_url: ReadOnly[str] + status: ReadOnly[str] + + +class _TaskJsonResponse(Protocol): + def json(self) -> _BFLTaskPayload: ... + + +def _task_payload(response: _TaskJsonResponse) -> _BFLTaskPayload: + """The JSON body of a BFL task submission or poll response.""" + return response.json() + + class BlackForestLabsImageGeneration: """ Black Forest Labs Image Generation handler. @@ -53,10 +72,10 @@ class BlackForestLabsImageGeneration: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, str] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, aimg_generation: bool = False, - ) -> ImageResponse | Any: + ) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Main entry point for image generation requests. @@ -187,7 +206,7 @@ class BlackForestLabsImageGeneration: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, str] | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: """ @@ -305,7 +324,7 @@ class BlackForestLabsImageGeneration: # Parse initial response to get polling URL try: - response_data: Final = initial_response.json() + response_data: Final = _task_payload(initial_response) except Exception as e: raise BlackForestLabsError( status_code=initial_response.status_code, @@ -350,7 +369,7 @@ class BlackForestLabsImageGeneration: message=f"Polling failed: {response.text}", ) - data = response.json() + data = _task_payload(response) status = data.get("status") verbose_logger.debug("BFL poll status: %s", status) @@ -396,7 +415,7 @@ class BlackForestLabsImageGeneration: # Parse initial response to get polling URL try: - response_data: Final = initial_response.json() + response_data: Final = _task_payload(initial_response) except Exception as e: raise BlackForestLabsError( status_code=initial_response.status_code, @@ -441,7 +460,7 @@ class BlackForestLabsImageGeneration: message=f"Polling failed: {response.text}", ) - data = response.json() + data = _task_payload(response) status = data.get("status") verbose_logger.debug("BFL poll status: %s", status) diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 8c08b2bc33c..f8486d3b274 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -4,9 +4,10 @@ import json from collections.abc import Callable from functools import partial -from typing import Final +from typing import Final, Protocol import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -23,6 +24,53 @@ from litellm.types.utils import TextChoices from litellm.utils import CustomStreamWrapper, TextCompletionResponse +class _CodestralChoiceMessage(TypedDict): + """`choices[].message` of a Codestral FIM completion.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[str | None]] + + +class _CodestralChoice(TypedDict): + """One entry of `choices` in a Codestral FIM completion.""" + + index: ReadOnly[int] + message: ReadOnly[NotRequired[_CodestralChoiceMessage]] + finish_reason: ReadOnly[NotRequired[str | None]] + logprobs: ReadOnly[NotRequired[dict[str, object] | None]] + + +class _CodestralUsage(TypedDict): + """Token accounting returned alongside a Codestral FIM completion.""" + + prompt_tokens: ReadOnly[NotRequired[int]] + completion_tokens: ReadOnly[NotRequired[int]] + total_tokens: ReadOnly[NotRequired[int]] + + +class _CodestralCompletionResponse(TypedDict): + """Body returned by the Codestral `/v1/fim/completions` endpoint.""" + + id: ReadOnly[NotRequired[str]] + created: ReadOnly[NotRequired[int]] + model: ReadOnly[NotRequired[str]] + object: ReadOnly[NotRequired[str]] + usage: ReadOnly[NotRequired[_CodestralUsage]] + choices: ReadOnly[NotRequired[list[_CodestralChoice]]] + + +class _CodestralHTTPResponse(Protocol): + """The Codestral completion response as this handler reads it.""" + + @property + def status_code(self) -> int: ... + + @property + def text(self) -> str: ... + + def json(self) -> _CodestralCompletionResponse: ... + + class TextCompletionCodestralError(Exception): def __init__( self, @@ -115,7 +163,7 @@ class CodestralTextCompletion: def process_text_completion_response( self, model: str, - response: httpx.Response, + response: _CodestralHTTPResponse, model_response: TextCompletionResponse, stream: bool, logging_obj: LiteLLMLogging, diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index e52c56af82b..a3d0482af0a 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -2,10 +2,11 @@ Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -24,6 +25,36 @@ from litellm.types.rerank import ( ) +class _DeepinfraInferenceStatus(TypedDict, total=False): + """The ``inference_status`` block of a DeepInfra rerank response.""" + + status: ReadOnly[str] + runtime_ms: ReadOnly[float] + cost: ReadOnly[float] + tokens_generated: ReadOnly[int] + tokens_input: ReadOnly[int] + + +class _DeepinfraRerankResponse(TypedDict, total=False): + """Body of a DeepInfra ``/rerank`` response.""" + + scores: ReadOnly[Sequence[float]] + input_tokens: ReadOnly[int] + request_id: ReadOnly[str | None] + inference_status: ReadOnly[_DeepinfraInferenceStatus] + + +class _DeepinfraRerankResponseSource(Protocol): + """The DeepInfra ``/rerank`` HTTP response, read for the body it decodes to.""" + + def json(self) -> _DeepinfraRerankResponse: ... + + +def _deepinfra_rerank_body(response: _DeepinfraRerankResponseSource) -> _DeepinfraRerankResponse: + """Decode the body of a DeepInfra ``/rerank`` response.""" + return response.json() + + class DeepinfraRerankConfig(BaseRerankConfig): """ Deepinfra Rerank - Follows the same Spec as Cohere Rerank @@ -95,7 +126,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -150,7 +181,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): litellm_params: dict = {}, ) -> RerankResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _deepinfra_rerank_body(raw_response) logging_obj.post_call(original_response=raw_response.text) # Extract the scores from the response diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index dcd2e4e3471..6d0f211ed7b 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -12,9 +12,10 @@ Schema versioning: litellm.use_legacy_interactions_schema = True. Remove flag after June 8, 2026. """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -41,6 +42,53 @@ else: LiteLLMLoggingObj = Any +_JsonObject: TypeAlias = dict[str, object] + + +class _InteractionPayload(TypedDict, total=False): + """JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields.""" + + id: ReadOnly[str | None] + object: ReadOnly[str | None] + model: ReadOnly[str | None] + agent: ReadOnly[str | None] + status: ReadOnly[str | None] + created: ReadOnly[str | None] + updated: ReadOnly[str | None] + outputs: ReadOnly[list[_JsonObject] | None] + steps: ReadOnly[list[_JsonObject] | None] + usage: ReadOnly[_JsonObject | None] + + +class _CancelPayload(TypedDict, total=False): + """JSON body of an Interactions API cancel response.""" + + id: ReadOnly[str | None] + status: ReadOnly[str | None] + + +class _InteractionPayloadSource(Protocol): + """An Interactions API HTTP response, read for the interaction body it decodes to.""" + + def json(self) -> _InteractionPayload: ... + + +class _CancelPayloadSource(Protocol): + """An Interactions API cancel HTTP response, read for the body it decodes to.""" + + def json(self) -> _CancelPayload: ... + + +def _interaction_body(response: _InteractionPayloadSource) -> _InteractionPayload: + """Decode the body of an Interactions API interaction response.""" + return response.json() + + +def _cancel_body(response: _CancelPayloadSource) -> _CancelPayload: + """Decode the body of an Interactions API cancel response.""" + return response.json() + + class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ Configuration for Google AI Studio Interactions API. @@ -143,7 +191,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ use_legacy: Final[bool] = litellm.use_legacy_interactions_schema - request_body: Final[dict[str, Any]] = {} + request_body: Final[dict[str, object]] = {} # Model or Agent (one required) if model: @@ -189,7 +237,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): and (not isinstance(response_format, dict) or "mime_type" not in response_format) ): # Wrap the legacy schema into the new polymorphic format. - new_rf: Final[dict[str, Any]] = { + new_rf: Final[dict[str, object]] = { "type": "text", "mime_type": response_mime_type, } @@ -215,7 +263,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if image_config is not None: # Move image_config to response_format with type=image. - image_rf: Final[dict[str, Any]] = {"type": "image", **image_config} + image_rf: Final[_JsonObject] = {"type": "image", **image_config} existing_rf: Final = request_body.get("response_format") if existing_rf is None: request_body["response_format"] = image_rf @@ -239,7 +287,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): original_response=raw_response.text, additional_args={"complete_input_dict": {}}, ) - raw_json: Final = raw_response.json() + raw_json: Final = _interaction_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, @@ -290,7 +338,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> InteractionsAPIResponse: try: - raw_json: Final = raw_response.json() + raw_json: Final = _interaction_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, @@ -355,7 +403,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> CancelInteractionResult: try: - raw_json: Final = raw_response.json() + raw_json: Final = _cancel_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 6a1fc144c42..ff4c675b02f 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -1,4 +1,5 @@ import base64 +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -54,8 +55,13 @@ def _convert_image_to_gemini_format(image_file) -> dict[str, str]: return {"bytesBase64Encoded": base64_encoded, "mimeType": mime_type} +def _json_payload(raw_response: httpx.Response) -> object: + """Read an HTTP response body as an opaque JSON payload.""" + return raw_response.json() + + def _usage_video_resolution_from_parameters( - parameters: dict[str, Any], + parameters: Mapping[str, object], ) -> str | None: """Normalize Veo ``parameters.resolution`` for usage and cost tracking.""" res: Final = parameters.get("resolution") @@ -97,7 +103,7 @@ class GeminiVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Map OpenAI-style parameters to Veo format. @@ -111,7 +117,7 @@ class GeminiVideoConfig(BaseVideoConfig): All other params are passed through as-is to support Gemini-specific parameters. """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Get supported OpenAI params (exclude "model" and "prompt" which are handled separately) supported_openai_params: Final = self.get_supported_openai_params(model) @@ -312,11 +318,11 @@ class GeminiVideoConfig(BaseVideoConfig): - status: "processing" - usage: includes duration_seconds and optional video_resolution for cost calculation """ - response_data: Final = raw_response.json() + response_data: Final = _json_payload(raw_response) # Parse response using Pydantic model for type safety try: - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) except Exception as e: raise ValueError(f"Failed to parse operation response: {e}") @@ -336,7 +342,7 @@ class GeminiVideoConfig(BaseVideoConfig): model=model, ) - usage_data: Final[dict[str, Any]] = {} + usage_data: Final[dict[str, float | str]] = {} if request_data: parameters: Final = request_data.get("parameters", {}) duration: Final = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS @@ -367,7 +373,7 @@ class GeminiVideoConfig(BaseVideoConfig): """ operation_name: Final = extract_original_video_id(video_id) url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return url, params @@ -403,9 +409,9 @@ class GeminiVideoConfig(BaseVideoConfig): } } """ - response_data: Final = raw_response.json() + response_data: Final = _json_payload(raw_response) # Parse response using Pydantic model for type safety - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) operation_name: Final = operation_response.name is_done: Final = operation_response.done @@ -443,9 +449,9 @@ class GeminiVideoConfig(BaseVideoConfig): client: Final = litellm.module_level_client status_response: Final = client.get(url=status_url, headers=headers) status_response.raise_for_status() - response_data: Final = status_response.json() + response_data: Final = _json_payload(status_response) - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) if not operation_response.done: raise ValueError( @@ -458,7 +464,7 @@ class GeminiVideoConfig(BaseVideoConfig): generated_samples: Final = operation_response.response.generateVideoResponse.generatedSamples download_url: Final = generated_samples[0].video.uri - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return download_url, params @@ -480,7 +486,7 @@ class GeminiVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Video remix is not supported by Veo API. @@ -506,7 +512,7 @@ class GeminiVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Video list is not supported by Veo API. @@ -547,7 +553,7 @@ class GeminiVideoConfig(BaseVideoConfig): """Video delete is not supported.""" raise NotImplementedError("Video delete is not supported by Google Veo.") - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for Gemini") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index d3db3530109..f6fe7f2fa10 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -1,8 +1,9 @@ import json import os import time +from collections.abc import Sequence from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -24,6 +25,8 @@ from litellm.utils import token_counter from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -31,6 +34,12 @@ else: LoggingClass = Any +class _TokenEncoding(Protocol): + """Tokenizer handle the caller passes in; only `encode` is used, to count completion tokens.""" + + def encode(self, text: str, /) -> Sequence[object]: ... + + tgi_models_cache = None conv_models_cache = None @@ -369,7 +378,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): model_response: ModelResponse, task: hf_tasks | None, optional_params: dict, - encoding: Any, + encoding: "_TokenEncoding | None", messages: list[AllMessageValues], model: str, ): @@ -439,9 +448,10 @@ class HuggingFaceEmbeddingConfig(BaseConfig): if output_text is not None and len(output_text) > 0: completion_tokens = 0 try: - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) - ) ##[TODO] use the llama2 tokenizer here + if encoding is not None: + completion_tokens = len( + encoding.encode(model_response["choices"][0]["message"].get("content", "")) + ) ##[TODO] use the llama2 tokenizer here except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass @@ -469,7 +479,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index d4747b2fb06..3a7f78fd5ba 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -325,7 +325,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] - ) -> Coroutine[Any, Any, list[AllMessageValues]]: + ) -> Coroutine[object, object, list[AllMessageValues]]: ... @overload @@ -341,7 +341,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: bool = False - ) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]: + ) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: """OpenAI no longer supports image_url as a string, so we need to convert it to a dict""" stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages) hoisted_messages: Final = hoist_images_from_tool_messages(stripped_messages) @@ -497,8 +497,12 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): return None tool_call_names: Final = get_tool_call_names(optional_params.get("tools", [])) try: - json_content: Final = json.loads(content) - if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: + json_content: Final[object] = json.loads(content) + if ( + isinstance(json_content, dict) + and json_content.get("type") == "function" + and json_content.get("name") in tool_call_names + ): return ChatCompletionMessageToolCall( function=Function( name=json_content.get("name"), @@ -622,7 +626,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ## RESPONSE OBJECT try: - completion_response: Final = raw_response.json() + completion_response: Final[dict[str, object]] = raw_response.json() except Exception as e: response_headers: Final = getattr(raw_response, "headers", None) raise OpenAIError( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index de15fefe943..ed628f55350 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -51,6 +51,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth class OpenAIChatCompletionsHandler(BaseTranslation): @@ -80,7 +81,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - ) -> Any: + ) -> dict: """ Process input messages by applying guardrails to text content. """ @@ -329,9 +330,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): response: "ModelResponse", guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> Any: + ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -436,7 +437,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, ) -> list["ModelResponseStream"]: @@ -486,7 +487,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", - user_api_key_dict: Any | None, + user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can @@ -589,8 +590,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def build_stream_error_items( self, exc: "HTTPException", - responses_so_far: Sequence[Any] | None = None, - ) -> Sequence[Any] | None: + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes] | None: import json from litellm.proxy.common_request_processing import sse_error_payload @@ -630,7 +631,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", - user_api_key_dict: Any | None, + user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, sink: StreamTransformSink, ) -> None: @@ -794,7 +795,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Determine content source and tool calls based on choice type content = None - tool_calls: list[Any] | None = None + tool_calls: Sequence[object] | None = None if isinstance(choice, litellm.Choices): content = choice.message.content tool_calls = choice.message.tool_calls diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index eadc087383a..09028b6dc5f 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,10 +1,11 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints import httpx from openai.types.responses import ResponseReasoningItem from pydantic import BaseModel, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -37,6 +38,36 @@ _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3 _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +class _DeleteResponseBody(TypedDict): + """Decoded body of the Responses API delete call.""" + + id: ReadOnly[str | None] + object: ReadOnly[str | None] + deleted: ReadOnly[bool | None] + + +class _DeleteResponse(Protocol): + """The delete call's HTTP response, read for the decoded body it carries.""" + + def json(self) -> _DeleteResponseBody: ... + + +class _JsonObjectResponse(Protocol): + """A Responses API HTTP response, read for the JSON object it decodes to.""" + + def json(self) -> dict[str, object]: ... + + +def _delete_response_body(response: _DeleteResponse) -> _DeleteResponseBody: + """Decode a delete response body into the id, object and deleted fields it carries.""" + return response.json() + + +def _json_object_body(response: _JsonObjectResponse) -> dict[str, object]: + """Decode a Responses API response body into its JSON object form.""" + return response.json() + + class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -469,7 +500,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return None @staticmethod - def get_event_model_class(event_type: str) -> Any: + def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ Returns the appropriate event model class based on the event type. @@ -583,7 +614,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the delete response API response into a DeleteResponseResult """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _delete_response_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) return DeleteResponseResult(**raw_response_json) @@ -618,7 +649,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the get response API response into a ResponsesAPIResponse """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers: Final = dict(raw_response.headers) @@ -646,7 +677,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}/input_items" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if before is not None: @@ -665,7 +696,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> dict: try: - return raw_response.json() + return _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) @@ -699,7 +730,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the cancel response API response into a ResponsesAPIResponse """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers: Final = dict(raw_response.headers) diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 8c548b6b0d6..855c49c320b 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -5,10 +5,11 @@ For handling OpenAI-like chat completions, like IBM WatsonX, etc. """ import json -from collections.abc import Callable -from typing import Any, Final +from collections.abc import Callable, Mapping, Sequence +from typing import Final, TypedDict import httpx +from typing_extensions import ReadOnly import litellm from litellm import LlmProviders @@ -25,6 +26,23 @@ from ..common_utils import OpenAILikeBase, OpenAILikeError from .transformation import OpenAILikeChatConfig +class _OpenAILikeChatCompletion(TypedDict, total=False): + """The chat-completion JSON body an OpenAI-like provider returns for a non-streamed call.""" + + id: ReadOnly[str] + choices: ReadOnly[Sequence[Mapping[str, object]]] + created: ReadOnly[int] + model: ReadOnly[str] + system_fingerprint: ReadOnly[str] + usage: ReadOnly[Mapping[str, object]] + object: ReadOnly[str] + + +def _fake_streamed_model_response(payload: _OpenAILikeChatCompletion) -> ModelResponse: + """Build the single response a fake-streamed provider call replays as one chunk.""" + return ModelResponse(**payload) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -42,9 +60,9 @@ async def make_call( response: Final = await client.post(api_base, headers=headers, data=data, stream=not fake_stream) if streaming_decoder is not None: - completion_stream: Any = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) + completion_stream = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = ModelResponse(**response.json()) + model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) @@ -82,7 +100,7 @@ def make_sync_call( if streaming_decoder is not None: completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = ModelResponse(**response.json()) + model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index cde65addb65..5913709c8a0 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -1,8 +1,10 @@ import asyncio import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import ( @@ -29,6 +31,16 @@ else: LiteLLMLoggingObj = Any +class _RunwayMLTask(TypedDict, total=False): + """The RunwayML task payload returned by POST /v1/text_to_image and GET /v1/tasks/{id}.""" + + id: ReadOnly[str] + status: ReadOnly[str] + output: ReadOnly[Sequence[str | Mapping[str, str]]] + failure: ReadOnly[str] + failureCode: ReadOnly[str] + + class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): """ Configuration for RunwayML image generation models. @@ -80,7 +92,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): @staticmethod def _transform_runwayml_response_to_openai( - response_data: dict[str, Any], + response_data: _RunwayMLTask, model_response: ImageResponse, ) -> ImageResponse: """ @@ -155,7 +167,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): raise TimeoutError(f"RunwayML task polling timed out after {timeout_secs} seconds") @staticmethod - def _check_task_status(response_data: dict[str, Any]) -> str: + def _check_task_status(response_data: _RunwayMLTask) -> str: """ Check RunwayML task status from response. @@ -227,7 +239,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayMLTask = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -276,7 +288,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayMLTask = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -322,7 +334,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): } """ try: - response_data = raw_response.json() + response_data: _RunwayMLTask = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -382,7 +394,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): We need to poll the task until it completes (status SUCCEEDED) using async polling. """ try: - response_data = raw_response.json() + response_data: _RunwayMLTask = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index d7743d4d337..a2a93b6114a 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -8,9 +8,10 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path from threading import Lock -from typing import Any, Final +from typing import Any, Final, Protocol import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -33,8 +34,8 @@ def _get_home() -> str: return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH) -def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any: - cur: Any = d +def _get_nested(d: object, path: Sequence[str]) -> object: + cur: object = d if isinstance(cur, str): # This shouldn't happen if service keys are pre-parsed correctly try: @@ -54,7 +55,7 @@ def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any: return cur -def _load_json_env(var_name: str) -> dict[str, Any] | None: +def _load_json_env(var_name: str) -> dict[str, object] | None: raw: Final = os.environ.get(var_name) if not raw: return None @@ -64,7 +65,7 @@ def _load_json_env(var_name: str) -> dict[str, Any] | None: return None -def _str_or_none(value) -> str | None: +def _str_or_none(value: object) -> str | None: try: return str(value) if value is not None else None except Exception: @@ -124,7 +125,7 @@ CREDENTIAL_VALUES: Final[list[CredentialsValue]] = [ ] -def init_conf(profile: str | None = None) -> dict[str, Any]: +def init_conf(profile: str | None = None) -> dict[str, object]: """ Loads config JSON from: 1) $AICORE_CONFIG if set, otherwise @@ -191,7 +192,7 @@ def resolve_resource_group(sources: list[Source]) -> str | None: def _parse_service_key_once( service_key: str | dict | None, -) -> dict[str, Any] | None: +) -> dict[str, object] | None: """ Pre-parse service_key if it's a string to avoid repeated JSON parsing. @@ -348,8 +349,33 @@ def validate_credentials( ) +class _TokenBody(TypedDict): + """Decoded body of the SAP AI Core OAuth2 token response.""" + + access_token: ReadOnly[str] + expires_in: ReadOnly[NotRequired[int]] + + +class _TokenResponse(Protocol): + """The token endpoint's HTTP response, read for the decoded token body it carries.""" + + def json(self) -> _TokenBody: ... + + +def _bearer_token_and_expiry(response: _TokenResponse) -> tuple[str, datetime]: + """Read a token response into the Authorization header value and the token's absolute expiry.""" + payload: Final = response.json() + expires_in: Final = int(payload.get("expires_in", 3600)) + access_token: Final = payload["access_token"] + return f"Bearer {access_token}", datetime.now(timezone.utc) + timedelta(seconds=expires_in) + + def _request_token( - client_id: str, auth_url: str, timeout: float, cert_pair=None, client_secret=None + client_id: str, + auth_url: str, + timeout: float, + cert_pair: tuple[str, str] | None = None, + client_secret: str | None = None, ) -> tuple[str, datetime]: data: Final = {"grant_type": "client_credentials", "client_id": client_id} if client_secret: @@ -361,15 +387,10 @@ def _request_token( with httpx.Client(cert=cert_pair) as raw_client: handler = HTTPHandler(client=raw_client) resp = handler.post(auth_url, data=data, timeout=timeout) - payload = resp.json() - else: - handler = _get_httpx_client() - resp = handler.post(auth_url, data=data, timeout=timeout) - payload = resp.json() - access_token: Final = payload["access_token"] - expires_in: Final = int(payload.get("expires_in", 3600)) - expiry_date: Final = datetime.now(timezone.utc) + timedelta(seconds=expires_in) - return f"Bearer {access_token}", expiry_date + return _bearer_token_and_expiry(resp) + handler = _get_httpx_client() + resp = handler.post(auth_url, data=data, timeout=timeout) + return _bearer_token_and_expiry(resp) except Exception as e: msg: Final = resp.text if resp is not None else getattr(e, "text", str(e)) raise RuntimeError(f"Token request failed: {msg}") from e diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b7f91bfba0d..b6ad9fbcc04 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -12,7 +12,7 @@ from urllib.parse import quote, unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, Required import litellm from litellm._uuid import uuid @@ -104,6 +104,27 @@ class _VertexBatchRow(TypedDict, total=False): processed_time: ReadOnly[str] +class _VertexEmbeddingVector(TypedDict): + values: ReadOnly[list[float]] + + +class _VertexEmbeddingUsageMetadata(TypedDict, total=False): + promptTokenCount: ReadOnly[int] + + +class _VertexEmbeddingResponse(TypedDict, total=False): + embedding: ReadOnly[Required[_VertexEmbeddingVector]] + usageMetadata: ReadOnly[_VertexEmbeddingUsageMetadata] + tokenCount: ReadOnly[int] + + +class _VertexEmbeddingBatchRow(TypedDict, total=False): + key: ReadOnly[str] + request: ReadOnly[Mapping[str, object]] + status: ReadOnly[Required[str]] + response: ReadOnly[Required[_VertexEmbeddingResponse]] + + class _OpenAIBatchOutputError(TypedDict): code: ReadOnly[str] message: ReadOnly[str] @@ -111,7 +132,7 @@ class _OpenAIBatchOutputError(TypedDict): class _OpenAIBatchOutputResponse(TypedDict): status_code: ReadOnly[int] - request_id: ReadOnly[str] + request_id: ReadOnly[object] body: ReadOnly[Mapping[str, object]] @@ -218,7 +239,7 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None return str(labels.get("litellm_custom_id", "unknown")) -def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: """ Whether a Vertex batch output row came from an `EmbedContentRequest`. @@ -237,7 +258,7 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) def _openai_batch_output_row( custom_id: str, - body: Mapping[str, Any] | None = None, + body: Mapping[str, object] | None = None, error_code: str | None = None, error_message: str = "", ) -> _OpenAIBatchOutputRow: @@ -259,7 +280,7 @@ def _openai_batch_output_row( } -def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]: +def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]: """ Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch output row. @@ -278,7 +299,7 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) -def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: +def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> int: """ Prompt tokens billed for one Vertex Gemini Embedding batch row. @@ -293,7 +314,7 @@ def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, - vertex_output_rows: tuple[Mapping[str, Any], ...], + vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...], element_indices: tuple[int, ...], element_count: int, model: str | None, @@ -348,7 +369,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( def _transform_vertex_embeddings_batch_output_to_openai( - vertex_output_rows: Iterable[Mapping[str, Any]], + vertex_output_rows: Iterable[_VertexEmbeddingBatchRow], model: str | None, ) -> tuple[_OpenAIBatchOutputRow, ...]: """ @@ -388,7 +409,7 @@ def _model_from_managed_gcs_url(url: str) -> str | None: return match.group(1) if match else None -def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: +def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: """ Whether an OpenAI batch JSONL line targets the embeddings endpoint. @@ -431,7 +452,7 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" -def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, object]) -> Mapping[str, object]: """ One Vertex Gemini Embedding batch input row. @@ -453,8 +474,8 @@ def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( - openai_entry: Mapping[str, Any], -) -> tuple[Mapping[str, Any], ...]: + openai_entry: Mapping[str, object], +) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding batch rows, one per requested embedding. @@ -512,7 +533,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]], -) -> tuple[Mapping[str, Any], ...]: +) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. @@ -533,7 +554,7 @@ def _openai_batch_jsonl_entry_to_vertex_rows( cached_content=None, ) - custom_id: Final = openai_entry.get("custom_id") + custom_id: Final[object] = openai_entry.get("custom_id") if custom_id is not None: if "labels" not in vertex_request_body: vertex_request_body["labels"] = {} diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 11c026010ee..e2d62be6a69 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -250,7 +250,7 @@ def _gs_uri_requires_content_type_metadata(url: str) -> bool: def _image_url_payload_may_need_sync_gcs_metadata_fetch( - raw_image_url: Any, + raw_image_url: object, ) -> bool: """ True when this image_url value (content-part image_url or assistant ``images[]`` @@ -326,7 +326,7 @@ def _openai_messages_may_need_sync_gcs_metadata_fetch( def _get_gcs_object_content_type( image_url: str, vertex_project: str | None = None, - vertex_credentials: Any | None = None, + vertex_credentials: object = None, ) -> str | None: """ Resolve content type from GCS object metadata. @@ -479,7 +479,7 @@ def _process_gemini_media( model: str | None = None, video_metadata: dict[str, Any] | None = None, vertex_project: str | None = None, - vertex_credentials: Any | None = None, + vertex_credentials: object = None, ) -> PartType: """ Given a media URL (image, audio, or video), return the appropriate PartType for Gemini @@ -1002,7 +1002,7 @@ def _gemini_convert_messages_with_history( if isinstance(_ss_invocations, list): for invocation in _ss_invocations: # Re-inject toolCall part - tc_part: dict[str, Any] = { + tc_part: dict[str, object] = { "toolCall": { "toolType": invocation.get("tool_type"), "id": invocation.get("id"), @@ -1015,13 +1015,13 @@ def _gemini_convert_messages_with_history( # Re-inject toolResponse part if response is present if "response" in invocation: - tr_dict: dict[str, Any] = { + tr_dict: dict[str, object] = { "id": invocation.get("id"), "response": invocation.get("response"), } if invocation.get("tool_type"): tr_dict["toolType"] = invocation["tool_type"] - tr_part: dict[str, Any] = {"toolResponse": tr_dict} + tr_part: dict[str, object] = {"toolResponse": tr_dict} if "response_thought_signature" in invocation: tr_part["thoughtSignature"] = invocation["response_thought_signature"] assistant_content.append(tr_part) @@ -1090,7 +1090,7 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: data_dict[k] = v -def _has_google_maps_tool(tools: Any | None) -> bool: +def _has_google_maps_tool(tools: object) -> bool: """Return True if any tool object in the list has a 'googleMaps' key.""" if not isinstance(tools, list): return False @@ -1127,7 +1127,7 @@ def _rewrite_mime_type_to_response_format(generation_config: GenerationConfig) - schema = generation_config.pop("response_schema", None) generation_config.pop("response_mime_type", None) - response_format: Final[dict[str, Any]] = {"text": {"mimeType": "APPLICATION_JSON"}} + response_format: Final[dict[str, dict[str, object]]] = {"text": {"mimeType": "APPLICATION_JSON"}} if schema is not None: response_format["text"]["schema"] = schema generation_config["responseFormat"] = response_format @@ -1316,7 +1316,7 @@ async def async_transform_request_body( timeout: float | httpx.Timeout | None, extra_headers: dict | None, optional_params: dict, - logging_obj: litellm.litellm_core_utils.litellm_logging.Logging, + logging_obj: LiteLLMLoggingObj, custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], litellm_params: dict, vertex_project: str | None, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index aca257dc095..1942bc850f1 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -9,7 +9,7 @@ import json import os import threading from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse import litellm @@ -47,6 +47,21 @@ else: GoogleCredentialsObject = Any +class _VertexCredentialsObject(Protocol): + """Structural view of the google-auth credentials handle that this class caches and refreshes.""" + + @property + def token(self) -> object: ... + + @property + def quota_project_id(self) -> str | None: ... + + @property + def expired(self) -> object: ... + + def refresh(self, request: object) -> None: ... + + class VertexBase: def __init__(self) -> None: super().__init__() @@ -55,7 +70,7 @@ class VertexBase: self._credentials: GoogleCredentialsObject | None = None self._credentials_project_mapping: dict[ tuple[VERTEX_CREDENTIALS_TYPES | None, str | None], - tuple[GoogleCredentialsObject, str | None], + tuple[_VertexCredentialsObject, str | None], ] = {} self.project_id: str | None = None self.async_handler: AsyncHTTPHandler | None = None @@ -109,7 +124,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - ) -> tuple[Any, str]: + ) -> tuple[_VertexCredentialsObject | None, str]: if credentials is not None: if isinstance(credentials, str): _is_path: Final = os.path.exists( @@ -209,7 +224,7 @@ class VertexBase: return creds, project_id # Google Auth Helpers -- extracted for mocking purposes in tests - def _credentials_from_identity_pool(self, json_obj, scopes): + def _credentials_from_identity_pool(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import identity_pool except ImportError: @@ -220,7 +235,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_pluggable(self, json_obj, scopes): + def _credentials_from_pluggable(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import pluggable except ImportError: @@ -231,7 +246,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_identity_pool_with_aws(self, json_obj, scopes): + def _credentials_from_identity_pool_with_aws(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import aws except ImportError: @@ -242,7 +257,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_authorized_user(self, json_obj, scopes): + def _credentials_from_authorized_user(self, json_obj, scopes) -> _VertexCredentialsObject: try: import google.oauth2.credentials except ImportError: @@ -250,7 +265,7 @@ class VertexBase: return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes) - def _credentials_from_service_account(self, json_obj, scopes): + def _credentials_from_service_account(self, json_obj, scopes) -> _VertexCredentialsObject: try: import google.oauth2.service_account except ImportError: @@ -258,7 +273,7 @@ class VertexBase: return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes) - def _credentials_from_default_auth(self, scopes): + def _credentials_from_default_auth(self, scopes) -> tuple[_VertexCredentialsObject, str | None]: try: import google.auth as google_auth except ImportError: @@ -350,7 +365,7 @@ class VertexBase: ) return api_base - def refresh_auth(self, credentials: Any) -> None: + def refresh_auth(self, credentials: _VertexCredentialsObject) -> None: try: from google.auth.transport.requests import ( Request, @@ -426,7 +441,7 @@ class VertexBase: self, credential_cache_key: tuple, project_id: str | None, - ) -> tuple[str, str, "TokenState", Any, str | None] | None: + ) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None: """ Look up cached credentials and return usable token info for FRESH or STALE tokens (both are still valid for outbound requests). STALE @@ -449,7 +464,9 @@ class VertexBase: return None return creds.token, resolved_project, token_state, creds, cached_project_id - def _unpack_cached_credentials(self, credential_cache_key: tuple) -> tuple[Any, str | None]: + def _unpack_cached_credentials( + self, credential_cache_key: tuple + ) -> tuple[_VertexCredentialsObject | None, str | None]: """ Return (credentials, project_id) from the cache, or (None, None) if not cached. Handles both tuple and legacy cache formats. @@ -461,7 +478,7 @@ class VertexBase: return cached_entry return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) - def _get_token_state(self, credentials: Any) -> "TokenState": + def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState": """ Return the token state using google-auth's TokenState enum. @@ -485,7 +502,7 @@ class VertexBase: credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, credential_cache_key: tuple, - ) -> tuple[Any, str | None]: + ) -> tuple[_VertexCredentialsObject, str | None]: """Load credentials via load_auth (in thread) and cache the result.""" try: _credentials, credential_project_id = await asyncify(self.load_auth)( @@ -505,7 +522,7 @@ class VertexBase: async def _background_refresh_credentials( self, - credentials: Any, + credentials: _VertexCredentialsObject, credential_cache_key: tuple, credential_project_id: str | None, ) -> None: @@ -557,7 +574,7 @@ class VertexBase: def _schedule_background_refresh( self, - credentials: Any, + credentials: _VertexCredentialsObject, credential_cache_key: tuple, credential_project_id: str | None, ) -> None: @@ -575,7 +592,7 @@ class VertexBase: self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id) ) - def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None: + def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None: if self._background_refresh_tasks.get(credential_cache_key) is _fut: self._background_refresh_tasks.pop(credential_cache_key, None) @@ -888,7 +905,7 @@ class VertexBase: # Convert dict credentials to string for caching cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) - _credentials: GoogleCredentialsObject | None = None + _credentials: _VertexCredentialsObject | None = None verbose_logger.debug("Checking cached credentials for project_id: %s", project_id) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9095cee15a9..c4bd03fb1c3 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -6,7 +6,7 @@ from __future__ import annotations import asyncio import contextvars -from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator, Iterator +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator from functools import partial from types import TracebackType from typing import Any, Final, cast @@ -27,19 +27,19 @@ base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils -async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, Any]: +async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]: async for chunk in iterable: yield chunk -def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, Any, Any]: +def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, bytes, None]: yield from iterable -class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): +class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): def __init__( self, - response: Coroutine[Any, Any, httpx.Response], + response: Awaitable[httpx.Response], litellm_logging_obj: LiteLLMLoggingObj, provider_config: BasePassthroughConfig, ) -> None: @@ -48,7 +48,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._headers = httpx.Headers() self._response_coro = response self._response: httpx.Response - self._iterator: AsyncGenerator[bytes, Any] + self._iterator: AsyncGenerator[bytes, bytes] self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks @@ -172,7 +172,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): pass -class PassthroughStreamingResponse(Generator[Any, Any, Any]): +class PassthroughStreamingResponse(Generator[bytes, bytes, None]): def __init__( self, response: httpx.Response, @@ -184,7 +184,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._iterator: Generator[bytes, Any, Any] = _as_generator(response.iter_bytes()) + self._iterator: Generator[bytes, bytes, None] = _as_generator(response.iter_bytes()) self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks self._flush_scheduled = False @@ -263,7 +263,7 @@ async def allm_passthrough_route( cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs, -) -> httpx.Response | AsyncGenerator[Any, Any]: +) -> httpx.Response | AsyncGenerator[bytes, bytes]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -390,10 +390,10 @@ def llm_passthrough_route( **kwargs, ) -> ( httpx.Response - | Coroutine[Any, Any, httpx.Response] - | Coroutine[Any, Any, httpx.Response | AsyncGenerator[Any, Any]] - | Generator[Any, Any, Any] - | AsyncGenerator[Any, Any] + | Coroutine[object, object, httpx.Response] + | Coroutine[object, object, httpx.Response | AsyncGenerator[bytes, bytes]] + | Generator[bytes, bytes, None] + | AsyncGenerator[bytes, bytes] ): """ Pass through requests to the LLM APIs. @@ -592,7 +592,7 @@ async def _async_passthrough_request( is_streaming_request: bool, litellm_logging_obj: LiteLLMLoggingObj, provider_config: BasePassthroughConfig, -) -> httpx.Response | AsyncGenerator[Any, Any]: +) -> httpx.Response | AsyncGenerator[bytes, bytes]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 7ec0f4b5192..dcf1b01bc25 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -5,6 +5,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints. """ import asyncio +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger @@ -74,7 +75,7 @@ class SemanticMCPToolFilter: self.router_instance = litellm_router_instance self.tool_router: SemanticRouter | None = None self.context_window_error: str | None = None - self._tool_map: dict[str, Any] = {} # MCPTool objects or OpenAI function dicts + self._tool_map: dict[str, object] = {} # MCPTool objects or OpenAI function dicts self._index_sync_lock = asyncio.Lock() async def build_router_from_mcp_registry(self) -> None: @@ -182,11 +183,11 @@ class SemanticMCPToolFilter: return raise - def _has_tools_missing_from_index(self, tools: list[Any]) -> bool: + def _has_tools_missing_from_index(self, tools: Sequence[object]) -> bool: """Allocation-free check for any named tool not yet in the semantic index.""" return any(name and name not in self._tool_map for name in (self._extract_tool_info(t)[0] for t in tools)) - def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]: + def _tools_missing_from_index(self, tools: Sequence[object]) -> Mapping[str, object]: """Map name -> tool for every named tool not yet in the semantic index.""" return { name: tool @@ -194,7 +195,7 @@ class SemanticMCPToolFilter: if name and name not in self._tool_map } - async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None: + async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: """ Index request-time tools the startup build never saw. @@ -385,7 +386,7 @@ class SemanticMCPToolFilter: separator: Final = client_name[-len(canonical) - 1] return separator in ("_", "-") - def _get_tools_by_names(self, tool_names: list[str], available_tools: list[Any]) -> list[Any]: + def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[object]) -> list[object]: """ Get tools from available_tools by their names, preserving the semantic router's ordering. @@ -401,14 +402,14 @@ class SemanticMCPToolFilter: # Exact matches win over suffix matches when both are present, and # each incoming tool is returned at most once even if two canonical # names happen to be tail-compatible with the same incoming name. - available_by_name: Final[dict[str, Any]] = {} + available_by_name: Final[dict[str, object]] = {} for tool in available_tools: client_name, _ = self._extract_tool_info(tool) if client_name and client_name not in available_by_name: available_by_name[client_name] = tool - matched: Final[list[Any]] = [] - used_ids: Final[set] = set() + matched: Final[list[object]] = [] + used_ids: Final[set[int]] = set() for canonical in tool_names: tool = available_by_name.get(canonical) if tool is None: @@ -430,7 +431,7 @@ class SemanticMCPToolFilter: used_ids.add(id(tool)) return matched - def extract_user_query(self, messages: list[dict[str, Any]]) -> str: + def extract_user_query(self, messages: Sequence[Mapping[str, object]]) -> str: """ Extract user query from messages for /chat/completions or /responses. diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index bd02cfdf907..31b05320cd3 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -14,7 +14,7 @@ import json from collections.abc import AsyncGenerator, Mapping from copy import deepcopy from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -215,11 +215,20 @@ def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None: ) +class _JsonRpcResponse(Protocol): + def json(self) -> dict[str, object]: ... + + +def _jsonrpc_body(response: _JsonRpcResponse) -> dict[str, object]: + """The decoded JSON-RPC body of ``response``.""" + return response.json() + + async def _forward_jsonrpc( agent_url: str, body: dict[str, object], extra_headers: Mapping[str, str] | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -230,7 +239,7 @@ async def _forward_jsonrpc( ) resp: Final = await handler.post(agent_url, json=body, headers=headers) try: - result: Final = resp.json() + result: Final = _jsonrpc_body(resp) except Exception: resp.raise_for_status() raise @@ -940,8 +949,8 @@ async def invoke_agent_a2a( ) result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers) if method == "agent/getAuthenticatedExtendedCard": - if isinstance(result.get("result"), dict): - card: Final = result["result"] + card: Final = result.get("result") + if isinstance(card, dict): proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}") # Rewrite the upstream agent URL in both 0.3 (top-level `url`) # and 1.0 (`supportedInterfaces[0].url`) wire formats so that diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 39e6ca9a369..0795cee7409 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,8 +14,8 @@ import hashlib import os import re import time -from collections.abc import Awaitable, Callable -from typing import Any, Final, Literal, NoReturn, TypeVar, cast +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast import httpx import jwt @@ -24,6 +24,7 @@ from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from fastapi import HTTPException, status from jwt.api_jwk import PyJWK +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value @@ -93,6 +94,47 @@ UNREACHABLE_CACHE_KEY_PREFIX: Final = "litellm_jwks_unreachable_" _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) +class _JWTAuthSettings(Protocol): + """The JWT auth settings block this handler reads back through ``getattr``, when one is configured.""" + + @property + def issuers(self) -> Sequence[JWTIssuerConfig] | None: ... + + @property + def public_key_ttl(self) -> float: ... + + @property + def public_key_stale_ttl(self) -> float: ... + + +class _OIDCDiscoveryBody(TypedDict, total=False): + """Decoded OIDC discovery document, read for the JWKS endpoint it advertises.""" + + jwks_uri: ReadOnly[str] + + +class _OIDCDiscoveryResponse(Protocol): + """The discovery endpoint's HTTP response, read for the decoded document it carries.""" + + def json(self) -> _OIDCDiscoveryBody: ... + + +class _UserInfoResponse(Protocol): + """The OIDC UserInfo endpoint's HTTP response, read for the identity document it carries.""" + + def json(self) -> dict[str, object]: ... + + +def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody: + """Decode an OIDC discovery response body.""" + return response.json() + + +def _userinfo_document(response: _UserInfoResponse) -> dict[str, object]: + """Decode an OIDC UserInfo response body into its JSON object form.""" + return response.json() + + def jwks_unavailable_exception(error: JWKSUnreachableError) -> ProxyException: return ProxyException( message=( @@ -794,7 +836,7 @@ class JWTHandler: f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}" ) try: - discovery: Final = response.json() + discovery: Final = _discovery_document(response) except Exception as e: raise Exception(f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}") @@ -806,13 +848,13 @@ class JWTHandler: return jwks_uri def _get_public_key_cache_ttl(self) -> float: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return 600 return litellm_jwtauth.public_key_ttl def _get_public_key_stale_ttl(self) -> float: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return DEFAULT_JWKS_STALE_TTL return litellm_jwtauth.public_key_stale_ttl @@ -938,7 +980,7 @@ class JWTHandler: if response.status_code != 200: raise Exception(f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}") - userinfo: Final = response.json() + userinfo: Final = _userinfo_document(response) verbose_proxy_logger.debug("Received OIDC UserInfo: %s", userinfo) # Cache the userinfo response @@ -996,7 +1038,7 @@ class JWTHandler: } def _get_configured_issuer(self, token: str) -> JWTIssuerConfig | None: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return None diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 3a1d18b48cc..554a6ae8d1a 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -6,9 +6,11 @@ import os import sys import tracemalloc from collections import Counter -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, NamedTuple, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query +from typing_extensions import ReadOnly from litellm import get_secret_str from litellm._logging import verbose_proxy_logger @@ -194,6 +196,42 @@ async def memory_usage_in_mem_cache_items( } +class _ProcessMemoryInfo(Protocol): + """The resident and virtual sizes psutil reports for a process.""" + + @property + def rss(self) -> int: ... + + @property + def vms(self) -> int: ... + + +class _ProcessHandle(Protocol): + """The psutil process handle members this module reads.""" + + def memory_info(self) -> _ProcessMemoryInfo: ... + + def memory_percent(self) -> float: ... + + +class _ProcessMemoryUsage(NamedTuple): + """Memory usage of a single worker process.""" + + resident_megabytes: float + virtual_megabytes: float + percent: float + + +def _process_memory_usage(process: _ProcessHandle) -> _ProcessMemoryUsage: + """Read resident/virtual megabytes and system memory share for ``process``.""" + memory_info: Final = process.memory_info() + return _ProcessMemoryUsage( + resident_megabytes=memory_info.rss / (1024 * 1024), + virtual_megabytes=memory_info.vms / (1024 * 1024), + percent=process.memory_percent(), + ) + + @router.get("/debug/memory/summary", include_in_schema=False) async def get_memory_summary( _: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -227,10 +265,9 @@ async def get_memory_summary( try: import psutil - process: Final = psutil.Process() - memory_info: Final = process.memory_info() - memory_mb: Final = memory_info.rss / (1024 * 1024) - memory_percent: Final = process.memory_percent() + usage: Final = _process_memory_usage(psutil.Process()) + memory_mb: Final = usage.resident_megabytes + memory_percent: Final = usage.percent process_memory = { "summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)", @@ -252,7 +289,7 @@ async def get_memory_summary( process_memory["error"] = str(e) # Get cache information - caches: Final[dict[str, Any]] = {} + caches: Final[dict[str, object]] = {} total_cache_items = 0 try: @@ -313,7 +350,7 @@ async def get_memory_summary( } -def _get_gc_statistics() -> dict[str, Any]: +def _get_gc_statistics() -> Mapping[str, object]: """Get garbage collector statistics.""" return { "enabled": gc.isenabled(), @@ -341,30 +378,42 @@ def _get_gc_statistics() -> dict[str, Any]: } -def _get_object_type_counts(top_n: int) -> tuple[int, list[dict[str, Any]]]: +class _ObjectTypeCount(TypedDict): + """One row of the tracked-object histogram.""" + + type: ReadOnly[str] + count: ReadOnly[int] + count_readable: ReadOnly[str] + + +def _type_name_counts(objects: Sequence[object]) -> Counter[str]: + """Count ``objects`` by the name of their type.""" + return Counter(type(obj).__name__ for obj in objects) + + +def _get_object_type_counts(top_n: int) -> tuple[int, list[_ObjectTypeCount]]: """Count objects by type and return total count and top N types.""" - type_counts: Final[Counter] = Counter() - total_objects = 0 + type_counts: Final = _type_name_counts(gc.get_objects()) - for obj in gc.get_objects(): - total_objects += 1 - obj_type = type(obj).__name__ - type_counts[obj_type] += 1 - - top_object_types: Final = [ + top_object_types: Final[list[_ObjectTypeCount]] = [ {"type": obj_type, "count": count, "count_readable": f"{count:,}"} for obj_type, count in type_counts.most_common(top_n) ] - return total_objects, top_object_types + return sum(type_counts.values()), top_object_types -def _get_uncollectable_objects_info() -> dict[str, Any]: +def _type_names(objects: Sequence[object]) -> Sequence[str]: + """The type name of each object in ``objects``.""" + return [type(obj).__name__ for obj in objects] + + +def _get_uncollectable_objects_info() -> Mapping[str, object]: """Get information about uncollectable objects (potential memory leaks).""" uncollectable: Final = gc.garbage return { "count": len(uncollectable), - "sample_types": [type(obj).__name__ for obj in uncollectable[:10]], + "sample_types": _type_names(uncollectable[:10]), "warning": ( "If count > 0, you may have reference cycles preventing garbage collection" if len(uncollectable) > 0 @@ -373,9 +422,11 @@ def _get_uncollectable_objects_info() -> dict[str, Any]: } -def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) -> dict[str, Any]: +def _get_cache_memory_stats( + user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache +) -> Mapping[str, object]: """Calculate memory usage for all caches.""" - cache_stats: Final[dict[str, Any]] = {} + cache_stats: Final[dict[str, object]] = {} try: # User API key cache user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) @@ -439,9 +490,9 @@ def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, r return cache_stats -def _get_router_memory_stats(llm_router) -> dict[str, Any]: +def _get_router_memory_stats(llm_router) -> Mapping[str, object]: """Get memory usage statistics for LiteLLM router.""" - litellm_router_memory: dict[str, Any] = {} + litellm_router_memory: dict[str, object] = {} try: if llm_router is not None: # Model list memory size @@ -505,7 +556,7 @@ def _get_router_memory_stats(llm_router) -> dict[str, Any]: return litellm_router_memory -def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> dict[str, Any] | None: +def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Mapping[str, object] | None: """Get process-level memory information using psutil.""" if not include_process_info: return None @@ -514,10 +565,10 @@ def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> dic import psutil process: Final = psutil.Process() - memory_info: Final = process.memory_info() - ram_usage_mb: Final = round(memory_info.rss / (1024 * 1024), 2) - virtual_memory_mb: Final = round(memory_info.vms / (1024 * 1024), 2) - memory_percent: Final = round(process.memory_percent(), 2) + usage: Final = _process_memory_usage(process) + ram_usage_mb: Final = round(usage.resident_megabytes, 2) + virtual_memory_mb: Final = round(usage.virtual_megabytes, 2) + memory_percent: Final = round(usage.percent, 2) return { "pid": worker_pid, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 202a95ba29b..e6880d521f1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -211,7 +211,7 @@ class DBSpendUpdateWriter: org_id: str | None, # Completion object fields kwargs: dict | None, - completion_response: litellm.ModelResponse | Any | Exception | None, + completion_response: object, start_time: datetime | None, end_time: datetime | None, response_cost: float | None, @@ -323,7 +323,7 @@ class DBSpendUpdateWriter: async def _enqueue_tool_usage_transaction( self, payload: SpendLogsPayload, - completion_response: "litellm.ModelResponse | Any | Exception | None", + completion_response: object, prisma_client: "PrismaClient | None", kwargs: "dict | None" = None, ) -> None: @@ -396,7 +396,7 @@ class DBSpendUpdateWriter: def _enqueue_tool_registry_upsert( self, kwargs: dict | None, - completion_response: Any | None, + completion_response: object, hashed_token: str | None = None, team_id: str | None = None, ) -> None: @@ -849,7 +849,7 @@ class DBSpendUpdateWriter: return # Parse tags from JSON string - tags = [] + tags: Sequence[object] = [] if isinstance(request_tags, str): tags = safe_json_loads(request_tags, default=[]) if not tags: @@ -2260,7 +2260,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.") return - request_tags = [] + request_tags: Sequence[str] = [] if isinstance(payload["request_tags"], str): request_tags = json.loads(payload["request_tags"]) elif isinstance(payload["request_tags"], list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py index 3716d00774f..2c27531cea1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py @@ -162,10 +162,10 @@ class AktoGuardrail(CustomGuardrail): def build_request_body( inputs: GenericGuardrailAPIInputs, request_data: dict | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the LLM request body from guardrail inputs (messages, model, tools).""" model: Final = inputs.get("model", "") or "" - body: Final[dict[str, Any]] = {"model": model} + body: Final[dict[str, object]] = {"model": model} structured: Final = inputs.get("structured_messages") if structured: @@ -194,7 +194,7 @@ class AktoGuardrail(CustomGuardrail): def build_response_body( inputs: GenericGuardrailAPIInputs, request_data: dict | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the LLM response body, preferring the actual model response if available.""" model_response: Final = request_data.get("response") if request_data else None if model_response is not None and hasattr(model_response, "model_dump"): @@ -224,7 +224,7 @@ class AktoGuardrail(CustomGuardrail): *, status_code: int = 200, include_response: bool = False, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the flat MIRRORING payload sent to Akto's HTTP proxy endpoint. All body fields use double-encoding: json.dumps({"body": json.dumps(actual_body)}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 955a868a0d6..48832f8ed5e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -2,9 +2,10 @@ import os import time -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -27,6 +28,33 @@ if TYPE_CHECKING: GRAYSWAN_BLOCK_ERROR_MSG: Final = "Blocked by Gray Swan Guardrail" +class _GraySwanMonitorResponse(TypedDict): + """Body returned by Gray Swan's `/cygnal/monitor` endpoint.""" + + violation: ReadOnly[NotRequired[float | None]] + violated_rules: ReadOnly[NotRequired[list[object]]] + violated_rule_descriptions: ReadOnly[NotRequired[list[object]]] + mutation: ReadOnly[NotRequired[bool | None]] + ipi: ReadOnly[NotRequired[bool | None]] + + +class _GraySwanMonitorHTTPResponse(Protocol): + def raise_for_status(self) -> object: ... + + def json(self) -> _GraySwanMonitorResponse: ... + + +class _GraySwanMonitorHTTPClient(Protocol): + async def post( + self, + *, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> _GraySwanMonitorHTTPResponse: ... + + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -77,7 +105,9 @@ class GraySwanGuardrail(CustomGuardrail): guardrail_timeout: float | None = 30.0, **kwargs: Any, ) -> None: - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler: _GraySwanMonitorHTTPClient = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) api_key_value: Final = api_key or os.getenv("GRAYSWAN_API_KEY") if not api_key_value: @@ -266,7 +296,7 @@ class GraySwanGuardrail(CustomGuardrail): # Legacy Test Interface (for backward compatibility) # ------------------------------------------------------------------ - async def run_grayswan_guardrail(self, payload: dict) -> dict[str, Any]: + async def run_grayswan_guardrail(self, payload: dict[str, object]) -> _GraySwanMonitorResponse: """ Run the GraySwan guardrail on a payload. @@ -285,7 +315,7 @@ class GraySwanGuardrail(CustomGuardrail): def _process_grayswan_response( self, - response_json: dict, + response_json: _GraySwanMonitorResponse, data: dict | None = None, hook_type: GuardrailEventHooks | None = None, ) -> None: @@ -385,7 +415,7 @@ class GraySwanGuardrail(CustomGuardrail): # Core GraySwan API interaction # ------------------------------------------------------------------ - async def _call_grayswan_api(self, payload: dict) -> dict[str, Any]: + async def _call_grayswan_api(self, payload: dict[str, object]) -> _GraySwanMonitorResponse: """Call the GraySwan monitoring API.""" headers: Final = self._prepare_headers() @@ -406,7 +436,7 @@ class GraySwanGuardrail(CustomGuardrail): def _process_response_internal( self, - response_json: dict[str, Any], + response_json: _GraySwanMonitorResponse, request_data: dict, inputs: GenericGuardrailAPIInputs, is_output: bool, @@ -534,8 +564,8 @@ class GraySwanGuardrail(CustomGuardrail): dynamic_body: dict, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"] = None, - ) -> dict[str, Any] | None: - payload: Final[dict[str, Any]] = {"messages": messages} + ) -> dict[str, object] | None: + payload: Final[dict[str, object]] = {"messages": messages} categories: Final = dynamic_body.get("categories") or self.categories if categories: @@ -563,13 +593,13 @@ class GraySwanGuardrail(CustomGuardrail): {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers ) if cleaned_litellm_metadata: - sanitized: Final = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized return payload - def _format_violation_message(self, detection_info: Any, is_output: bool = False) -> str: + def _format_violation_message(self, detection_info: object, is_output: bool = False) -> str: """ Format detection info into a user-friendly violation message. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index ea022510309..cf5da27e9ca 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -8,6 +8,7 @@ import json import os import uuid +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict try: @@ -128,7 +129,7 @@ class LassoGuardrail(CustomGuardrail): @staticmethod def _extract_tool_call_fields( - call: Any, + call: object, ) -> tuple[str | None, str | None, dict[str, object] | None]: """Extract (call_id, name, parsed_input) from a tool call. @@ -476,7 +477,7 @@ class LassoGuardrail(CustomGuardrail): def _map_masked_messages_back( self, original_messages: list[dict[str, Any]], - masked_messages: list[dict[str, Any]], + masked_messages: Sequence[Mapping[str, object]], ) -> list[dict[str, object]]: """Map Lasso-format masked messages back onto the original OpenAI-format messages. @@ -638,7 +639,7 @@ class LassoGuardrail(CustomGuardrail): }, ) - def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, object]]: """ Convert raw OpenAI-format messages to Lasso API format with content blocks. @@ -646,7 +647,7 @@ class LassoGuardrail(CustomGuardrail): - role=tool messages → developer role + tool_result block - plain text messages pass through unchanged """ - expanded: Final[list[dict[str, Any]]] = [] + expanded: Final[list[dict[str, object]]] = [] for msg in messages: role = msg.get("role", "") content = msg.get("content") @@ -917,7 +918,7 @@ class LassoGuardrail(CustomGuardrail): def _apply_masking_to_model_response( self, model_response: litellm.ModelResponse, - masked_messages: list[dict[str, Any]], + masked_messages: Sequence[Mapping[str, object]], ) -> None: """Apply masking to the actual model response when mask=True and masked content is available.""" # Index masked tool_use blocks by id for O(1) lookup. diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 78639ce4fd0..7021d41475b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -8,11 +8,12 @@ # Standard library imports import json import os -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import quote # Third-party imports from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict # LiteLLM imports from litellm import DualCache @@ -42,7 +43,34 @@ if TYPE_CHECKING: MAX_PILLAR_HEADER_VALUE_BYTES: Final = 8 * 1024 -def _encode_json_for_header(data: Any) -> str: +class _PillarProtectResponse(TypedDict): + """Body returned by Pillar's `/api/v1/protect` endpoint.""" + + flagged: ReadOnly[NotRequired[bool]] + session_id: ReadOnly[NotRequired[str]] + scanners: ReadOnly[NotRequired[dict[str, object]]] + evidence: ReadOnly[NotRequired[list[object]]] + masked_session_messages: ReadOnly[NotRequired[list[object]]] + + +class _PillarProtectHTTPResponse(Protocol): + def raise_for_status(self) -> object: ... + + def json(self) -> _PillarProtectResponse: ... + + +class _PillarProtectHTTPClient(Protocol): + async def post( + self, + *, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> _PillarProtectHTTPResponse: ... + + +def _encode_json_for_header(data: object) -> str: """ JSON-serialize and URL-encode data for safe header transmission. """ @@ -50,7 +78,9 @@ def _encode_json_for_header(data: Any) -> str: return quote(json_payload, safe="") -def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES) -> tuple[Any, str, bool]: +def _truncate_evidence_payload( + evidence: object, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES +) -> tuple[object, str, bool]: """ Truncate evidence payload so the encoded header value stays within max_bytes. @@ -66,12 +96,12 @@ def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER truncated_value: Final = "[truncated]" return truncated_value, _encode_json_for_header(truncated_value), True - truncated: Final[list[Any]] = [] + truncated: Final[list[object]] = [] encoded = _encode_json_for_header(truncated) truncated_flag = False for entry in evidence: - working_entry: Any + working_entry: object if isinstance(entry, dict): working_entry = dict(entry) else: @@ -105,7 +135,7 @@ def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER return truncated, encoded, truncated_flag -def build_pillar_response_headers(metadata_store: dict[str, Any]) -> dict[str, str]: +def build_pillar_response_headers(metadata_store: dict[str, object]) -> dict[str, str]: """ Create URL-safe Pillar response headers and apply truncation metadata. """ @@ -191,7 +221,9 @@ class PillarGuardrail(CustomGuardrail): LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler: _PillarProtectHTTPClient = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") if self.api_key is None: @@ -686,7 +718,7 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> dict[str, Any]: + async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> _PillarProtectResponse: """ Call the Pillar API and return the response. @@ -714,7 +746,7 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Analysis complete - flagged=%s, session=%s", flagged, session_id) return res - def _process_pillar_response(self, pillar_response: dict[str, Any], original_data: dict) -> None: + def _process_pillar_response(self, pillar_response: _PillarProtectResponse, original_data: dict) -> None: """ Process the Pillar API response and handle detections based on configuration. @@ -774,7 +806,7 @@ class PillarGuardrail(CustomGuardrail): build_pillar_response_headers(metadata_store) - def _raise_pillar_detection_exception(self, pillar_response: dict[str, Any]) -> None: + def _raise_pillar_detection_exception(self, pillar_response: _PillarProtectResponse) -> None: """ Raise an HTTPException for Pillar security detections. @@ -784,7 +816,7 @@ class PillarGuardrail(CustomGuardrail): Raises: HTTPException: Always raises with security detection details """ - pillar_response_dict: Final = { + pillar_response_dict: Final[dict[str, object]] = { "session_id": pillar_response.get("session_id"), } diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index e34beec4d3e..2fbd50b5863 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -6,7 +6,7 @@ via embedding similarity. Smarter than regex (understands intent), lighter than an LLM call (~20-50ms per request for embedding). """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_guardrail import ( @@ -50,7 +50,7 @@ class SemanticGuardrail(CustomGuardrail): similarity_threshold: float, route_templates: list[str] | None = None, custom_routes_file: str | None = None, - custom_routes: list[dict[str, Any]] | None = None, + custom_routes: list[dict[str, object]] | None = None, on_flagged_action: str = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, default_on: bool = False, @@ -157,7 +157,14 @@ class SemanticGuardrail(CustomGuardrail): return response -def _get_top_route_choice(result: Any) -> Any: +class _RouteChoice(Protocol): + """The semantic-router match this guardrail reads: the route that fired, if any.""" + + @property + def name(self) -> str | None: ... + + +def _get_top_route_choice(result: _RouteChoice | list[_RouteChoice] | None) -> _RouteChoice | None: """Extract the top RouteChoice from SemanticRouter result. SemanticRouter.__call__ can return RouteChoice or List[RouteChoice]. @@ -194,7 +201,7 @@ def _extract_response_text(response: Any) -> str: return "" -def _content_to_text(content: Any) -> str: +def _content_to_text(content: object) -> str: if isinstance(content, str): return content if isinstance(content, list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 3c5625bc272..a8b33109900 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,9 +1,10 @@ import json import re from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict from fastapi import HTTPException +from typing_extensions import ReadOnly, Required from litellm import ChatCompletionToolParam from litellm._logging import verbose_proxy_logger @@ -51,6 +52,27 @@ def _object_list(value: object) -> Sequence[object] | None: return value if isinstance(value, list) else None +class _ToolPermissionRuleFields(TypedDict, total=False): + """The config-file shape a :class:`ToolPermissionRule` is built from.""" + + id: ReadOnly[Required[str]] + tool_name: ReadOnly[str | None] + tool_type: ReadOnly[str | None] + decision: ReadOnly[Required[Literal["allow", "deny"]]] + allowed_param_patterns: ReadOnly[dict[str, str] | None] + + +def _rule_from_fields(fields: _ToolPermissionRuleFields) -> ToolPermissionRule: + """Validate one config-file rule entry into a :class:`ToolPermissionRule`.""" + return ToolPermissionRule(**fields) + + +def _is_tool_use_block(block: object) -> bool: + """Whether ``block`` is an Anthropic ``tool_use`` content block.""" + fields: Final = _object_mapping(block) + return fields is not None and fields.get("type") == "tool_use" + + class ToolPermissionGuardrail(CustomGuardrail): def __init__( self, @@ -101,7 +123,7 @@ class ToolPermissionGuardrail(CustomGuardrail): compiled_patterns: Final[dict[str, dict[str, re.Pattern]]] = {} for rule_item in rules or []: - rule = rule_item if isinstance(rule_item, ToolPermissionRule) else ToolPermissionRule(**rule_item) + rule = rule_item if isinstance(rule_item, ToolPermissionRule) else _rule_from_fields(rule_item) target_patterns: dict[str, re.Pattern | None] = { "tool_name": None, @@ -440,7 +462,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return is_allowed, None, message @staticmethod - def _get_mapping_value(item: Any, key: str) -> Any: + def _get_mapping_value(item: object, key: str) -> Any: if isinstance(item, dict): return item.get(key) return getattr(item, key, None) @@ -450,7 +472,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return f"legacy_function_call_{choice_index}" def _legacy_function_call_to_tool_call( - self, function_call: Any, choice_index: int + self, function_call: object, choice_index: int ) -> ChatCompletionMessageToolCall | None: if function_call is None: return None @@ -549,7 +571,7 @@ class ToolPermissionGuardrail(CustomGuardrail): def _modify_anthropic_content_with_permission_errors( self, response: object, - content: tuple[Any, ...], + content: tuple[object, ...], denied_tools: tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...], ) -> None: if not denied_tools or not isinstance(response, dict): @@ -557,27 +579,33 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tools)) - error_by_tool_use_id: Final = { # mutable-ok: read-only lookup, never mutated after construction + error_by_tool_use_id: Final[ + Mapping[object, str] + ] = { # mutable-ok: read-only lookup, never mutated after construction tool_call.id: self._create_permission_error_result(tool_call, error).content for tool_call, error in denied_tools } - denied_block_ids: Final = frozenset(error_by_tool_use_id) - def _is_denied(block: object) -> bool: - return isinstance(block, dict) and block.get("type") == "tool_use" and block.get("id") in denied_block_ids + def _denied_message(block: object) -> str | None: + fields: Final = _object_mapping(block) + if fields is None or fields.get("type") != "tool_use": + return None + return error_by_tool_use_id.get(fields.get("id")) - error_messages: Final = tuple(error_by_tool_use_id[block["id"]] for block in content if _is_denied(block)) - kept_blocks: Final = tuple(block for block in content if not _is_denied(block)) + error_messages: Final = tuple( + message for message in (_denied_message(block) for block in content) if message is not None + ) + kept_blocks: Final = tuple(block for block in content if _denied_message(block) is None) new_content: Final = [ # mutable-ok: response content is a JSON array on the wire *kept_blocks, {"type": "text", "text": "\n".join(error_messages)}, # mutable-ok: content block is a JSON object ] response["content"] = new_content # rebind-ok: the guardrail rewrites the provider response in place - if not any(isinstance(block, dict) and block.get("type") == "tool_use" for block in kept_blocks): + if not any(_is_tool_use_block(block) for block in kept_blocks): response["stop_reason"] = "end_turn" # rebind-ok: dropping every tool_use ends the turn - def _get_request_tool_name(self, tool: Any) -> tuple[str | None, str | None]: + def _get_request_tool_name(self, tool: object) -> tuple[str | None, str | None]: tool_type: Final = self._get_mapping_value(tool, "type") if tool_type != "function": return None, tool_type @@ -586,7 +614,7 @@ class ToolPermissionGuardrail(CustomGuardrail): tool_name: Final = self._get_mapping_value(function, "name") return tool_name, tool_type - def _get_legacy_function_name(self, function: Any) -> str | None: + def _get_legacy_function_name(self, function: object) -> str | None: return self._get_mapping_value(function, "name") def _get_named_tool_choice(self, data: dict) -> str | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index 6b8148645aa..a5945a39589 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -433,7 +433,7 @@ class VigilGuardGuardrail(CustomGuardrail): return collected @staticmethod - def _clamp_metadata_value(value: Any) -> _MetadataValue | None: + def _clamp_metadata_value(value: object) -> _MetadataValue | None: if isinstance(value, bool): return None if isinstance(value, str): diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 569ec32c1a0..9edbc6dbf1c 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -67,6 +67,24 @@ class _ChatMessage(Protocol): def tool_calls(self) -> Sequence[_ChatToolCall] | None: ... +class _ChatChoice(Protocol): + @property + def message(self) -> _ChatMessage: ... + + @property + def finish_reason(self) -> str | None: ... + + +class _ChatCompletion(Protocol): + @property + def choices(self) -> Sequence[_ChatChoice]: ... + + +def _first_choice(response: _ChatCompletion) -> _ChatChoice: + """The first choice of an OpenAI shaped completion response.""" + return response.choices[0] + + class SkillsInjectionHook(CustomLogger): """ Pre/Post-call hook that processes skills from container.skills parameter. @@ -738,8 +756,9 @@ print('No executable skill module found') for iteration in range(self.max_iterations): # OpenAI format response has choices[0].message - assistant_message: _ChatMessage = current_response.choices[0].message - stop_reason: str | None = current_response.choices[0].finish_reason + choice: _ChatChoice = _first_choice(current_response) + assistant_message: _ChatMessage = choice.message + stop_reason: str | None = choice.finish_reason # Build assistant message for conversation history assistant_msg_dict: dict[str, object] = { diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1e65da5b867..63129602082 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -8,7 +8,7 @@ import asyncio import binascii import os import uuid -from collections.abc import Callable, Mapping, Sequence, Set +from collections.abc import Awaitable, Callable, Mapping, Sequence, Set from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime @@ -386,6 +386,12 @@ CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None] ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes +class _AsyncLuaScript(Protocol): + """A Lua script registered against the async Redis client, called with KEYS and ARGV.""" + + def __call__(self, *, keys: Sequence[str], args: Sequence[object]) -> Awaitable[list[CacheCounterValue]]: ... + + class RateLimitDescriptorRateLimitObject(TypedDict, total=False): requests_per_unit: int | None tokens_per_unit: int | None @@ -577,6 +583,14 @@ def _parse_output_cap_value(raw_value: object) -> int | None: class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): + batch_rate_limiter_script: _AsyncLuaScript | None + token_increment_script: _AsyncLuaScript | None + check_and_increment_by_n_script: _AsyncLuaScript | None + window_guarded_token_increment_script: _AsyncLuaScript | None + parallel_acquire_script: _AsyncLuaScript | None + parallel_release_script: _AsyncLuaScript | None + parallel_count_script: _AsyncLuaScript | None + def __init__( self, internal_usage_cache: InternalUsageCache, @@ -3855,7 +3869,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): expected_window_start = operation.get("expected_window_start") if window_key is None or expected_window_start is None: continue - active_window_start = await self.internal_usage_cache.async_get_cache( + active_window_start: CacheCounterValue | None = await self.internal_usage_cache.async_get_cache( key=window_key, litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -4144,7 +4158,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_tpm_scope_targets( self, standard_logging_metadata: dict[str, Any], - kwargs: Any, + kwargs: object, model_group: str | None, ) -> list[tuple[str, str]]: """ @@ -4301,8 +4315,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_success_event_pipeline_operations( self, - kwargs: Any, - response_obj: Any, + kwargs: dict[str, Any], + response_obj: object, rate_limit_type: Literal["output", "input", "total"], ) -> list[RedisPipelineIncrementOperation]: """Build Redis pipeline increment ops for TPM / parallel-request counters.""" diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 012aec38458..14d2332a7eb 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -543,7 +543,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) - merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) + merged_model_info: Final[dict[str, object]] = db_model.model_info.model_dump(exclude_none=True) # update litellm params if updated_patch.litellm_params: @@ -1982,7 +1982,7 @@ async def update_model( ### MERGE WITH EXISTING DATA ### merged_dictionary: Final = {} - _mp: Final = model_params.litellm_params.dict() + _mp: Final[dict[str, object]] = model_params.litellm_params.dict() for key, value in _mp.items(): if value is not None: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 9198aa35f3f..5e38a016099 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -487,12 +487,11 @@ async def new_organization( for m in data.models: await can_user_call_model(m, llm_router=llm_router, user_object=user_object_correct_type) - organization_row: Final = LiteLLM_OrganizationTable( - **data.json(exclude_none=True), - object_permission_id=object_permission_id, - created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, - updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, - ) + organization_payload: Final = _STR_OBJECT_DICT_ADAPTER.validate_python(data.json(exclude_none=True)) + organization_payload["object_permission_id"] = object_permission_id + organization_payload["created_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + organization_payload["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + organization_row: Final = LiteLLM_OrganizationTable.model_validate(organization_payload) for field in LiteLLM_ManagementEndpoint_MetadataFields: if getattr(data, field, None) is not None: @@ -644,7 +643,7 @@ async def update_organization( ) # Transform UI payload to expected format - raw_data: Final = await request.json() + raw_data: Final[dict[str, object]] = await request.json() raw_data_with_flat_budget_fields: Final = handle_nested_budget_structure_in_organization_update_request(raw_data) # Create validated data model @@ -691,7 +690,7 @@ async def update_organization( # Merge metadata from existing organization with updated metadata if updated_organization_row_json.get("metadata") is not None: existing_metadata: Final = existing_organization_row.metadata or {} - updated_metadata: Final = updated_organization_row_json.get("metadata", {}) + updated_metadata: Final[dict[str, object]] = updated_organization_row_json.get("metadata", {}) merged_metadata: Final[Mapping[str, object]] = _update_dictionary( existing_dict=cast( # cast-ok: prisma de-serializes a Json column to the plain python dict it stores "dict[str, object]", existing_metadata diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 613508da22b..606569c5b8b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -502,7 +502,7 @@ def _set_nested_metadata_value(metadata: dict[str, object], key_path: str, value placeholder: Final = "\x00" parts = key_path.replace("\\.", placeholder).split(".") parts = [p.replace(placeholder, ".") for p in parts] - current: Any = metadata + current: dict[str, object] = metadata for part in parts[:-1]: existing = current.get(part) if not isinstance(existing, dict): @@ -4076,7 +4076,7 @@ class SSOAuthenticationHandler: ) if resp.status_code == 200: try: - userinfo_raw: Final = resp.json() + userinfo_raw: Final[dict[str, object] | None] = resp.json() if not userinfo_raw: # JSON null (None) or empty dict ({}) — no identity claims. # Treat as failure so id_token fallback can be attempted. @@ -4406,7 +4406,7 @@ class MicrosoftSSOHandler: ) -> tuple[list[str], str | None]: """Helper function to fetch and parse group data from a URL""" response: Final = await async_client.get(url, headers=headers) - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() response_typed: Final = await MicrosoftSSOHandler._cast_graph_api_response_dict(response=response_json) group_ids: Final = MicrosoftSSOHandler._get_group_ids_from_graph_api_response(response=response_typed) return group_ids, response_typed.get("odata_nextLink") diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index ee9a5d94440..49ec18013b5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -267,7 +267,7 @@ class VertexPassthroughLoggingHandler: model: Final = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) - _json_response: Final = httpx_response.json() + _json_response: Final[dict[str, object]] = httpx_response.json() litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse() if vertex_image_generation_class.is_image_generation_response(_json_response): @@ -422,7 +422,7 @@ class VertexPassthroughLoggingHandler: - Creates standard logging object - Logs in litellm callbacks """ - kwargs: dict[str, Any] = {} + kwargs: dict[str, object] = {} vertex_location: Final = get_vertex_location_from_url(url_route) if vertex_location is not None: litellm_logging_obj.optional_params["vertex_location"] = vertex_location diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index aa7595ed13d..5907ffc64eb 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -52,7 +52,7 @@ _TOOL_PAYLOAD_KEYS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( "function": ("name", "description", "parameters", "strict"), } ) -_EMPTY_TOOL_PAYLOAD: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_TOOL_PAYLOAD: Final[Mapping[str, object]] = MappingProxyType({}) def _convert_tool_payload_value(key: str, value: object, *, to_chat: bool) -> object: @@ -105,7 +105,7 @@ def _normalize_tool_dialect( return {**data, **{key: value for key, value in replaceable if key in data}} # mutable-ok: plain body dict -def _is_chat_completions_body(data: Mapping[str, Any]) -> bool: +def _is_chat_completions_body(data: Mapping[str, object]) -> bool: messages: Final = data.get("messages") if isinstance(messages, list) and messages: return True @@ -1373,7 +1373,7 @@ async def _enforce_responses_ws_first_frame_model_auth( request: Request, model: str, user_api_key_dict: UserAPIKeyAuth, - llm_router: Any | None, + llm_router: "Router | None", ) -> None: from litellm.proxy.auth.user_api_key_auth import ( _enforce_key_and_fallback_model_access, @@ -1417,7 +1417,7 @@ async def _enforce_responses_ws_first_frame_model_auth( async def responses_websocket_endpoint( websocket: WebSocket, model: str | None = fastapi.Query(None, description="The model to use for the responses WebSocket session."), - user_api_key_dict=Depends(user_api_key_auth_websocket), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): """ Responses API WebSocket mode endpoint. @@ -1462,7 +1462,7 @@ async def responses_websocket_endpoint( return model, first_message = result - data: dict[str, Any] = { + data: dict[str, object] = { "model": model, "websocket": websocket, } @@ -1471,7 +1471,7 @@ async def responses_websocket_endpoint( # Construct a synthetic Request for pre-call processing headers_list: Final = list(websocket.scope.get("headers") or []) - scope: Final[dict[str, Any]] = { + scope: Final[dict[str, object]] = { "type": "http", "method": "POST", "path": "/v1/responses", diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 91a0c68fd58..3d0bd5e61c9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -50,12 +50,12 @@ def _route_user_config_request(data: dict, route_type: str): return ret_val -def _is_a2a_agent_model(model_name: Any) -> bool: +def _is_a2a_agent_model(model_name: object) -> bool: """Check if the model name is for an A2A agent (a2a/ prefix).""" return isinstance(model_name, str) and model_name.startswith("a2a/") -def _raise_if_model_fully_blocked(llm_router: LitellmRouter, model_name: Any, team_id: str | None) -> None: +def _raise_if_model_fully_blocked(llm_router: LitellmRouter, model_name: object, team_id: str | None) -> None: if not isinstance(model_name, str) or not model_name: return if not isinstance(llm_router, litellm.Router): diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index d985a546fa7..66071c05b4f 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -1,6 +1,6 @@ #### Video Endpoints ##### -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse @@ -161,7 +161,7 @@ async def video_list( # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params} + data: Final[dict[str, object]] = {"query_params": query_params} # Extract custom_llm_provider from headers, query params, or body custom_llm_provider: Final = ( @@ -246,7 +246,7 @@ async def video_status( ) # Create data with video_id - data: Final[dict[str, Any]] = {"video_id": video_id} + data: Final[dict[str, object]] = {"video_id": video_id} decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") @@ -345,7 +345,7 @@ async def video_content( ) # Create data with video_id - data: Final[dict[str, Any]] = {"video_id": video_id} + data: Final[dict[str, object]] = {"video_id": video_id} decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") @@ -653,7 +653,7 @@ async def video_get_character( ) original_requested_character_id: Final = character_id - data: Final[dict[str, Any]] = {"character_id": character_id} + data: Final[dict[str, object]] = {"character_id": character_id} decoded: Final = decode_character_id_with_provider(character_id) provider_from_id: Final = decoded.get("custom_llm_provider") diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 2dcaa200cc6..7bc1a6a52a3 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -29,6 +29,7 @@ from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.ingestion.vertex_ai_ingestion import VertexAIRAGIngestion from litellm.rag.rag_query import RAGQuery +from litellm.types.llms.openai import AllMessageValues from litellm.types.rag import ( RAGIngestOptions, RAGIngestResponse, @@ -204,7 +205,7 @@ def _suppressed_sub_call_billing() -> Iterator[None]: async def _execute_query_pipeline( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, @@ -311,7 +312,7 @@ async def _execute_query_pipeline( @client async def aquery( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, @@ -358,12 +359,12 @@ async def aquery( @client def query( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, **kwargs, -) -> ModelResponse | Coroutine[Any, Any, ModelResponse]: +) -> ModelResponse | Coroutine[None, None, ModelResponse]: """ Query a RAG pipeline. """ @@ -410,7 +411,7 @@ def ingest( file_id: str | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> RAGIngestResponse | Coroutine[Any, Any, RAGIngestResponse]: +) -> RAGIngestResponse | Coroutine[None, None, RAGIngestResponse]: """ Ingest a document into a vector store. diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 197d0c02ba8..367915156d1 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -399,7 +399,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_without_openai_transform( - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth | None", mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: str | None = None, mcp_auth_header: str | None = None, @@ -636,7 +636,7 @@ class LiteLLM_Proxy_MCP_Handler: async def _execute_tool_calls( tool_server_map: dict[str, str], tool_calls: Sequence[object], - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth | None", mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index d57d7da0410..a8d51f95e45 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -20,6 +20,7 @@ anthropic: import asyncio import builtins +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Any, Final @@ -54,19 +55,19 @@ class _LiteLLMParamsDictView: __slots__ = ("_params",) - def __init__(self, params: dict[str, Any]): + def __init__(self, params: Mapping[str, object]): self._params = params - def __getattr__(self, key: str) -> Any: + def __getattr__(self, key: str) -> object: return self._params.get(key) - def __getitem__(self, key: str) -> Any: + def __getitem__(self, key: str) -> object: return self._params.get(key) def __contains__(self, key: str) -> bool: return key in self._params - def get(self, key: str, default: Any = None) -> Any: + def get(self, key: str, default: object = None) -> object: return self._params.get(key, default) def keys(self): @@ -84,10 +85,10 @@ class _LiteLLMParamsDictView: def __len__(self) -> int: return len(self._params) - def dict(self) -> dict[str, Any]: + def dict(self) -> builtins.dict[str, object]: return dict(self._params) - def model_dump(self) -> builtins.dict[str, Any]: + def model_dump(self) -> builtins.dict[str, object]: return dict(self._params) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index be7653902a7..577cee0920d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -282,7 +282,7 @@ def _response_cost_or_none(response: ModelResponse) -> float | None: return float(cost) -def _effective_turn_off_message_logging(request_kwargs: Mapping[str, Any] | None) -> bool | None: +def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None: from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( initialize_standard_callback_dynamic_params, ) @@ -1925,7 +1925,7 @@ class ComplexityRouter(CustomLogger): ceiling_severity: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None best_model: str | None = None best_score = float("-inf") - candidate_scores: Final[list[dict[str, Any]]] = [] + candidate_scores: Final[list[dict[str, object]]] = [] for model in candidates: if floor_severity is not None and all( self._active_tier_severity(model_tier) < floor_severity diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index e2662d96b52..8f677b54700 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -1,7 +1,9 @@ import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -17,6 +19,72 @@ from litellm.proxy._types import KeyManagementSystem from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name +class _VaultAuthData(TypedDict): + """The ``auth`` block Vault returns from a login endpoint.""" + + client_token: ReadOnly[str] + lease_duration: ReadOnly[int] + + +class _VaultLoginResponse(TypedDict): + """Body of a Vault ``/v1/auth/.../login`` response.""" + + auth: ReadOnly[_VaultAuthData] + + +class _VaultSecretTarget(TypedDict): + """Resolved coordinates of one Vault KV v2 secret.""" + + url: ReadOnly[str] + data_key: ReadOnly[str] + secret_name: ReadOnly[str] + + +class _VaultSecretDataBlock(TypedDict, total=False): + """The inner ``data`` block of a Vault KV v2 read body.""" + + data: ReadOnly[Mapping[str, object]] + + +class _VaultSecretReadResponse(TypedDict, total=False): + """Body of a Vault KV v2 secret read, narrowed to the nesting this module walks.""" + + data: ReadOnly[_VaultSecretDataBlock] + + +class _VaultLoginResponseSource(Protocol): + """A Vault login call's HTTP response, read for the auth block it carries.""" + + def json(self) -> _VaultLoginResponse: ... + + +class _VaultSecretReadSource(Protocol): + """A Vault KV v2 read response, read for the nested secret data it carries.""" + + def json(self) -> _VaultSecretReadResponse: ... + + +class _JsonObjectSource(Protocol): + """A Vault response whose body is a JSON object nothing further is assumed about.""" + + def json(self) -> dict[str, object]: ... + + +def _vault_login_body(response: _VaultLoginResponseSource) -> _VaultLoginResponse: + """Decode the body of a Vault login response.""" + return response.json() + + +def _vault_secret_read_body(response: _VaultSecretReadSource) -> _VaultSecretReadResponse: + """Decode the body of a Vault KV v2 secret read response.""" + return response.json() + + +def _json_object_body(response: _JsonObjectSource) -> dict[str, object]: + """Decode a Vault response body as a plain JSON object.""" + return response.json() + + class HashicorpSecretManager(BaseSecretManager): def __init__(self): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user @@ -130,7 +198,8 @@ class HashicorpSecretManager(BaseSecretManager): ) resp.raise_for_status() - auth_data: Final = resp.json()["auth"] + login_response: Final = _vault_login_body(resp) + auth_data: Final = login_response["auth"] token: Final = auth_data["client_token"] _lease_duration: Final = auth_data["lease_duration"] @@ -191,8 +260,10 @@ class HashicorpSecretManager(BaseSecretManager): json=self._get_tls_cert_auth_body(), ) resp.raise_for_status() - token: Final = resp.json()["auth"]["client_token"] - _lease_duration: Final = resp.json()["auth"]["lease_duration"] + token_response: Final = _vault_login_body(resp) + token: Final = token_response["auth"]["client_token"] + lease_response: Final = _vault_login_body(resp) + _lease_duration: Final = lease_response["auth"]["lease_duration"] verbose_logger.debug("Successfully obtained Vault token via TLS cert auth.") self.cache.set_cache(key="hcp_vault_token", value=token, ttl=_lease_duration) return token @@ -205,9 +276,9 @@ class HashicorpSecretManager(BaseSecretManager): def get_url( self, secret_name: str, - namespace: str | None = None, - mount_name: str | None = None, - path_prefix: str | None = None, + namespace: object = None, + mount_name: object = None, + path_prefix: object = None, ) -> str: """ Constructs the Vault URL for KV v2 secrets. @@ -238,7 +309,7 @@ class HashicorpSecretManager(BaseSecretManager): _url += secret_name return _url - def _sanitize_plain_value(self, value: str | int | None) -> str | None: + def _sanitize_plain_value(self, value: object) -> str | None: if value is None: return None value_str: Final = str(value).strip() @@ -246,23 +317,23 @@ class HashicorpSecretManager(BaseSecretManager): return None return value_str - def _sanitize_path_component(self, value: str | int | None) -> str | None: + def _sanitize_path_component(self, value: object) -> str | None: sanitized_value = self._sanitize_plain_value(value) if sanitized_value is None: return None sanitized_value = sanitized_value.strip("/") return sanitized_value or None - def _extract_secret_manager_settings(self, optional_params: dict | None) -> dict[str, Any]: + def _extract_secret_manager_settings(self, optional_params: dict | None) -> dict[str, object]: if not isinstance(optional_params, dict): return {} candidate: Final = optional_params.get("secret_manager_settings") - source: Final = candidate if isinstance(candidate, dict) else optional_params + source: Final[Mapping[str, object]] = candidate if isinstance(candidate, dict) else optional_params allowed_keys: Final = {"namespace", "mount", "path_prefix", "data"} return {k: source[k] for k in allowed_keys if k in source} - def _build_secret_target(self, secret_name: str, optional_params: dict | None) -> dict[str, Any]: + def _build_secret_target(self, secret_name: str, optional_params: dict | None) -> _VaultSecretTarget: settings: Final = self._extract_secret_manager_settings(optional_params) namespace: Final = settings.get("namespace", self.vault_namespace) @@ -331,7 +402,7 @@ class HashicorpSecretManager(BaseSecretManager): response.raise_for_status() # For KV v2, the secret is in response.json()["data"]["data"] - json_resp: Final = response.json() + json_resp: Final = _json_object_body(response) _value: Final = self._get_secret_value_from_json_response(json_resp) self.cache.set_cache(secret_name, _value) return _value @@ -362,7 +433,7 @@ class HashicorpSecretManager(BaseSecretManager): response.raise_for_status() # For KV v2, the secret is in response.json()["data"]["data"] - json_resp: Final = response.json() + json_resp: Final = _json_object_body(response) _value: Final = self._get_secret_value_from_json_response(json_resp) self.cache.set_cache(secret_name, _value) return _value @@ -379,7 +450,7 @@ class HashicorpSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Writes a secret to Vault KV v2 using an async HTTPX client. @@ -413,7 +484,7 @@ class HashicorpSecretManager(BaseSecretManager): json=data, ) response.raise_for_status() - return response.json() + return _json_object_body(response) except Exception as e: verbose_logger.exception("Error writing secret to Hashicorp Vault: %s", e) return {"status": "error", "message": str(e)} @@ -500,7 +571,7 @@ class HashicorpSecretManager(BaseSecretManager): headers=self._get_request_headers(), ) response.raise_for_status() - json_resp: Final = response.json() + json_resp: Final = _vault_secret_read_body(response) # Use data_key from target to get the correct value data_key: Final = new_target["data_key"] new_secret_value_from_vault: Final = json_resp.get("data", {}).get("data", {}).get(data_key, None) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index fcade835cce..32d88da0085 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -327,6 +327,10 @@ class BatchGuardrailReport(BaseModel): """Every record that was redacted or dropped, in file order.""" +_JsonValue: TypeAlias = object +"""Alias for ``object``, usable inside model bodies that declare a field named ``object``.""" + + BATCH_GUARDRAIL_RESPONSE_FIELD: Final = "litellm_batch_guardrail" @@ -1191,7 +1195,7 @@ class ShellToolParam(TypedDict, total=False): type: Required[Literal["shell"] | str] """The type of tool. Use ``\"shell\"``.""" - environment: Required[dict[str, Any]] + environment: Required[dict[str, object]] """Environment config: ``type`` (e.g. ``\"container_auto\"``, ``\"container_reference\"``, ``\"local\"``), optional ``container_id``, ``network_policy``, ``domain_secrets``, ``skills``.""" @@ -1308,7 +1312,7 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): @field_validator("cost", mode="before") @classmethod - def parse_cost(cls, v: Any) -> float | None: + def parse_cost(cls, v: object) -> object: """Normalise cost: accept either a float or a dict with a ``total_cost`` key.""" if isinstance(v, dict): return v.get("total_cost") @@ -1805,7 +1809,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): type: str # e.g., 'invalid_request_error' code: str # e.g., 'context_length_exceeded' message: str - param: str | dict[str, Any] | None = None + param: str | dict[str, object] | None = None class ErrorEvent(BaseLiteLLMOpenAIResponseObject): @@ -2418,7 +2422,7 @@ class OpenAIVideoObject(BaseModel): expires_at: int | None = None """Unix timestamp (seconds) for when the downloadable assets expire, if set.""" - error: dict[str, Any] | None = None + error: dict[str, _JsonValue] | None = None """Error payload that explains why generation failed, if applicable.""" progress: int | None = None @@ -2436,15 +2440,15 @@ class OpenAIVideoObject(BaseModel): model: str | None = None """The video generation model that produced the job.""" - _hidden_params: dict[str, Any] = {} + _hidden_params: dict[str, _JsonValue] = {} def __contains__(self, key) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key, default=None) -> _JsonValue: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key) -> _JsonValue: return getattr(self, key) def json(self, **kwargs): diff --git a/litellm/types/router.py b/litellm/types/router.py index ab6c807ba20..e0957383aac 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -369,7 +369,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): @model_validator(mode="before") @classmethod - def preprocess_input_data(cls, data: Any) -> Any: + def preprocess_input_data(cls, data: object) -> object: """ Pre-process input data before validation: 1. Filter out reserved Python keywords ('self', 'params', '__class__') to prevent @@ -627,6 +627,11 @@ class AlertingConfig(BaseModel): alerting_threshold: float | None = 300 +def _resolved_annotations(model_class: type[object]) -> Mapping[str, object]: + """Resolve a class's annotations, keeping each resolved annotation opaque.""" + return get_type_hints(model_class) + + class ModelGroupInfo(BaseModel): model_group: str providers: list[str] @@ -655,7 +660,7 @@ class ModelGroupInfo(BaseModel): configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None def __init__(self, **data) -> None: - for field_name, field_type in get_type_hints(self.__class__).items(): + for field_name, field_type in _resolved_annotations(self.__class__).items(): if field_type is bool and data.get(field_name) is None: data[field_name] = False super().__init__(**data) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 22d27bc3266..b71d6784873 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -112,7 +112,9 @@ class VectorStoreRegistry: Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS. """ # Get the list of supported param names from the Literal type - supported_params: Final = get_args(VECTOR_STORE_OPENAI_PARAMS) + supported_params: Final = tuple( + param for param in get_args(VECTOR_STORE_OPENAI_PARAMS) if isinstance(param, str) + ) # Extract only the params that exist in the tool kwargs: Final = {param: tool.get(param) for param in supported_params if param in tool} @@ -503,7 +505,7 @@ class VectorStoreRegistry: vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db - def get_credentials_for_vector_store(self, vector_store_id: str) -> dict[str, Any]: + def get_credentials_for_vector_store(self, vector_store_id: str) -> dict[str, object]: """ Get the credentials for a vector store diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 569c23cd03f..2dfa92ae694 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2991 + "limit": 2985 }, "ANN002": { "limit": 71 @@ -9,13 +9,13 @@ "limit": 809 }, "ANN201": { - "limit": 2002 + "limit": 2001 }, "ANN202": { - "limit": 841 + "limit": 835 }, "ANN204": { - "limit": 694 + "limit": 693 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 387 + "limit": 307 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1084 + "limit": 1073 }, "TRY002": { "limit": 524 @@ -246,7 +246,7 @@ "limit": 113 }, "TRY300": { - "limit": 855 + "limit": 854 }, "UP028": { "limit": 2 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 83c49afb538..3d2e97d55a5 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22403 + "limit": 22367 }, "LIT002": { - "limit": 26780 + "limit": 26777 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16512 + "limit": 16507 }, "LIT011": { - "limit": 5537 + "limit": 5535 }, "LIT012": { "limit": 4495 From 5f44bdd1c1230e4ce18515453c8747cb55d7cf4a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:06:08 +0000 Subject: [PATCH 41/92] test: trim mcp fixture docstring and reload comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/conftest.py | 11 ++--------- .../mcp_server/test_mcp_server_identity_env.py | 8 +++----- 2 files changed, 5 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index c559e023c47..2ccba2b2055 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -9,15 +9,8 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @pytest.fixture(autouse=True) def _hermetic_mcp_server_registry(): - """Snapshot and restore the global manager's server-registry state around every test. - - ``global_mcp_server_manager`` is a module-global singleton, and many tests in this - package seed ``registry``/``config_mcp_servers`` (or clear them) without cleaning up. - In a shared CI shard the leaked entries poison later tests in the same worker, e.g. - the ``all_proxy_servers`` sentinel expansion in ``auth/`` suddenly sees a bridge - server registered by a discovery test, so the outcome depends on xdist scheduling. - Restoring the state here makes ordering irrelevant. - """ + """Restore the singleton ``global_mcp_server_manager``'s registry state around every + test, so entries seeded by one test never leak into another on a shared shard.""" saved_registry = dict(global_mcp_server_manager.registry) saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py index 934810c305f..1c65adac4c6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py @@ -25,11 +25,9 @@ def _env_and_reload(**env): saved = {key: os.environ.get(key) for key in env} utils_module = importlib.import_module(UTILS_MODULE) mgmt_module = importlib.import_module(MGMT_MODULE) - # Restore the pre-reload module attributes afterwards instead of reloading - # a third time: a reload re-creates every class in the module, so modules - # that imported names like MCPMissingUserEnvVarsError before this test - # would keep raising the old class while pytest.raises in later tests - # matches the new one + # Restore pre-reload module attributes afterwards instead of reloading again: + # a reload re-creates the module's classes, breaking exception identity for + # modules that imported them earlier snapshots = {module: dict(vars(module)) for module in (utils_module, mgmt_module)} def _apply_env(values): From 054acb2223cb4e4a0c8c2a2b57fc0ad00a854fa5 Mon Sep 17 00:00:00 2001 From: yatishgoel Date: Tue, 1 Sep 2026 17:05:08 +0530 Subject: [PATCH 42/92] fix(ui): stop checkboxes stretching to the full width of a form field --- .../src/components/CreateUserButton.test.tsx | 15 +++++++++++++++ .../src/components/CreateUserButton.tsx | 2 +- .../SSOSettings/Modals/BaseSSOSettingsForm.tsx | 4 ++-- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 3777c46973f..46c66ee9ef6 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -294,6 +294,21 @@ describe("CreateUserButton", () => { }); }); + it("lays the send invitation email checkbox out beside its label", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + const checkbox = within(dialog).getByRole("checkbox"); + + expect(checkbox.closest('[data-slot="field"]')).toHaveAttribute("data-orientation", "horizontal"); + }); + describe("organizations", () => { it("should send organizations list in POST body when organizations are selected", async () => { const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index e052a9e0818..0f7c356b8cc 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -270,7 +270,7 @@ export const CreateUserButton: React.FC = ({ ); const sendInviteEmailField = ( - + {({ id, value, onChange, onBlur }) => ( )} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx index 5216da382d0..cb97304f77a 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -303,7 +303,7 @@ const SSOProviderField = ({ field }: { field: SSOProviderConfig["fields"][number if (field.type === "checkbox") { return ( - + {({ value, onChange, onBlur, id, ...rest }) => ( (); return ( - + {({ value, onChange, onBlur, id, ...rest }) => ( Date: Tue, 1 Sep 2026 13:15:22 +0000 Subject: [PATCH 43/92] Registry audit: Fireworks DeepSeek V4 Flash 0731 pricing, Databricks DeepSeek V4 entries, provider deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 66 ++++++++++++++++++- model_prices_and_context_window.json | 66 ++++++++++++++++++- .../test_fireworks_serverless_model_costs.py | 38 +++++++++++ 3 files changed, 164 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4229d80a671..adc4a557e63 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -554,6 +554,7 @@ "supports_vision": true }, "amazon.nova-sonic-v1:0": { + "deprecation_date": "2026-09-14", "input_cost_per_audio_token": 3.4e-06, "input_cost_per_token": 6e-08, "litellm_provider": "bedrock", @@ -3045,6 +3046,7 @@ "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { + "deprecation_date": "2027-12-05", "supports_mid_conversation_system": true, "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -3078,6 +3080,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5": { + "deprecation_date": "2027-07-08", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3110,6 +3113,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-4-8": { + "deprecation_date": "2027-09-01", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3188,6 +3192,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-5": { + "deprecation_date": "2027-06-30", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -15101,6 +15106,60 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/databricks-deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "input_dbu_cost_per_token": 2e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "output_dbu_cost_per_token": 4e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "databricks/databricks-deepseek-v4-pro-0813": { + "cache_read_input_token_cost": 1.3202e-07, + "input_cost_per_token": 1.31999e-06, + "input_dbu_cost_per_token": 1.8857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 3.95997e-06, + "output_dbu_cost_per_token": 5.6571e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "databricks/databricks-gemini-2-5-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, @@ -20631,6 +20690,7 @@ "supports_image_size": false }, "gemini-live-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -52237,14 +52297,14 @@ "supports_vision": true }, "fireworks_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 6.6e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4229d80a671..adc4a557e63 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -554,6 +554,7 @@ "supports_vision": true }, "amazon.nova-sonic-v1:0": { + "deprecation_date": "2026-09-14", "input_cost_per_audio_token": 3.4e-06, "input_cost_per_token": 6e-08, "litellm_provider": "bedrock", @@ -3045,6 +3046,7 @@ "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { + "deprecation_date": "2027-12-05", "supports_mid_conversation_system": true, "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -3078,6 +3080,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5": { + "deprecation_date": "2027-07-08", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3110,6 +3113,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-4-8": { + "deprecation_date": "2027-09-01", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3188,6 +3192,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-5": { + "deprecation_date": "2027-06-30", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -15101,6 +15106,60 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/databricks-deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "input_dbu_cost_per_token": 2e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "output_dbu_cost_per_token": 4e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "databricks/databricks-deepseek-v4-pro-0813": { + "cache_read_input_token_cost": 1.3202e-07, + "input_cost_per_token": 1.31999e-06, + "input_dbu_cost_per_token": 1.8857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 3.95997e-06, + "output_dbu_cost_per_token": 5.6571e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "databricks/databricks-gemini-2-5-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, @@ -20631,6 +20690,7 @@ "supports_image_size": false }, "gemini-live-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -52237,14 +52297,14 @@ "supports_vision": true }, "fireworks_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 6.6e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index 0458af0da0e..a7a9e0fc37d 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -84,3 +84,41 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) assert info["max_input_tokens"] == expected["max_input_tokens"] assert info["max_output_tokens"] == expected["max_output_tokens"] + + +TWIN_PINNED_PRICES = { + "deepseek-v4-flash-0731": { + "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 7e-09, + "output_cost_per_token": 6.6e-07, + }, +} + + +def test_deepseek_v4_flash_0731_twins_pin_published_pricing(model_data): + """Both 0731 entries carry the price published at docs.fireworks.ai/serverless/pricing.""" + for bare_suffix, expected in TWIN_PINNED_PRICES.items(): + for key in ( + f"fireworks_ai/{bare_suffix}", + f"fireworks_ai/accounts/fireworks/models/{bare_suffix}", + ): + entry = model_data[key] + for field, value in expected.items(): + assert entry[field] == pytest.approx(value), f"{key}.{field}" + + +def test_fireworks_account_prefixed_twins_agree_on_price(model_data): + """Every accounts/fireworks/models/X entry prices identically to its bare fireworks_ai/X twin.""" + prefix = "fireworks_ai/accounts/fireworks/models/" + pairs_checked = 0 + for key, entry in model_data.items(): + if not key.startswith(prefix): + continue + bare_key = f"fireworks_ai/{key[len(prefix):]}" + bare_entry = model_data.get(bare_key) + if bare_entry is None: + continue + pairs_checked += 1 + for field in sorted({f for f in (*entry, *bare_entry) if "cost" in f}): + assert entry.get(field) == bare_entry.get(field), f"{key} vs {bare_key}: {field}" + assert pairs_checked >= 20 From 5263570e68e66bad407b0e16cfad25a73eee9e47 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:33:31 +0000 Subject: [PATCH 44/92] Add cerebras/zai-glm-4.7 deprecation_date per Cerebras deprecations page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index adc4a557e63..fae9699541a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12292,6 +12292,7 @@ "supports_tool_choice": true }, "cerebras/zai-glm-4.7": { + "deprecation_date": "2026-08-17", "input_cost_per_token": 2.25e-06, "litellm_provider": "cerebras", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index adc4a557e63..fae9699541a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12292,6 +12292,7 @@ "supports_tool_choice": true }, "cerebras/zai-glm-4.7": { + "deprecation_date": "2026-08-17", "input_cost_per_token": 2.25e-06, "litellm_provider": "cerebras", "max_input_tokens": 128000, From 9a1aebc14673a6150c7d47c3d0c2d797f697ff92 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:50:21 +0000 Subject: [PATCH 45/92] fix(registry): declare databricks deepseek cache-write rate at the input rate per repo convention Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ .../llms/databricks/test_databricks_cost_calculator.py | 2 ++ 3 files changed, 6 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fae9699541a..dd45495884d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15108,6 +15108,7 @@ "supports_vision": true }, "databricks/databricks-deepseek-v4-flash-0731": { + "cache_creation_input_token_cost": 1.4e-07, "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.4e-07, "input_dbu_cost_per_token": 2e-06, @@ -15135,6 +15136,7 @@ "supports_vision": false }, "databricks/databricks-deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 1.31999e-06, "cache_read_input_token_cost": 1.3202e-07, "input_cost_per_token": 1.31999e-06, "input_dbu_cost_per_token": 1.8857e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fae9699541a..dd45495884d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15108,6 +15108,7 @@ "supports_vision": true }, "databricks/databricks-deepseek-v4-flash-0731": { + "cache_creation_input_token_cost": 1.4e-07, "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.4e-07, "input_dbu_cost_per_token": 2e-06, @@ -15135,6 +15136,7 @@ "supports_vision": false }, "databricks/databricks-deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 1.31999e-06, "cache_read_input_token_cost": 1.3202e-07, "input_cost_per_token": 1.31999e-06, "input_dbu_cost_per_token": 1.8857e-05, diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 29ad8ee4b6e..e72642f7a04 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -62,6 +62,8 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-gemini-2-5-pro": ("22.321", "178.571", "22.321", "2.232"), "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), "databricks/databricks-kimi-k3": ("42.857", "214.286", "42.857", "4.286"), + "databricks/databricks-deepseek-v4-flash-0731": ("2.000", "4.000", "2.000", "0.400"), + "databricks/databricks-deepseek-v4-pro-0813": ("18.857", "56.571", "18.857", "1.886"), "databricks/databricks-glm-5-2": ("20.000", "62.857", "20.000", "3.714"), } PROMOTIONAL_DISCOUNT: Final = 0.80 From e3b5cf13c6565f7a2b34b44b071a408354e4ff43 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 1 Sep 2026 14:14:50 +0000 Subject: [PATCH 46/92] feat(helm): add Argo CD PreSync hook and rollout strategy knobs to the componentized chart Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm/templates/backend/deployment.yaml | 4 ++ .../litellm/templates/gateway/deployment.yaml | 4 ++ helm/litellm/templates/migrations-job.yaml | 12 +++- helm/litellm/templates/ui/deployment.yaml | 4 ++ .../tests/migration_job_hooks_tests.yaml | 63 ++++++++++++++++++ .../litellm/tests/rollout_strategy_tests.yaml | 66 +++++++++++++++++++ helm/litellm/values.yaml | 29 ++++++++ 7 files changed, 181 insertions(+), 1 deletion(-) create mode 100644 helm/litellm/tests/migration_job_hooks_tests.yaml create mode 100644 helm/litellm/tests/rollout_strategy_tests.yaml diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 5c0431fc0bd..0db2f0b3d43 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend spec: + {{- with .Values.backend.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.backend.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index d5363d0096e..5030ba2c9dc 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: gateway spec: + {{- with .Values.gateway.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.gateway.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index 9cd8397f794..8d33081e72f 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -7,6 +7,8 @@ # # Running this pre-upgrade closes the window where new application pods would # otherwise serve traffic against the previous release's unmigrated schema. +# Argo CD users can swap the Helm hook for a PreSync hook through +# `migrationJob.hooks`, which re-runs the Job on every sync. apiVersion: batch/v1 kind: Job metadata: @@ -14,10 +16,18 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: migrations + {{- if or .Values.migrationJob.hooks.helm.enabled .Values.migrationJob.hooks.argocd.enabled }} annotations: + {{- if .Values.migrationJob.hooks.helm.enabled }} helm.sh/hook: pre-install,pre-upgrade helm.sh/hook-delete-policy: before-hook-creation - helm.sh/hook-weight: "0" + helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "0" | quote }} + {{- end }} + {{- if .Values.migrationJob.hooks.argocd.enabled }} + argocd.argoproj.io/hook: PreSync + argocd.argoproj.io/hook-delete-policy: BeforeHookCreation + {{- end }} + {{- end }} spec: backoffLimit: {{ .Values.migrationJob.backoffLimit }} ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index 91d6de39ea6..b992b347bad 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui spec: + {{- with .Values.ui.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.ui.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/tests/migration_job_hooks_tests.yaml b/helm/litellm/tests/migration_job_hooks_tests.yaml new file mode 100644 index 00000000000..650d2700429 --- /dev/null +++ b/helm/litellm/tests/migration_job_hooks_tests.yaml @@ -0,0 +1,63 @@ +suite: test migrations Job hook annotations +templates: + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: runs as a Helm pre-install / pre-upgrade hook by default + asserts: + - equal: + path: metadata.annotations["helm.sh/hook"] + value: pre-install,pre-upgrade + - equal: + path: metadata.annotations["helm.sh/hook-delete-policy"] + value: before-hook-creation + - equal: + path: metadata.annotations["helm.sh/hook-weight"] + value: "0" + - notExists: + path: metadata.annotations["argocd.argoproj.io/hook"] + + - it: adds the Argo CD PreSync hook when asked + set: + migrationJob.hooks.argocd.enabled: true + asserts: + - equal: + path: metadata.annotations["argocd.argoproj.io/hook"] + value: PreSync + - equal: + path: metadata.annotations["argocd.argoproj.io/hook-delete-policy"] + value: BeforeHookCreation + + - it: drops the Helm hook so Argo CD owns the Job + set: + migrationJob.hooks.argocd.enabled: true + migrationJob.hooks.helm.enabled: false + asserts: + - equal: + path: metadata.annotations["argocd.argoproj.io/hook"] + value: PreSync + - notExists: + path: metadata.annotations["helm.sh/hook"] + - notExists: + path: metadata.annotations["helm.sh/hook-delete-policy"] + - notExists: + path: metadata.annotations["helm.sh/hook-weight"] + + - it: renders an ordinary Job when both hooks are disabled + set: + migrationJob.hooks.helm.enabled: false + asserts: + - notExists: + path: metadata.annotations + - equal: + path: kind + value: Job + + - it: honours a custom Helm hook weight + set: + migrationJob.hooks.helm.weight: "-5" + asserts: + - equal: + path: metadata.annotations["helm.sh/hook-weight"] + value: "-5" diff --git a/helm/litellm/tests/rollout_strategy_tests.yaml b/helm/litellm/tests/rollout_strategy_tests.yaml new file mode 100644 index 00000000000..b12e2073c7c --- /dev/null +++ b/helm/litellm/tests/rollout_strategy_tests.yaml @@ -0,0 +1,66 @@ +suite: test rolling update strategy on the component deployments +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: leaves the strategy to Kubernetes defaults when unset + asserts: + - notExists: + path: spec.strategy + + - it: renders the configured strategy on each deployment + set: + gateway.strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + backend.strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: "25%" + maxSurge: 2 + ui.strategy: + type: Recreate + asserts: + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + template: gateway/deployment.yaml + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 25% + maxSurge: 2 + template: backend/deployment.yaml + - equal: + path: spec.strategy + value: + type: Recreate + template: ui/deployment.yaml + + - it: keeps a component on the cluster default when only another one sets a strategy + set: + gateway.strategy: + type: Recreate + asserts: + - equal: + path: spec.strategy.type + value: Recreate + template: gateway/deployment.yaml + - notExists: + path: spec.strategy + template: backend/deployment.yaml + - notExists: + path: spec.strategy + template: ui/deployment.yaml diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index d0c80fd6f6f..378c3b7a618 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -75,6 +75,22 @@ serviceAccounts: # generate` — the migration engine doesn't need the generated client. migrationJob: enabled: true + # Which controller is responsible for running the Job. + # + # `helm.enabled` renders the Helm pre-install / pre-upgrade hook, so the Job + # runs whenever `helm upgrade` sees a change to apply. `argocd.enabled` + # renders an Argo CD PreSync hook instead, which runs the Job on every sync + # even when the rendered manifests are unchanged: the way to re-run + # migrations on demand from a GitOps pipeline. Turning the Helm hook off + # while the Argo CD hook is on leaves the Job out of Helm's own upgrade + # path, which is what Argo CD users want since Argo, not Helm, applies the + # manifests. + hooks: + helm: + enabled: true + weight: "0" + argocd: + enabled: false backoffLimit: 4 ttlSecondsAfterFinished: 120 # Wall-clock budget for the whole Job, shared across every `backoffLimit` @@ -257,6 +273,15 @@ gateway: initialDelaySeconds: 5 periodSeconds: 10 timeoutSeconds: 10 + # Rolling update tuning for the gateway Deployment. Empty by default, so + # Kubernetes applies its own RollingUpdate defaults (25% maxSurge / + # 25% maxUnavailable). Example, for a surge-only rollout behind a load + # balancer that must never lose capacity: + # type: RollingUpdate + # rollingUpdate: + # maxUnavailable: 0 + # maxSurge: 1 + strategy: {} # Optional startupProbe. Empty by default, so existing installs are unchanged # and liveness/readiness apply from container start. Set it to gate # liveness/readiness until a slow cold start finishes — a high failureThreshold @@ -369,6 +394,8 @@ backend: initialDelaySeconds: 5 periodSeconds: 10 timeoutSeconds: 10 + # Same shape as gateway.strategy. + strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} hpa: @@ -433,6 +460,8 @@ ui: httpGet: { path: /, port: http } initialDelaySeconds: 2 periodSeconds: 10 + # Same shape as gateway.strategy. + strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} hpa: From 0d7035989c8acb1ba3d8d1d5b149f9d8bc3dc3fd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 08:10:36 -0700 Subject: [PATCH 47/92] fix(mcp): persist alias MCP grants verbatim instead of rewriting to local server ids Since PR #29128, key create/update/regenerate resolved every object_permission.mcp_servers entry against the saving instance's DB + config registry and persisted the resolved server ids. For config-loaded servers the id is derived from a hash of the regional URL, so in a shared-database multi-region deployment the rewrite baked one region's ids into the row and every other region denied the key. Grants written before v1.88.0 kept the raw alias and kept working, which is why only newly provisioned keys broke. Keep the validation and the stale-entry drop (the LIT-3278 fix), but persist the caller's original identifiers for everything that resolves. Read-time expand_permission_list already maps a name to each region's local server id. --- .../object_permission_utils.py | 40 ++++++------ .../test_object_permission_utils.py | 61 +++++++++++++++---- 2 files changed, 66 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 13080a6cf83..a2fbf80422c 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -286,7 +286,7 @@ async def _resolve_mcp_server_identifiers_to_ids( return resolved -def _rewrite_object_permission_mcp_servers( +def _drop_stale_object_permission_mcp_servers( object_permission: ObjectPermissionDict, identifier_to_server_ids: dict[str, set[str]], ) -> None: @@ -294,16 +294,18 @@ def _rewrite_object_permission_mcp_servers( if not isinstance(mcp_servers, list): return - normalized_servers: Final[list[str]] = [] - for identifier in mcp_servers: - if identifier == SpecialMCPServerNames.no_mcp_servers.value: - normalized_servers.append(SpecialMCPServerNames.no_mcp_servers.value) - continue - normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, []))) - object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers) + # Persist original identifiers, never resolved ids: shared-DB multi-region + # instances each expand a name/alias to their own local server id at read + # time. Only entries resolving to nothing (deleted servers, typos) drop. + kept_servers: Final = [ + identifier + for identifier in mcp_servers + if identifier == SpecialMCPServerNames.no_mcp_servers.value or identifier_to_server_ids.get(identifier) + ] + object_permission["mcp_servers"] = _dedupe_preserving_order(kept_servers) -def _rewrite_object_permission_mcp_tool_permissions( +def _drop_stale_object_permission_mcp_tool_permissions( object_permission: ObjectPermissionDict, identifier_to_server_ids: dict[str, set[str]], ) -> None: @@ -311,31 +313,25 @@ def _rewrite_object_permission_mcp_tool_permissions( if not isinstance(mcp_tool_permissions, dict): return - normalized_tool_permissions: Final[dict[str, list[str]]] = {} - for identifier, tools in mcp_tool_permissions.items(): - if not isinstance(tools, list): - tools = [] - for server_id in sorted(identifier_to_server_ids.get(identifier, [])): - normalized_tool_permissions.setdefault(server_id, []) - normalized_tool_permissions[server_id].extend(tools) - object_permission["mcp_tool_permissions"] = { - server_id: _dedupe_preserving_order(tools) for server_id, tools in normalized_tool_permissions.items() + identifier: _dedupe_preserving_order(tools if isinstance(tools, list) else []) + for identifier, tools in mcp_tool_permissions.items() + if identifier_to_server_ids.get(identifier) } -def _rewrite_object_permission_mcp_identifiers( +def _drop_stale_object_permission_mcp_identifiers( object_permission: ObjectPermissionDict | None, identifier_to_server_ids: dict[str, set[str]], ) -> None: if not object_permission or not isinstance(object_permission, dict): return - _rewrite_object_permission_mcp_servers( + _drop_stale_object_permission_mcp_servers( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) - _rewrite_object_permission_mcp_tool_permissions( + _drop_stale_object_permission_mcp_tool_permissions( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) @@ -615,7 +611,7 @@ async def validate_key_mcp_servers_against_team( "validate_key_mcp_servers_against_team: ignoring stale MCP server identifiers (no longer in registry or DB): %s", sorted(stale_identifiers), ) - _rewrite_object_permission_mcp_identifiers( + _drop_stale_object_permission_mcp_identifiers( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 5ef83344c1a..d6a48f53ddd 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -1,11 +1,9 @@ import json +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException - -from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import ( LiteLLM_ObjectPermissionBase, LiteLLM_ObjectPermissionTable, @@ -13,10 +11,10 @@ from litellm.proxy._types import ( SpecialMCPServerName, ) from litellm.proxy.management_helpers.object_permission_utils import ( + _drop_stale_object_permission_mcp_servers, _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, _resolve_team_allowed_mcp_servers, - _rewrite_object_permission_mcp_servers, _set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, validate_key_mcp_servers_against_team, @@ -153,10 +151,10 @@ def test_extract_requested_mcp_server_ids_excludes_no_mcp_servers_sentinel(): assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"} -def test_rewrite_object_permission_mcp_servers_preserves_sentinel(): - obj_perm = {"mcp_servers": ["no-mcp-servers", "alias-1"]} - _rewrite_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}}) - assert obj_perm["mcp_servers"] == ["no-mcp-servers", "server-1"] +def test_drop_stale_object_permission_mcp_servers_preserves_sentinel_and_alias(): + obj_perm = {"mcp_servers": ["no-mcp-servers", "alias-1", "gone-id"]} + _drop_stale_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}, "gone-id": set()}) + assert obj_perm["mcp_servers"] == ["no-mcp-servers", "alias-1"] @pytest.mark.asyncio @@ -692,9 +690,10 @@ async def test_validate_mcp_server_alias_outside_team_scope_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_mcp_server_alias_is_normalized_before_save( - mock_access_groups, mock_allow_all -): +async def test_validate_mcp_server_alias_persists_verbatim(mock_access_groups, mock_allow_all): + """Regression for the multi-region shared-DB setup: an alias grant must be + stored as the alias, so every instance can expand it to its own local id. + Rewriting to this instance's server_id breaks access on the other region.""" team_obj = _make_team_obj(mcp_servers=["allowed-server-id"]) object_permission = { "mcp_servers": ["allowed-alias"], @@ -706,8 +705,44 @@ async def test_validate_mcp_server_alias_is_normalized_before_save( team_obj=team_obj, ) - assert object_permission["mcp_servers"] == ["allowed-server-id"] - assert object_permission["mcp_tool_permissions"] == {"allowed-server-id": ["tool1"]} + assert object_permission["mcp_servers"] == ["allowed-alias"] + assert object_permission["mcp_tool_permissions"] == {"Allowed Server": ["tool1"]} + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_alias_grant_expands_on_other_region_after_save(mock_access_groups, mock_allow_all): + """Full cross-region flow: save a key on the west instance (alias resolves to + west's hash-derived id), then expand the persisted grant on the central + instance, whose registry maps the same alias to a different id.""" + west_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("west-id", alias="github-mcp")]) + central_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("central-id", alias="github-mcp")]) + + team_obj = _make_team_obj(mcp_servers=["west-id"]) + object_permission = {"mcp_servers": ["github-mcp"]} + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + west_mgr, + ): + await validate_key_mcp_servers_against_team( + object_permission=object_permission, + team_obj=team_obj, + ) + assert object_permission["mcp_servers"] == ["github-mcp"] + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + expand = MCPServerManager.expand_permission_list + assert expand(west_mgr, object_permission["mcp_servers"]) == ["west-id"] + assert expand(central_mgr, object_permission["mcp_servers"]) == ["central-id"] @pytest.mark.asyncio From 7a761ccf5a22ef1db10c429b0c1960044d8a4659 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 08:10:41 -0700 Subject: [PATCH 48/92] test(e2e): cover alias MCP grant persisting verbatim on key generate A single-instance run cannot reproduce the two-region setup, but the regression is fully visible in one: the alias must survive to /key/info unrewritten, and the alias-granted key must still list the server's tools. The broken write path stored the resolved server id instead. --- tests/e2e/coverage_registry/mcp.yaml | 8 +++++++ tests/e2e/mcp/test_mcp_key_access_e2e.py | 29 ++++++++++++++++++++++++ tests/e2e/models.py | 1 + 3 files changed, 38 insertions(+) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..7ed2a7ab2ad 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -15,6 +15,14 @@ assertions: [access_group_scoped] source: "test_mcp_access_group_e2e.py" rationale: "A key granted an MCP access group sees the tagged server's tools; a key with a different group does not. Access-group-scoped tool selection at key creation" +- id: mcp.list_tools.api_key.alias_grant_persists + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [alias_grant_persists] + source: "object_permission_utils.py validate_key_mcp_servers_against_team" + rationale: "A key granted an MCP server by alias keeps the alias verbatim in its stored object_permission (shared-DB multi-region instances each resolve it to their local server id at read time) and still lists the server's tools" - id: mcp.list_tools.api_key.denied_without_permission module: mcp tier: P0 diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 88ab5666084..7fe40795f1c 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -30,6 +30,35 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str return key +class TestMcpKeyGrantByAlias: + @pytest.mark.covers("mcp.list_tools.api_key.alias_grant_persists") + def test_alias_grant_persists_verbatim_and_lists_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + """A key granted an MCP server by its alias must store the alias, not the + resolved server_id: in a shared-DB multi-region deployment each instance + derives a different id for the same config server, so only the alias + grants access on every region. The same key must still see the server's + tools, proving the alias grant is honored at request time.""" + server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) + alias = next(row.alias for row in client.registered_servers() if row.server_id == server_id) + assert alias, f"registered server {server_id} has no alias to grant by" + + key = _key(client, resources, mcp_servers=[alias]) + + stored = client.proxy.key_info(key).object_permission + assert stored is not None and stored.mcp_servers == [alias], ( + f"alias grant was rewritten before persisting (expected [{alias!r}]): " + f"{stored.mcp_servers if stored else None}. A stored server_id is region-local " + f"and breaks the grant on every other instance sharing this database" + ) + + _ = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) + + class TestMcpKeyWithoutAccessIsDenied: @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission") def test_list_tools_denied_without_permission( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 6d9ccad9a24..967144dfb16 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -114,6 +114,7 @@ class KeyInfo(BaseModel): budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None budget_limits: list[BudgetWindowState] | None = None + object_permission: ObjectPermission | None = None class KeyInfoResponse(BaseModel): From 4118db5a5076e62de66d48266830da13b610841b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 08:23:26 -0700 Subject: [PATCH 49/92] test: drop litellm-internal patches from the cross-region alias test (TQ008) The pure helpers express the same regression: the save-side drop must leave the alias in place, and two registries must expand it to their own ids. The full validate path is already covered by the persists-verbatim test and the live e2e test. --- .../test_object_permission_utils.py | 29 ++++--------------- 1 file changed, 6 insertions(+), 23 deletions(-) diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index d6a48f53ddd..f2b6b799271 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -709,33 +709,16 @@ async def test_validate_mcp_server_alias_persists_verbatim(mock_access_groups, m assert object_permission["mcp_tool_permissions"] == {"Allowed Server": ["tool1"]} -@pytest.mark.asyncio -@patch( - "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", - return_value=set(), -) -@patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", - new_callable=AsyncMock, - return_value=[], -) -async def test_validate_alias_grant_expands_on_other_region_after_save(mock_access_groups, mock_allow_all): - """Full cross-region flow: save a key on the west instance (alias resolves to - west's hash-derived id), then expand the persisted grant on the central - instance, whose registry maps the same alias to a different id.""" +def test_alias_grant_expands_on_other_region_after_save(): + """Cross-region flow: the west instance saves an alias grant (its resolver maps + the alias to west's hash-derived id), then the central instance, whose registry + maps the same alias to a different id, expands the persisted grant. Rewriting + to west's id at save time is exactly the regression this guards against.""" west_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("west-id", alias="github-mcp")]) central_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("central-id", alias="github-mcp")]) - team_obj = _make_team_obj(mcp_servers=["west-id"]) object_permission = {"mcp_servers": ["github-mcp"]} - with patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - west_mgr, - ): - await validate_key_mcp_servers_against_team( - object_permission=object_permission, - team_obj=team_obj, - ) + _drop_stale_object_permission_mcp_servers(object_permission, {"github-mcp": {"west-id"}}) assert object_permission["mcp_servers"] == ["github-mcp"] from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager From 190b8c7d8e18f1edd642727e1754c157bab5645b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 09:03:47 -0700 Subject: [PATCH 50/92] test(e2e): drop coverage registry cell for the alias-grant test --- tests/e2e/coverage_registry/mcp.yaml | 8 -------- tests/e2e/mcp/test_mcp_key_access_e2e.py | 1 - 2 files changed, 9 deletions(-) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 7ed2a7ab2ad..ab644118a47 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -15,14 +15,6 @@ assertions: [access_group_scoped] source: "test_mcp_access_group_e2e.py" rationale: "A key granted an MCP access group sees the tagged server's tools; a key with a different group does not. Access-group-scoped tool selection at key creation" -- id: mcp.list_tools.api_key.alias_grant_persists - module: mcp - tier: P1 - operation: list_tools - auth_family: api_key - assertions: [alias_grant_persists] - source: "object_permission_utils.py validate_key_mcp_servers_against_team" - rationale: "A key granted an MCP server by alias keeps the alias verbatim in its stored object_permission (shared-DB multi-region instances each resolve it to their local server id at read time) and still lists the server's tools" - id: mcp.list_tools.api_key.denied_without_permission module: mcp tier: P0 diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 7fe40795f1c..68005ae3f6a 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -31,7 +31,6 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str class TestMcpKeyGrantByAlias: - @pytest.mark.covers("mcp.list_tools.api_key.alias_grant_persists") def test_alias_grant_persists_verbatim_and_lists_tools( self, client: McpClient, From 2757399c99d45e3b9ae1143a9bd3d3a981bfdfae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 09:43:47 -0700 Subject: [PATCH 51/92] fix(ui): render the skill detail page with theme tokens The page painted every surface, border and text color inline with a fixed light palette (#202124, #5f6368, #dadce0, #f8f9fa, #fff), so in dark mode it drew dark text on hardcoded white cards. Move the whole component to the foreground/muted/border/card/info tokens, which already resolve for both themes. --- .../claude_code_plugins/skill_detail.tsx | 310 +++++------------- 1 file changed, 78 insertions(+), 232 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 25b35c34861..809272c2ac0 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; import { ArrowLeft, Check, Copy, Link2 } from "lucide-react"; +import { cn } from "@/lib/cva.config"; import { buildMarketplaceSettingsSnippet, formatInstallCommand } from "./helpers"; import { Plugin } from "./types"; @@ -50,48 +51,37 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { ]; return ( -

+
{/* Back link */}
Skills
{/* Header */} -
-

{skill.name}

+
+

{skill.name}

{skill.description && ( -

{skill.description}

+

{skill.description}

)}
{/* Tab bar */} -
-
+
+
{tabs.map((tab) => (
setActiveTab(tab.key)} - style={{ - padding: "12px 20px", - fontSize: 14, - color: activeTab === tab.key ? "#1a73e8" : "#5f6368", - borderBottom: activeTab === tab.key ? "3px solid #1a73e8" : "3px solid transparent", - cursor: "pointer", - fontWeight: activeTab === tab.key ? 500 : 400, - marginBottom: -1, - }} + className={cn( + "-mb-px cursor-pointer border-b-[3px] px-5 py-3 text-sm", + activeTab === tab.key + ? "border-info font-medium text-info" + : "border-transparent font-normal text-muted-foreground", + )} > {tab.label}
@@ -101,27 +91,23 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Overview tab */} {activeTab === "overview" && ( -
+
{/* Left column */} -
-

Skill Details

-

Metadata registered with this skill

- +
+

Skill Details

+

Metadata registered with this skill

+
- - - + + + {detailRows.map((row, i) => ( - - - + + + ))} @@ -129,38 +115,27 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Right sidebar */} -
-
-
Status
+
+
+
Status
{skill.enabled ? "Public" : "Draft"}
{sourceUrl && ( -
-
Source
+
+
Source
{sourceUrl.replace("https://", "")} @@ -169,20 +144,13 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { )} {skill.keywords && skill.keywords.length > 0 && ( -
-
Tags
-
+
+
Tags
+
{skill.keywords.map((kw) => ( {kw} @@ -192,10 +160,8 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { )}
-
Skill ID
-
- {skill.id} -
+
Skill ID
+
{skill.id}
@@ -203,93 +169,43 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* How to Use tab */} {activeTab === "usage" && ( -
-

Using this skill

-

+

+

Using this skill

+

Once your proxy is set as a marketplace, enable this skill in Claude Code with one command:

{/* Install command */} -
-
- Run in Claude Code +
+
+ Run in Claude Code
- 
-              {installCommand}
-
+ 
{installCommand}
{/* Shown when the marketplace catalog is stale and the plugin isn't found yet */} -
-

+

+

If you see "Plugin {skill.name} not found in marketplace", update the catalog first:

- 
+            
               /plugin marketplace update litellm
-

+

Don't have the marketplace configured yet?{" "} - setActiveTab("setup")} style={{ color: "#1a73e8", cursor: "pointer" }}> + setActiveTab("setup")} className="cursor-pointer text-info"> See one-time setup →

@@ -298,126 +214,56 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Setup tab (linked from usage) */} {activeTab === "setup" && ( -
-

- One-time marketplace setup -

+
+

One-time marketplace setup

{/* Option 1: single command — fastest path for most users */} -

+

Run this command in Claude Code to register the marketplace:

-
-
- Run in Claude Code +
+
+ Run in Claude Code
- 
+            
               {`/plugin marketplace add ${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`}
{/* Option 2: settings.json — for persistent config or managed deployments. extraKnownMarketplaces requires source to be a nested object, not a flat string. */} -

- Or add this to{" "} - - ~/.claude/settings.json - {" "} +

+ Or add this to ~/.claude/settings.json{" "} for a persistent configuration:

-
-
- ~/.claude/settings.json +
+
+ ~/.claude/settings.json
- 
-              {settingsSnippet}
-
+ 
{settingsSnippet}
)} From a5f941068144ff7f0acc3838e2ace35202f74ce2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 09:46:43 -0700 Subject: [PATCH 52/92] test(ui): wait for the select popup before clicking its option key_edit_view opened a select and then clicked the option it found by title text or by raw text. Both queries match the moment the option enters the DOM, which is one render before the popup finishes entering. Until then the positioner still carries an inline pointer-events: none, and user-event refuses to click through it. That is a race, and a fast machine loses it. Five of the file's 84 tests failed on every local run while CI stayed green, which is the worst shape for a test to have: it is only ever red on the machine of whoever is trying to change the code. tests/test-utils.tsx already ships chooseSelectOption for exactly this. It finds the option by role and waits for the positioner to release pointer events before clicking. The five call sites now use it, and the helper takes the direct user-event API as well as a setup() instance so callers do not have to restructure to use it. Five consecutive full-file runs pass where every previous run failed. Also finishes this file's screen queries, which brings prefer-screen-queries to its target of 18. --- ui/litellm-dashboard/eslint-budgets.json | 2 +- .../templates/key_edit_view.test.tsx | 31 +++++++------------ ui/litellm-dashboard/tests/test-utils.tsx | 2 +- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index e7545eb2383..e8207d179bd 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -7,5 +7,5 @@ "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 716, "target": 500 }, - "testing-library/prefer-screen-queries": { "max": 21, "target": 18 } + "testing-library/prefer-screen-queries": { "max": 18, "target": 18 } } diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index bf8f43b8b5f..97b09e00808 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1,7 +1,7 @@ import { fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "../key_team_helpers/ModelMaxBudgetEditor"; import { @@ -300,7 +300,7 @@ describe("KeyEditView", () => { }); it("should render", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -313,12 +313,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("Save Changes")).toBeInTheDocument(); + expect(screen.getByText("Save Changes")).toBeInTheDocument(); }); }); it("should render tags", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -331,12 +331,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("test-tag")).toBeInTheDocument(); + expect(screen.getByText("test-tag")).toBeInTheDocument(); }); }); it("should not render tags in metadata textarea", async () => { - const { getByLabelText } = renderWithProviders( + renderWithProviders( {}} @@ -348,7 +348,7 @@ describe("KeyEditView", () => { />, ); - const metadataTextarea = getByLabelText("Metadata") as HTMLTextAreaElement; + const metadataTextarea = screen.getByLabelText("Metadata") as HTMLTextAreaElement; await waitFor(() => { expect(metadataTextarea).toHaveValue("{}"); }); @@ -963,10 +963,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - - const weeklyOption = await screen.findByText("weekly"); - await userEvent.click(weeklyOption); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "weekly"); const submitButton = screen.getByRole("button", { name: /save changes/i }); await userEvent.click(submitButton); @@ -1042,8 +1039,7 @@ describe("KeyEditView", () => { ); const resetBudget = await screen.findByLabelText("Reset Budget"); - await userEvent.click(resetBudget); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, resetBudget, "Never resets"); await waitFor(() => { expect(resetBudget).toHaveTextContent("Never resets"); @@ -1074,8 +1070,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "Never resets"); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -1946,8 +1941,7 @@ describe("KeyEditView", () => { await userEvent.clear(duration); await userEvent.type(duration, "45d"); - await userEvent.click(screen.getByLabelText(/TPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/TPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -2103,8 +2097,7 @@ describe("KeyEditView", () => { renderForPayload(onSubmitMock); await screen.findByRole("button", { name: /save changes/i }); - await userEvent.click(screen.getByLabelText(/RPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/RPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index 66966201a9c..162b8a3df7b 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -52,7 +52,7 @@ const pointerBlocked = (element: HTMLElement): boolean => { * the option text alone is a race that React 19's flush timing loses. */ export const chooseSelectOption = async ( - user: ReturnType, + user: Pick, "click">, trigger: HTMLElement, optionName: string | RegExp, ) => { From 7abed91523f8a1bfb79697949f070e367a1069c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:11:25 -0700 Subject: [PATCH 53/92] feat(cost): support day-of-week qualified off-peak windows --- .../litellm_core_utils/llm_cost_calc/utils.py | 108 ++++++++++++-- litellm/types/utils.py | 27 +++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 136 ++++++++++++++++++ 3 files changed, 259 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 21680129ed4..c968fac254e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -4,9 +4,10 @@ import re from collections.abc import Mapping, Sequence from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timezone, tzinfo from types import MappingProxyType from typing import Any, Final, Literal, TypedDict, cast +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm from litellm._logging import verbose_logger @@ -321,6 +322,99 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ return False +_WEEKDAY_NUMBERS: Final = MappingProxyType( + { + "mon": 1, + "monday": 1, + "tue": 2, + "tues": 2, + "tuesday": 2, + "wed": 3, + "wednesday": 3, + "thu": 4, + "thur": 4, + "thurs": 4, + "thursday": 4, + "fri": 5, + "friday": 5, + "sat": 6, + "saturday": 6, + "sun": 7, + "sunday": 7, + } +) + + +def _normalize_weekday(value: object) -> int | None: + if isinstance(value, bool): + return None + if isinstance(value, int): + return value if 1 <= value <= 7 else None + if isinstance(value, str): + return _WEEKDAY_NUMBERS.get(value.strip().lower()) + return None + + +def _weekday_calendar(weekday_timezone: object) -> tzinfo: + if isinstance(weekday_timezone, str) and weekday_timezone.strip(): + try: + return ZoneInfo(weekday_timezone.strip()) + except (ValueError, ZoneInfoNotFoundError): + return timezone.utc + return timezone.utc + + +def _matches_weekdays(reference_utc: datetime, weekdays: object, weekday_timezone: object) -> bool: + """Return True when reference_utc falls on one of the rule's weekdays, read on the calendar + named by weekday_timezone (default UTC). An absent weekdays means every day. The calendar + matters even when UTC and vendor-local weekdays agree at every currently priced hour: a + window past 16:00 UTC is where an Asia/Shanghai weekday diverges from the UTC one. + """ + if weekdays is None: + return True + if isinstance(weekdays, str) or not isinstance(weekdays, Sequence): + return False + allowed: Final = frozenset(day for day in map(_normalize_weekday, weekdays) if day is not None) + return reference_utc.astimezone(_weekday_calendar(weekday_timezone)).isoweekday() in allowed + + +def _as_window_strings(value: object) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, Sequence): + return tuple(entry for entry in value if isinstance(entry, str)) + return () + + +def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = None) -> bool: + """Return True when current_time (UTC, defaulting to now) is off-peak under the block's + rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose + hours apply only on its weekdays. + """ + reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference_utc: Final = ( + reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc) + ) + flat_windows: Final = _as_window_strings(off_peak.get("hours_utc")) + if flat_windows and _is_within_off_peak_window(flat_windows, reference_utc): + return True + windows: Final = off_peak.get("windows") + if isinstance(windows, str) or not isinstance(windows, Sequence): + return False + weekday_timezone: Final = off_peak.get("weekday_timezone") + for rule in windows: + if not isinstance(rule, Mapping): + continue + rule_windows = _as_window_strings(rule.get("hours_utc")) + if not rule_windows: + continue + if not _matches_weekdays(reference_utc, rule.get("weekdays"), weekday_timezone): + continue + if _is_within_off_peak_window(rule_windows, reference_utc): + return True + return False + + def _coerce_off_peak_rate(value: object, default: float) -> float: if isinstance(value, bool): return default @@ -342,16 +436,14 @@ def _apply_off_peak_pricing( cache_read_cost: float, ) -> tuple[float, float, float]: """Swap in off-peak per-token rates when the current UTC time is inside one of the model's - off_peak_pricing windows. An off-peak rate replaces the rate that would otherwise apply - rather than discounting it, so a model that also has tiered or above-threshold pricing bills - the flat off-peak rate for the whole request while the window is open. Any rate left unset in + off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in + windows. An off-peak rate replaces the rate that would otherwise apply rather than + discounting it, so a model that also has tiered or above-threshold pricing bills the flat + off-peak rate for the whole request while the window is open. Any rate left unset in off_peak_pricing falls back to the standard rate. """ off_peak: Final = model_info.get("off_peak_pricing") - if not off_peak: - return prompt_base_cost, completion_base_cost, cache_read_cost - hours_utc: Final = off_peak.get("hours_utc") - if not hours_utc or not _is_within_off_peak_window(hours_utc, current_time): + if not off_peak or not _is_off_peak(off_peak, current_time): return prompt_base_cost, completion_base_cost, cache_read_cost return ( _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3ab3a2382dc..17051714c25 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -193,14 +193,33 @@ class AgenticLoopParams(TypedDict, total=False): """The LLM provider name (e.g., 'bedrock', 'anthropic')""" -class OffPeakPricing(TypedDict, total=False): - """Time-windowed off-peak rates for providers that discount by time of day (e.g. DeepSeek). +class OffPeakWindow(TypedDict, total=False): + """One off-peak rule: UTC time-of-day windows, optionally restricted to weekdays. - hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them for multiple daily windows; - a window may wrap past midnight. Any rate left unset falls back to the standard rate. + hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them; a window may wrap past + midnight and an equal-ended window covers the whole day. weekdays is a list of days the + rule applies on, as ISO-8601 numbers (1 = Monday .. 7 = Sunday) or English day names; + omitted means every day. The weekday is read on the calendar named by the block's + weekday_timezone. """ hours_utc: ReadOnly[str | Sequence[str]] + weekdays: ReadOnly[Sequence[int | str]] + + +class OffPeakPricing(TypedDict, total=False): + """Time-windowed off-peak rates for providers that discount by time of day (e.g. DeepSeek). + + hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them for multiple daily windows, + applying on every day of the week; a window may wrap past midnight. windows adds + day-of-week-qualified rules (e.g. weekend-only whole-day off-peak), matched as a union + with hours_utc. weekday_timezone names the IANA calendar weekdays are read on, defaulting + to UTC. Any rate left unset falls back to the standard rate. + """ + + hours_utc: ReadOnly[str | Sequence[str]] + windows: ReadOnly[Sequence[OffPeakWindow]] + weekday_timezone: ReadOnly[str] input_cost_per_token: ReadOnly[float] output_cost_per_token: ReadOnly[float] cache_read_input_token_cost: ReadOnly[float] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index c9ed5936f03..501a9314c90 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, + _is_off_peak, _is_within_off_peak_window, calculate_cache_writing_cost, generic_cost_per_token, @@ -479,6 +480,141 @@ def test_is_within_off_peak_window_malformed_returns_false(): assert _is_within_off_peak_window("25:00-26:00", now) is False +def test_is_off_peak_weekday_qualified_windows_deepseek_schedule(): + """DeepSeek since 2026-08-23: peak is 01:00-04:00 and 06:00-10:00 UTC on weekdays only, with + weekends off-peak around the clock. The weekday axis is not a filter on one window set; on + two days of seven the off-peak window becomes the whole day, so the schedule needs two + day-qualified rules. The weekend instants inside would-be peak hours are the ones a + time-only implementation bills wrong.""" + from datetime import datetime, timezone + + deepseek = { + "windows": [ + {"hours_utc": ["00:00-01:00", "04:00-06:00", "10:00-00:00"], "weekdays": [1, 2, 3, 4, 5]}, + {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, + ], + } + peak_instants = [ + datetime(2026, 8, 24, 1, 30, tzinfo=timezone.utc), + datetime(2026, 8, 26, 7, 0, tzinfo=timezone.utc), + datetime(2026, 8, 28, 9, 59, tzinfo=timezone.utc), + ] + off_peak_instants = [ + datetime(2026, 8, 23, 1, 30, tzinfo=timezone.utc), + datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc), + datetime(2026, 8, 30, 8, 0, tzinfo=timezone.utc), + datetime(2026, 8, 26, 5, 0, tzinfo=timezone.utc), + datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc), + datetime(2026, 8, 24, 0, 30, tzinfo=timezone.utc), + ] + for when in peak_instants: + assert _is_off_peak(deepseek, when) is False, f"{when.isoformat()} should bill peak" + for when in off_peak_instants: + assert _is_off_peak(deepseek, when) is True, f"{when.isoformat()} should bill off-peak" + + +def test_is_off_peak_weekday_timezone_reads_vendor_calendar(): + """The UTC and Asia/Shanghai calendars only disagree about the date over 16:00-24:00 UTC, so + a window in that stretch is the one place a vendor-local weekday differs from a UTC one: + 2026-08-28T16:30Z is Friday in UTC but already Saturday in Beijing.""" + from datetime import datetime, timezone + + shanghai_saturday = { + "weekday_timezone": "Asia/Shanghai", + "windows": [{"hours_utc": "16:00-17:00", "weekdays": [6]}], + } + assert _is_off_peak(shanghai_saturday, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(shanghai_saturday, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_weekdays_default_utc_calendar_and_accept_names(): + from datetime import datetime, timezone + + named_weekend = {"windows": [{"hours_utc": "00:00-00:00", "weekdays": ["Sat", "sunday"]}]} + assert _is_off_peak(named_weekend, datetime(2026, 8, 29, 12, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(named_weekend, datetime(2026, 8, 28, 12, 0, tzinfo=timezone.utc)) is False + + utc_friday = {"windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(utc_friday, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(utc_friday, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_naive_current_time_read_as_utc(): + from datetime import datetime + + block = {"windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(block, datetime(2026, 8, 28, 16, 30)) is True + assert _is_off_peak(block, datetime(2026, 8, 29, 16, 30)) is False + + +def test_is_off_peak_invalid_weekday_timezone_falls_back_to_utc(): + from datetime import datetime, timezone + + block = {"weekday_timezone": "Not/AZone", "windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(block, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_ignores_malformed_weekday_rules(): + from datetime import datetime, timezone + + when = datetime(2026, 8, 29, 12, 0, tzinfo=timezone.utc) + assert _is_off_peak({"windows": [{"hours_utc": "00:00-00:00", "weekdays": []}]}, when) is False + assert _is_off_peak({"windows": [{"hours_utc": "00:00-00:00", "weekdays": [0, 8, "noday", True]}]}, when) is False + assert _is_off_peak({"windows": [{"weekdays": [6]}]}, when) is False + assert _is_off_peak({"windows": [{"hours_utc": 1630}]}, when) is False + assert _is_off_peak({"windows": ["00:00-00:00"]}, when) is False + assert _is_off_peak({"windows": "00:00-00:00"}, when) is False + assert _is_off_peak({"hours_utc": 1630}, when) is False + assert _is_off_peak({}, when) is False + + +def test_is_off_peak_flat_hours_and_windows_are_a_union(): + from datetime import datetime, timezone + + block = { + "hours_utc": "04:00-06:00", + "windows": [{"hours_utc": "00:00-00:00", "weekdays": [7]}], + } + assert _is_off_peak(block, datetime(2026, 8, 28, 5, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 30, 20, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 28, 20, 0, tzinfo=timezone.utc)) is False + + +def test_get_token_base_cost_weekend_only_off_peak_rate(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": { + "windows": [ + {"hours_utc": ["00:00-01:00", "04:00-06:00", "10:00-00:00"], "weekdays": [1, 2, 3, 4, 5]}, + {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, + ], + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + saturday_peak_hours = _get_token_base_cost( + model_info, usage, current_time=datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc) + ) + assert saturday_peak_hours[:2] == (5e-7, 1e-6) + + monday_same_hours = _get_token_base_cost( + model_info, usage, current_time=datetime(2026, 8, 24, 2, 0, tzinfo=timezone.utc) + ) + assert monday_same_hours[:2] == (1e-6, 2e-6) + + def test_get_token_base_cost_applies_off_peak_pricing(): from datetime import datetime, timezone from typing import cast From 06d4521fc0b7a69aa0db0d6ba1aa54dc4ab1af28 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:15:18 -0700 Subject: [PATCH 54/92] fix(registry): add vertex veo 3.1 resolution tier pricing per vertex pricing page --- ...odel_prices_and_context_window_backup.json | 10 ++++-- model_prices_and_context_window.json | 10 ++++-- tests/test_litellm/test_video_generation.py | 35 +++++++++++++++++++ 3 files changed, 51 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dd45495884d..2cfa1f93f79 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -43509,7 +43509,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43523,6 +43524,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" @@ -43538,7 +43541,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43553,6 +43557,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dd45495884d..2cfa1f93f79 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -43509,7 +43509,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43523,6 +43524,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" @@ -43538,7 +43541,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43553,6 +43557,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index b166e902d6e..2a60ff9c4b5 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,41 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): + """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, provider: str, resolution: str | None, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": duration, + **({"video_resolution": resolution} if resolution else {}), + } + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider=provider, + ) + + for provider in ("gemini", "vertex_ai"): + for suffix in ("generate-preview", "generate-001"): + standard = f"{provider}/veo-3.1-{suffix}" + fast = f"{provider}/veo-3.1-fast-{suffix}" + assert abs(cost_for(standard, provider, None, 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "1080p", 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "4k", 8.0) - 4.8) < 1e-6 + assert abs(cost_for(fast, provider, "720p", 8.0) - 0.8) < 1e-6 + assert abs(cost_for(fast, provider, "1080p", 8.0) - 0.96) < 1e-6 + assert abs(cost_for(fast, provider, "4k", 8.0) - 2.4) < 1e-6 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() From 8acdb9208756fb306680878d06f53bbd75c2f04a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 10:27:43 -0700 Subject: [PATCH 55/92] fix(ui): keep the skill detail copy buttons transparent bg-none only clears background-image, so the buttons fell back to the browser's default button background instead of the transparent one the inline style had. --- .../src/components/claude_code_plugins/skill_detail.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 809272c2ac0..7cbbb08b623 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -182,7 +182,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => {
- Property - - {skill.name} -
Property{skill.name}
{row.property}{row.value}
{row.property}{row.value}
+
- - - + + + {detailRows.map((row, i) => ( - - - + + + ))} @@ -122,37 +116,30 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* Right column — metadata sidebar like Vertex */} -
+
{/* Guardrail ID */} -
-
Guardrail ID
-
litellm/{card.id}
+
+
Guardrail ID
+
litellm/{card.id}
{/* Type */} -
-
Type
-
+
+
Type
+
{card.category === "litellm" ? "Content Filter" : "Partner"}
{/* Tags — pill style like Vertex */} {card.tags.length > 0 && ( -
-
Tags
-
+
+
Tags
+
{card.tags.map((tag) => ( {tag} @@ -166,19 +153,19 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {activeTab === "eval" && (
-

Eval Results

-
- Property - - {card.name} -
Property{card.name}
{row.property}{row.value}
{row.property}{row.value}
+

Eval Results

+
- - - + + + {evalRows.map((row, i) => ( - - - + + + ))} From eb53639ecbd7af0c59e1ef225af679c973da1051 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 12:42:27 -0700 Subject: [PATCH 88/92] test(ui): pick select options by role instead of by text Clicking a Base UI select entry found by text or by a title attribute is a race. The text node exists one render before the popup finishes entering, and until then the positioner still carries pointer-events: none, so user-event refuses the click and the test throws. Querying by role only matches once the popup is exposed to the accessibility tree, which is after that window closes. Route the 37 remaining select interactions through chooseSelectOption, which does the role query. Instrumenting the converted files shows the text query resolving while the popup was still pointer-blocked on 6 of 41 samples; the role query was never blocked. Seven files kept their text queries because their popup entries carry no accessible role, so there is nothing to query by. --- .../add_agent_form.integration.test.tsx | 4 ++-- .../budget_modal.integration.test.tsx | 7 +++---- .../edit_budget_modal.integration.test.tsx | 4 ++-- .../CoordinationRedisTypeSelector.test.tsx | 4 ++-- .../_components/ShadowEvalSection.test.tsx | 10 ++++------ .../_components/ToolTestPanel.test.tsx | 4 ++-- .../_components/policy_test_panel.test.tsx | 4 ++-- .../prompts/_components/index.test.tsx | 10 ++++------ .../components/TeamsPage/TeamsTable.test.tsx | 5 ++--- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 11 ++++------- .../add_pass_through.integration.test.tsx | 5 ++--- .../KeyLifecycleSettings.test.tsx | 17 ++++++----------- .../common_components/ModelSelector.test.tsx | 4 ++-- .../shared/DataTable/DataTable.test.tsx | 16 ++++++---------- .../DataTable/DataTableSortHeader.test.tsx | 10 ++++------ .../shared/PaginatedSearchSelect.test.tsx | 10 ++++------ .../src/components/shared/SearchSelect.test.tsx | 4 ++-- .../view_logs/RequestLogsFilters.test.tsx | 5 ++--- ui/litellm-dashboard/tests/test-utils.tsx | 13 ++++++++----- 19 files changed, 63 insertions(+), 84 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index bffc7b0fbc8..dad8599e967 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -5,6 +5,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import AddAgentForm from "./add_agent_form"; import * as networking from "@/components/networking"; import type { AgentCreateInfo } from "@/components/networking"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ createAgentCall: vi.fn(), @@ -309,8 +310,7 @@ describe("AddAgentForm submit payload", () => { await user.type(await screen.findByLabelText("Allowed Models"), "gpt-4o,"); await user.keyboard("{Escape}"); - await user.click(screen.getByLabelText("Allowed Agents (Sub-Agents)")); - await user.click(await screen.findByTitle("Sub Agent One")); + await chooseSelectOption(user, screen.getByLabelText("Allowed Agents (Sub-Agents)"), "Sub Agent One"); await user.keyboard("{Escape}"); await user.click(screen.getByText(/Configure which models, agents, and MCP tools/)); await user.click(screen.getByRole("button", { name: /^Next/ })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx index bc920e55abd..17a297d203d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx @@ -4,6 +4,7 @@ import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import BudgetModal from "./budget_modal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); @@ -63,8 +64,7 @@ describe("BudgetModal", () => { await openOptionalSettings(user); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await create(user); @@ -80,8 +80,7 @@ describe("BudgetModal", () => { await openOptionalSettings(user); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await user.click(screen.getByText("Optional Settings")); await waitFor(() => expect(screen.queryByLabelText("Max Budget (USD)")).not.toBeInTheDocument()); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx index 3fa96b54f1d..fe0ecfc10dd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx @@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { components } from "@/lib/http/schema"; import EditBudgetModal from "./edit_budget_modal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const { updateMock } = vi.hoisted(() => ({ updateMock: vi.fn() })); @@ -73,8 +74,7 @@ describe("EditBudgetModal", () => { await user.clear(screen.getByLabelText("Max Budget (USD)")); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await save(user); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx index 76287e6e724..563c684237a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx @@ -5,6 +5,7 @@ import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import CoordinationRedisTypeSelector from "./CoordinationRedisTypeSelector"; import { COORDINATION_REDIS_TYPE_DESCRIPTIONS } from "./coordinationRedisFields"; +import { chooseSelectOption } from "../../../../../../tests/test-utils"; describe("CoordinationRedisTypeSelector", () => { it("labels the control and shows the current selection", () => { @@ -36,8 +37,7 @@ describe("CoordinationRedisTypeSelector", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Cluster")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Cluster"); expect(onTypeChange).toHaveBeenCalledTimes(1); expect(onTypeChange.mock.calls[0][0]).toBe("cluster"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 64e03985f57..a1de608d0bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -97,6 +97,7 @@ import { useStopShadowEval, type ShadowEvalJob, } from "./useShadowEval"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const job = (overrides: Partial = {}): ShadowEvalJob => ({ job_id: "job-1", @@ -437,8 +438,7 @@ describe("ShadowEvalSection", () => { await user.click(within(keyList).getByText("prod-alpha")); await user.click(keyInput); await user.click(within(keyList).getByText("staging-beta")); - await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); expect(screen.getByText("Start shadow eval")).toBeDisabled(); @@ -470,8 +470,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search teams by alias")); const teamList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(teamList).getByText("engineering")); - await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); await user.click(screen.getByText("Start shadow eval")); @@ -502,8 +501,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search keys by alias")); const keyList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(keyList).getByText("prod-alpha")); - await user.click(screen.getByPlaceholderText("Select up to 4 auto-routers")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx index 26e66c8338c..81c55e7982a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx @@ -6,6 +6,7 @@ import { describe, expect, it, vi } from "vitest"; import { ToolTestPanel } from "./ToolTestPanel"; import { InputSchema, MCPTool } from "@/components/mcp_tools/types"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const buildTool = (schema: InputSchema | string): MCPTool => ({ name: "demo-tool", @@ -229,8 +230,7 @@ describe("ToolTestPanel argument payload", () => { const onSubmit = await submitPanel( { type: "object", properties: { active: { type: "boolean", default: false } } }, async (user) => { - await user.click(screen.getByLabelText("active")); - await user.click(await screen.findByText("True")); + await chooseSelectOption(user, screen.getByLabelText("active"), "True"); }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx index c8779be06c4..f23145a917a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx @@ -5,6 +5,7 @@ import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event" import { renderWithProviders } from "@/../tests/test-utils"; import * as networking from "@/components/networking"; import PolicyTestPanel from "./policy_test_panel"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; vi.mock("@/components/networking"); @@ -24,8 +25,7 @@ const setup = () => { }; const pickOption = async (user: ReturnType, label: string, option: string) => { - await user.click(screen.getByLabelText(label)); - await user.click(await screen.findByTitle(option)); + await chooseSelectOption(user, screen.getByLabelText(label), option); }; const simulate = async (user: ReturnType) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx index d6a4aaea2e1..3b0ee3a3ce2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { deletePromptCall, getPromptsList } from "@/components/networking"; import PromptsPanel from "./index"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ getPromptsList: vi.fn(), @@ -118,8 +119,7 @@ describe("PromptsPanel toolbar", () => { expect(screen.getByText("All Environments")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Production")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Production"); await waitFor(() => expect(mockGetPromptsList).toHaveBeenLastCalledWith("sk-test", "production")); }); @@ -131,12 +131,10 @@ describe("PromptsPanel toolbar", () => { renderPanel("Admin"); await screen.findByText("table-loaded"); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Production")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Production"); await waitFor(() => expect(screen.getByRole("combobox")).toHaveTextContent("Production")); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("All Environments")); + await chooseSelectOption(user, screen.getByRole("combobox"), "All Environments"); await waitFor(() => expect(screen.getByRole("combobox")).toHaveTextContent("All Environments")); await waitFor(() => expect(mockGetPromptsList).toHaveBeenLastCalledWith("sk-test", undefined)); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx index 9375673ba10..d3b9a55a7d5 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -2,7 +2,7 @@ import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, MockedFunction, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { Team } from "../key_team_helpers/key_list"; import { TeamsResponse, useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams"; import { TeamsTable } from "./TeamsTable"; @@ -243,8 +243,7 @@ describe("row actions", () => { await user.click(await screen.findByText("Edit team")); expect(onEditTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); - await user.click(screen.getByTestId("team-actions-team-1")); - await user.click(await screen.findByText("Delete team")); + await chooseSelectOption(user, screen.getByTestId("team-actions-team-1"), "Delete team", "menuitem"); expect(onDeleteTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index d09ccbf8826..b06448912d0 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -2,7 +2,7 @@ import { screen, waitFor, within, fireEvent } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; @@ -327,8 +327,7 @@ it("sorts by the backend max_budget field when 'Budget descending' is chosen fro const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith( @@ -343,8 +342,7 @@ it("emphasizes the active field in the Spend / Budget header so the sorted colum const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); await waitFor(() => { expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" })).toHaveClass("font-semibold"); @@ -356,8 +354,7 @@ it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / B const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Spend ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Spend ascending", "menuitem"); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "spend", sortOrder: "asc" })); diff --git a/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx b/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx index df4dbff4111..daaacd7f0e2 100644 --- a/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen, waitFor } from "../../tests/test-utils"; +import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "../../tests/test-utils"; import AddPassThroughEndpoint from "./add_pass_through"; const createPassThroughEndpoint = vi.fn(); @@ -153,8 +153,7 @@ describe("add_pass_through submit payload", () => { await openModal(user); await fillRequiredFields(user); - await user.click(screen.getByLabelText(/HTTP Methods/)); - await user.click(await screen.findByTitle("POST")); + await chooseSelectOption(user, screen.getByLabelText(/HTTP Methods/), "POST"); await submit(user); diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx index 6e627e3b45c..df21444cbad 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { Controller, useForm } from "react-hook-form"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; -import { fireEvent, renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import KeyLifecycleSettings from "./KeyLifecycleSettings"; const CREATE_PLACEHOLDER = "e.g., 30d or leave empty to never expire"; @@ -167,8 +167,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("90 days")); + await chooseSelectOption(user, screen.getByRole("combobox"), "90 days"); await waitFor(() => expect(screen.getAllByTitle("90 days").some(isRenderedSelection)).toBe(true)); expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("90d"); @@ -181,8 +180,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); expect(await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).toBeInTheDocument(); expect(screen.getByText("Supported formats: seconds (s), minutes (m), hours (h), days (d)")).toBeInTheDocument(); @@ -196,8 +194,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); fireEvent.change(customInput, { target: { value: "14d" } }); @@ -213,14 +210,12 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); fireEvent.change(customInput, { target: { value: "14d" } }); await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("14d")); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("7 days")); + await chooseSelectOption(user, screen.getByRole("combobox"), "7 days"); await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("7d")); expect(screen.queryByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx index bd7b56f1817..7a5db4667e9 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx @@ -2,6 +2,7 @@ import { act, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import ModelSelector from "./ModelSelector"; +import { chooseSelectOption } from "../../../tests/test-utils"; vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([]), @@ -9,8 +10,7 @@ vi.mock("@/components/llm_calls/fetch_models", () => ({ const openCustomModelInput = async () => { const user = userEvent.setup(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Enter custom model")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Enter custom model"); return screen.getByPlaceholderText("Enter custom model name"); }; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 7ead9bb64d4..f502ea77127 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -7,6 +7,7 @@ import { describe, expect, it, vi } from "vitest"; import { DataTable } from "./DataTable"; import { DataTableMultiSortHeader, DataTableSortHeader } from "./DataTableSortHeader"; import { DataTableViewOptions } from "./DataTableViewOptions"; +import { chooseSelectOption } from "../../../../tests/test-utils"; interface Person { id: string; @@ -176,16 +177,13 @@ describe("DataTable sorting", () => { const user = userEvent.setup(); render(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Ascending", "menuitem"); expect(names()).toEqual(["Alice", "Bob", "Charlie"]); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Descending", "menuitem"); expect(names()).toEqual(["Charlie", "Bob", "Alice"]); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Reset")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Reset", "menuitem"); expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); @@ -202,12 +200,10 @@ describe("DataTable sorting", () => { />, ); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "max_budget", desc: true }]); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Spend ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Spend ascending", "menuitem"); expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "spend", desc: false }]); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx index 70403ba2efe..aeb6fa44e91 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx @@ -13,6 +13,7 @@ import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; import { DataTableSortHeader, type DataTableSortVariant } from "./DataTableSortHeader"; +import { chooseSelectOption } from "../../../../tests/test-utils"; interface Item { name: string; @@ -84,16 +85,13 @@ describe("DataTableSortHeader", () => { const user = userEvent.setup(); render(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Descending", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="desc"]')).not.toBeNull(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Ascending", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="asc"]')).not.toBeNull(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Reset")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Reset", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="none"]')).not.toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx index d248cf311cc..04c8d3e9012 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx @@ -5,6 +5,7 @@ import { describe, expect, it, vi } from "vitest"; import { PaginatedSearchSelect } from "./PaginatedSearchSelect"; import type { SearchSelectOption } from "./SearchSelect"; +import { chooseSelectOption } from "../../../tests/test-utils"; const OPTIONS: SearchSelectOption[] = [ { label: "alias-alpha", value: "alias-alpha" }, @@ -63,8 +64,7 @@ describe("PaginatedSearchSelect", () => { } render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("alias-beta")); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); expect(screen.getByRole("combobox")).toHaveValue("alias-beta"); await new Promise((resolve) => setTimeout(resolve, 400)); @@ -136,8 +136,7 @@ describe("PaginatedSearchSelect", () => { const onValueChange = vi.fn(); renderSelect({ onValueChange }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("alias-beta")); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); expect(onValueChange).toHaveBeenCalledWith("alias-beta"); }); @@ -255,8 +254,7 @@ describe("PaginatedSearchSelect", () => { } render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Beta Team")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Beta Team"); await user.click(screen.getByRole("button", { name: "refetch" })); expect(screen.getByRole("combobox")).toHaveValue("Beta Team"); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx index c0b9bc7d8ae..62a510268dd 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import { SearchSelect } from "./SearchSelect"; +import { chooseSelectOption } from "../../../tests/test-utils"; const OPTIONS = [ { label: "Acme Prod", value: "team-1" }, @@ -71,8 +72,7 @@ describe("SearchSelect", () => { const onValueChange = vi.fn(); const user = userEvent.setup(); render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Growth")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Growth"); expect(onValueChange).toHaveBeenCalledWith("team-2"); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index e0ec66e1ae2..5a35c7ae16b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import { ERROR_CODE_OPTIONS } from "./constants"; import { LOG_FILTER_IDS } from "./log_filter_logic"; import { RequestLogsFilters } from "./RequestLogsFilters"; @@ -125,8 +125,7 @@ describe("RequestLogsFilters", () => { const user = userEvent.setup(); const { set } = renderFilters(); - await user.click(await screen.findByPlaceholderText("Search an internal user")); - await user.click(await screen.findByText("alice@example.com")); + await chooseSelectOption(user, await screen.findByPlaceholderText("Search an internal user"), "alice@example.com"); expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "alice@example.com"); }); diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index 162b8a3df7b..31af51f3f32 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -45,19 +45,22 @@ const pointerBlocked = (element: HTMLElement): boolean => { }; /** - * Opens a Base UI Select and picks an option by its accessible name. + * Opens a Base UI popup and picks an entry by its accessible name. * - * The option is in the DOM one render before the popup finishes entering, and until then its - * positioner still carries `pointer-events: none`, which user-event refuses to click. Waiting on - * the option text alone is a race that React 19's flush timing loses. + * Querying the entry by text or by a title attribute matches the moment the node exists, which is + * one render before the popup finishes entering. Until then the positioner still carries + * `pointer-events: none` and user-event refuses to click, so that shape is a race a fast machine + * loses. The role query only matches once the popup is open to the accessibility tree, which is + * what makes this wait correct rather than lucky. */ export const chooseSelectOption = async ( user: Pick, "click">, trigger: HTMLElement, optionName: string | RegExp, + role: "option" | "menuitem" | "menuitemradio" = "option", ) => { await user.click(trigger); - const option = await screen.findByRole("option", { name: optionName }); + const option = await screen.findByRole(role, { name: optionName }); await waitFor(() => expect(pointerBlocked(option)).toBe(false)); await user.click(option); }; From 0ec3e936b73a9ebc4445c025a1e4d15a940e29de Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:47:26 -0700 Subject: [PATCH 89/92] fix(cost): inherit the backend's full price structure for off-peak-only deployments Copying only the flat token rates dropped threshold, tiered, service-tier, cache, character, and per-second rates from peak-hour billing once cost lookup switched to the deployment entry, and get_model_info synthesizes zero flat rates for backends without one, which would have marked tiered-only backends explicitly priced free. Copy every price-bearing field instead, deep-copied, rejecting the synthesized zeros the way _inherit_builtin_tiered_output_rate already does. --- litellm/router.py | 28 ++++++--- .../test_router_model_cost_isolation.py | 61 +++++++++++++++++++ 2 files changed, 80 insertions(+), 9 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f3563d9c387..dd1cb51d9af 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8169,16 +8169,23 @@ class Router: backend_model: str, custom_llm_provider: str | None, ) -> None: - """Fill missing base token rates on a deployment entry that only sets + """Fill missing pricing fields on a deployment entry that only sets ``off_peak_pricing``, from the backend model's built-in cost map entry. Cost lookup selects the deployment-scoped entry over the shared backend entry only when the deployment entry carries a base pricing field, and ``off_peak_pricing`` is deliberately kept off the shared entry, so a deployment spelling out only its off-peak schedule would otherwise - never receive the discount. User-specified rates always win; no-op when - any base pricing field is already set or the backend model has no - canonical entry. + never receive the discount. Every price-bearing backend field is + copied, not just the flat token rates: threshold, tiered, service-tier, + cache, character, and per-second rates all carry over, so peak-hour + billing through the deployment entry matches the shared backend entry + exactly. Values are deep-copied to keep the builtin entry isolated, and + a flat token rate ``get_model_info`` synthesized as zero for a backend + without one is rejected, like ``_inherit_builtin_tiered_output_rate`` + does, so a tiered-only backend is never marked explicitly priced free. + User-specified rates always win; no-op when any base pricing field is + already set or the backend model has no canonical entry. """ if not model_info.get("off_peak_pricing"): return @@ -8191,11 +8198,14 @@ class Router: backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider) except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model return - for field in ("input_cost_per_token", "output_cost_per_token"): - if model_info.get(field) is None: - backend_value = backend_info.get(field) - if backend_value is not None: - model_info[field] = backend_value + for field, backend_value in backend_info.items(): + if "cost" not in field and field != "tiered_pricing": + continue + if model_info.get(field) is not None or backend_value is None: + continue + if field in ("input_cost_per_token", "output_cost_per_token") and not backend_value: + continue + model_info[field] = copy.deepcopy(backend_value) @staticmethod def _inherit_builtin_tiered_output_rate( diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index c8adc0af431..a1040e972f2 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -564,6 +564,67 @@ def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates(): assert model_info["off_peak_pricing"] == off_peak_block +def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): + """A backend with above-threshold pricing hands the whole rate structure to + the deployment entry, so peak-hour billing of large prompts through that + entry matches the shared backend entry instead of flattening to the base + rate. + """ + backend_model = "gemini/gemini-2.5-pro" + builtin_info = litellm.get_model_info(model=backend_model) + assert builtin_info["input_cost_per_token_above_200k_tokens"] is not None + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="gemini", + ) + + assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] + assert ( + model_info["input_cost_per_token_above_200k_tokens"] + == builtin_info["input_cost_per_token_above_200k_tokens"] + ) + assert ( + model_info["output_cost_per_token_above_200k_tokens"] + == builtin_info["output_cost_per_token_above_200k_tokens"] + ) + + +def test_inherit_builtin_base_rates_for_off_peak_tiered_only_backend_stores_no_zero(): + """A tiered-only backend has no flat token rates; get_model_info synthesizes + zeros for them, and storing those would mark the deployment explicitly + priced free. The tier table itself must carry over as an isolated copy so + mutating the deployment entry never touches the shared cost map. + """ + backend_model = "dashscope/qwen-flash" + raw_tiers = litellm.model_cost[backend_model]["tiered_pricing"] + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="dashscope", + ) + + assert model_info.get("input_cost_per_token") != 0 + assert model_info.get("output_cost_per_token") != 0 + assert model_info["tiered_pricing"] == raw_tiers + assert model_info["tiered_pricing"] is not raw_tiers + assert model_info["tiered_pricing"][0] is not raw_tiers[0] + + original_first_tier = copy.deepcopy(raw_tiers[0]) + model_info["tiered_pricing"][0]["input_cost_per_token"] = 123.0 + assert raw_tiers[0] == original_first_tier + + def test_inherit_builtin_base_rates_for_off_peak_leaves_explicit_rates_alone(): """An entry that sets its own base rate beside the block already counts as a full custom pricing entry; the helper must not mix builtin rates into it. From 529ac12ba5850fc097877dfd1e9d56d503fdbed0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 1 Sep 2026 13:01:46 -0700 Subject: [PATCH 90/92] test(ui): drop the helper docblock The repo does not take explanatory comments. The reason the helper queries by role lives in the commit that introduced it and in the PR description. --- ui/litellm-dashboard/tests/test-utils.tsx | 9 --------- 1 file changed, 9 deletions(-) diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index 31af51f3f32..553726faff0 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -44,15 +44,6 @@ const pointerBlocked = (element: HTMLElement): boolean => { return false; }; -/** - * Opens a Base UI popup and picks an entry by its accessible name. - * - * Querying the entry by text or by a title attribute matches the moment the node exists, which is - * one render before the popup finishes entering. Until then the positioner still carries - * `pointer-events: none` and user-event refuses to click, so that shape is a race a fast machine - * loses. The role query only matches once the popup is open to the accessibility tree, which is - * what makes this wait correct rather than lucky. - */ export const chooseSelectOption = async ( user: Pick, "click">, trigger: HTMLElement, From 60de5468d748c576b49455391dce897d1e5e800f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:05:25 -0700 Subject: [PATCH 91/92] fix(cost): inherit the backend's raw cost map entry for off-peak-only deployments Filtering copied fields by name dropped companion billing rules like web_search_billing_unit and the regional uplift multipliers, so grounding and uplifts billed differently through the deployment entry. Copy the backend's raw litellm.model_cost entry wholesale instead, which also removes the synthesized-zero special case since the raw entry only holds real values. --- litellm/router.py | 32 ++++++++++--------- .../test_router_model_cost_isolation.py | 24 ++++++++++++++ 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index dd1cb51d9af..c50fab5fe0f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8176,16 +8176,19 @@ class Router: entry only when the deployment entry carries a base pricing field, and ``off_peak_pricing`` is deliberately kept off the shared entry, so a deployment spelling out only its off-peak schedule would otherwise - never receive the discount. Every price-bearing backend field is - copied, not just the flat token rates: threshold, tiered, service-tier, - cache, character, and per-second rates all carry over, so peak-hour - billing through the deployment entry matches the shared backend entry - exactly. Values are deep-copied to keep the builtin entry isolated, and - a flat token rate ``get_model_info`` synthesized as zero for a backend - without one is rejected, like ``_inherit_builtin_tiered_output_rate`` - does, so a tiered-only backend is never marked explicitly priced free. - User-specified rates always win; no-op when any base pricing field is - already set or the backend model has no canonical entry. + never receive the discount. The backend model's entire canonical cost + map entry is copied, field by field, so threshold, tiered, + service-tier, cache, character, and per-second rates as well as + companion billing fields like ``web_search_billing_unit`` and the + regional uplift multipliers all carry over, and peak-hour billing + through the deployment entry matches the shared backend entry exactly. + The raw ``litellm.model_cost`` entry is the copy source rather than + ``get_model_info``'s view of it, since that view synthesizes zero flat + token rates for backends without one and storing those would mark a + tiered-only backend explicitly priced free. Values are deep-copied to + keep the builtin entry isolated. User-specified fields always win; + no-op when any base pricing field is already set or the backend model + has no canonical entry. """ if not model_info.get("off_peak_pricing"): return @@ -8198,13 +8201,12 @@ class Router: backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider) except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model return - for field, backend_value in backend_info.items(): - if "cost" not in field and field != "tiered_pricing": - continue + backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "") + if not isinstance(backend_entry, dict): + return + for field, backend_value in backend_entry.items(): if model_info.get(field) is not None or backend_value is None: continue - if field in ("input_cost_per_token", "output_cost_per_token") and not backend_value: - continue model_info[field] = copy.deepcopy(backend_value) @staticmethod diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index a1040e972f2..7b7a962bf00 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -595,6 +595,30 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): ) +def test_inherit_builtin_base_rates_for_off_peak_carries_companion_billing_fields(): + """Billing rules that are not literal cost rates, like the web search + billing unit, must ride along, or grounding and regional uplifts would + bill differently through the deployment entry than through the shared + backend entry. + """ + backend_model = "gemini-3-pro-image" + raw_entry = litellm.model_cost[backend_model] + assert raw_entry.get("web_search_billing_unit") is not None + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider=None, + ) + + assert model_info["web_search_billing_unit"] == raw_entry["web_search_billing_unit"] + assert model_info["input_cost_per_token"] == raw_entry["input_cost_per_token"] + + def test_inherit_builtin_base_rates_for_off_peak_tiered_only_backend_stores_no_zero(): """A tiered-only backend has no flat token rates; get_model_info synthesizes zeros for them, and storing those would mark the deployment explicitly From cdb1245e7419fa8b0294da3946cbad68b3d567d5 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 1 Sep 2026 13:30:02 -0700 Subject: [PATCH 92/92] fix(s3): bound s3 object keys and download filenames for long Responses API ids (#39164) * fix(s3): bound object keys and download filenames to s3 limits Long OpenAI-compatible Responses API ids pushed the s3 object key past s3's 1024 UTF-8 byte cap, so the PUT failed with a 400 and the log record was dropped. Keys that still fit are unchanged, byte for byte. An oversized one now keeps a readable head of the file name and appends the sha256 of the full name. A configured path/alias prefix that is long enough to overflow on its own keeps whole leading path segments, so a prefix-scoped IAM policy or lifecycle rule still matches, and ends in a short digest of the full configured value so two operators do not land in the same folder. The Content-Disposition filename carried the same unbounded id and hit s3's 2048 byte metadata-header cap, so the upload still failed with MetadataTooLarge once the key was bounded. It is bounded the same way, head plus digest, so two records downloaded from the console stay distinct files. The full response id stays in the uploaded JSON payload. * fix(s3): keep the configured prefix whole and spend the whole key budget Shorten the response id first and only trim the operator's configured prefix when the prefix itself is what does not fit, so prefix scoped IAM policies and lifecycle rules keep matching. Trim by bytes rather than whole segments so the longest possible string prefix survives, and route the audit log key through the same shared builder. * chore(s3): trim the comments and docstrings the review flagged Keep the two external facts that are not visible from the code, the 1024 byte object key cap and the 2048 byte metadata header cap, and drop the rest. --- litellm/constants.py | 6 + litellm/integrations/s3.py | 81 ++++- litellm/integrations/s3_v2.py | 20 +- tests/test_litellm/integrations/test_s3.py | 42 ++- tests/test_litellm/integrations/test_s3_v2.py | 288 ++++++++++++++++++ 5 files changed, 409 insertions(+), 28 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index b9bea466247..9a50797f517 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -13,6 +13,12 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) +# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html +MAX_S3_OBJECT_KEY_BYTES: Final = 1024 +S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 +S3_PREFIX_DIGEST_CHARS: Final = 16 +# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against +MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index ddeb410c54a..8ce461eea5b 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -1,11 +1,18 @@ #### What this does #### # On success + failure, log events to Supabase +import hashlib from datetime import datetime from typing import Final, cast import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.constants import ( + MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, + MAX_S3_OBJECT_KEY_BYTES, + S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, + S3_PREFIX_DIGEST_CHARS, +) from litellm.types.utils import StandardLoggingPayload @@ -133,9 +140,7 @@ class S3Logger: s3_file_name, ) - s3_object_download_filename: Final = ( - "time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json" - ) + s3_object_download_filename: Final = get_s3_object_download_filename(start_time, payload["id"]) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -198,6 +203,47 @@ def resolve_sse_params( return algorithm, valid_key_id +S3_MIN_BOUNDED_FILE_NAME_BYTES: Final = 64 + + +def _truncate_to_utf8_bytes(value: str, max_bytes: int) -> str: + """Trim `value` so its UTF-8 encoding fits `max_bytes`, never splitting a character.""" + if max_bytes <= 0: + return "" + encoded: Final = value.encode("utf-8") + if len(encoded) <= max_bytes: + return value + return encoded[:max_bytes].decode("utf-8", errors="ignore") + + +def get_s3_object_download_filename(start_time: datetime, response_id: str) -> str: + """Content-Disposition filename for the uploaded object, bounded to the metadata header cap.""" + sanitized_response_id: Final = response_id.replace("/", "_").replace('"', "_") + file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{response_id}" + sanitized_file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{sanitized_response_id}" + budget: Final = MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES - len(b".json") + if len(sanitized_file_name.encode("utf-8")) <= budget: + return sanitized_file_name + ".json" + return _bounded_s3_file_name(file_name, sanitized_file_name, budget) + ".json" + + +def _bounded_s3_file_name(s3_file_name: str, sanitized_s3_file_name: str, max_bytes: int) -> str: + """As much of the file name as `max_bytes` allows, then the sha256 of the whole name.""" + digest: Final = hashlib.sha256(s3_file_name.encode("utf-8")).hexdigest() + head_budget: Final = min(S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, max_bytes - len(digest) - 1) + head: Final = _truncate_to_utf8_bytes(sanitized_s3_file_name, head_budget) + return f"{head}_{digest}" if head else digest + + +def _bounded_s3_prefix(configured_prefix: str, max_bytes: int) -> str: + """As much of the configured prefix as fits, then a digest segment naming the full prefix.""" + digest_segment: Final = hashlib.sha256(configured_prefix.encode("utf-8")).hexdigest()[:S3_PREFIX_DIGEST_CHARS] + "/" + if max_bytes < len(digest_segment): + return "" + head: Final = _truncate_to_utf8_bytes(configured_prefix, max_bytes - len(digest_segment) - 1).rstrip("/") + return f"{head}/{digest_segment}" if head else digest_segment + + def get_s3_object_key( s3_path: str, prefix: str, @@ -205,12 +251,23 @@ def get_s3_object_key( s3_file_name: str, ) -> str: sanitized_s3_file_name: Final = s3_file_name.replace("/", "_") - s3_object_key = ( - (s3_path.rstrip("/") + "/" if s3_path else "") - + prefix - + start_time.strftime("%Y-%m-%d") - + "/" - + sanitized_s3_file_name - ) # we need the s3 key to include the time, so we log cache hits too - s3_object_key += ".json" - return s3_object_key + configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix + date_segment: Final = start_time.strftime("%Y-%m-%d") + "/" + # we need the s3 key to include the time, so we log cache hits too + s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json" + if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES: + return s3_object_key + + # shorten the response id first and only trim the configured prefix if that is what does not + # fit, so prefix scoped IAM policies and lifecycle rules keep matching + budget: Final = MAX_S3_OBJECT_KEY_BYTES - len(date_segment.encode("utf-8")) - len(b".json") + prefix_bytes: Final = len(configured_prefix.encode("utf-8")) + if prefix_bytes + S3_MIN_BOUNDED_FILE_NAME_BYTES <= budget: + bounded_file_name: Final = _bounded_s3_file_name(s3_file_name, sanitized_s3_file_name, budget - prefix_bytes) + return configured_prefix + date_segment + bounded_file_name + ".json" + + shortest_file_name: Final = _bounded_s3_file_name( + s3_file_name, sanitized_s3_file_name, S3_MIN_BOUNDED_FILE_NAME_BYTES + ) + bounded_prefix: Final = _bounded_s3_prefix(configured_prefix, budget - len(shortest_file_name.encode("utf-8"))) + return bounded_prefix + date_segment + shortest_file_name + ".json" diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 9f6ae72fb3a..712ce41d09e 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -16,7 +16,11 @@ from urllib.parse import quote import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS -from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params +from litellm.integrations.s3 import ( + get_s3_object_download_filename, + get_s3_object_key, + resolve_sse_params, +) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -259,11 +263,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): now: Final = datetime.now(timezone.utc) audit_log_id: Final = audit_log.get("id", "unknown") - s3_path = cast(str | None, self.s3_path) or "" - s3_path = s3_path.rstrip("/") + "/" if s3_path else "" - - s3_object_key: Final = ( - f"{s3_path}audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log_id}.json" + s3_object_key: Final = get_s3_object_key( + cast(str | None, self.s3_path) or "", + "audit_logs/", + now, + f"{now.strftime('%H-%M-%S')}_{audit_log_id}", ) element: Final = s3BatchLoggingElement( @@ -463,9 +467,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) verbose_logger.debug("s3_object_key=%s", s3_object_key) - s3_object_download_filename: Final = ( - f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" - ) + s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"]) return s3BatchLoggingElement( payload=dict(standard_logging_payload), diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py index 7e997870852..58b15b79e76 100644 --- a/tests/test_litellm/integrations/test_s3.py +++ b/tests/test_litellm/integrations/test_s3.py @@ -2,26 +2,27 @@ from datetime import datetime from unittest.mock import MagicMock, patch import litellm +from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import S3Logger TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id" -def _standard_logging_payload() -> dict: +def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict: return { - "id": "chatcmpl-test-id", + "id": response_id, "metadata": {"user_api_key_team_alias": None}, } -def _log_event_kwargs() -> dict: +def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict: return { "litellm_params": {"metadata": {}}, - "standard_logging_object": _standard_logging_payload(), + "standard_logging_object": _standard_logging_payload(response_id), } -def _run_log_event(callback_params: dict) -> MagicMock: +def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock: original = litellm.s3_callback_params litellm.s3_callback_params = callback_params try: @@ -30,8 +31,8 @@ def _run_log_event(callback_params: dict) -> MagicMock: mock_boto3_client.return_value = mock_s3_client logger = S3Logger() logger.log_event( - kwargs=_log_event_kwargs(), - response_obj={}, + kwargs=_log_event_kwargs(response_id), + response_obj={"id": response_id}, start_time=datetime(2026, 7, 30, 12, 0, 0), end_time=datetime(2026, 7, 30, 12, 0, 1), print_verbose=lambda *args, **kwargs: None, @@ -154,3 +155,30 @@ def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): put_object_kwargs = mock_s3_client.put_object.call_args.kwargs assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" assert "SSEKMSKeyId" not in put_object_kwargs + + +def test_put_object_key_and_filename_are_bounded_for_an_oversized_response_id(): + """The sync logger bounds both the key and the Content-Disposition filename.""" + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": "logs"}, + response_id="resp_" + "A" * 1100, + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert len(put_object_kwargs["Key"].encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert put_object_kwargs["Key"].startswith("logs/2026-07-30/time-12-00-00-000000_resp_") + filename = put_object_kwargs["ContentDisposition"].removeprefix('inline; filename="').removesuffix('"') + assert len(filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + +def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shrink(): + """A long configured s3_path survives whole when the id can be shortened instead.""" + long_path = "litellm-prod-logs/" + "t" * 921 + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": long_path}, + response_id="resp_" + "B" * 100, + ) + + key = mock_s3_client.put_object.call_args.kwargs["Key"] + assert key.startswith(long_path + "/2026-07-30/") + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 51671d5101e..a037284d7c1 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1170,6 +1170,294 @@ def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json") +# -------------------------------------------------------------- +# object keys bounded to S3's 1024 UTF-8 byte limit +# -------------------------------------------------------------- +def _oversized_response_id() -> str: + return "resp_" + "A" * 1100 + + +def test_s3_object_key_at_the_byte_limit_is_left_alone(): + """A key that still fits is left byte-identical.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + fixed_len = len("input/2026-08-24/.json") + file_name = "x" * (MAX_S3_OBJECT_KEY_BYTES - fixed_len) + + key = get_s3_object_key(s3_path="input", prefix="", start_time=start_time, s3_file_name=file_name) + + assert key == f"input/2026-08-24/{file_name}.json" + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def test_s3_object_key_is_bounded_for_oversized_response_id(): + """An oversized Responses API id is shortened to a readable head plus a digest.""" + import hashlib + + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + file_name = f"time-06-18-41-948021_{_oversized_response_id()}" + + key = get_s3_object_key(s3_path="input", prefix="DefaultTeamProd/", start_time=start_time, s3_file_name=file_name) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("input/DefaultTeamProd/2026-08-24/time-06-18-41-948021_resp_") + assert key.endswith(f"_{hashlib.sha256(file_name.encode('utf-8')).hexdigest()}.json") + + +@pytest.mark.parametrize( + "s3_path,prefix", + [ + ("input", ""), + ("a" * 900, ""), + ("input", "team-" + "b" * 900 + "/"), + ("c" * 600, "team-" + "d" * 600 + "/key-" + "e" * 600 + "/"), + # many short segments, so the trim lands exactly on the budget edge + ("", "ssss/" * 200), + ], +) +def test_s3_object_key_is_bounded_for_long_paths_and_aliases(s3_path: str, prefix: str): + """Long paths, team aliases and key aliases stay within the cap.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + key = get_s3_object_key( + s3_path=s3_path, + prefix=prefix, + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.endswith(".json") + assert "/2026-08-24/" in key or key.startswith("2026-08-24/") + assert "/" not in key.rsplit("2026-08-24/", 1)[1] + + +def test_s3_object_key_trimmed_prefixes_stay_distinct_per_operator(): + """Prefixes that differ only past the trim point keep separate folders.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + keys = [ + get_s3_object_key( + s3_path="input", + prefix="team-" + "b" * 1000 + suffix + "/", + start_time=start_time, + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + for suffix in ("-one", "-two") + ] + + assert keys[0] != keys[1] + assert all(key.startswith("input/team-" + "b" * 900) for key in keys) + assert all(len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES for key in keys) + + +def test_s3_object_key_bounded_prefix_never_splits_a_multibyte_character(): + """A multibyte prefix is trimmed on a character boundary.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + s3_path = "\u65e5\u672c\u8a9e" * 200 + + key = get_s3_object_key( + s3_path=s3_path, + prefix="\u30c1\u30fc\u30e0" * 200 + "/", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.startswith(s3_path[:100]) + assert "\ufffd" not in key + + +def test_s3_object_key_stays_unique_for_ids_sharing_a_head(): + """Ids sharing a visible head still get distinct keys.""" + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + keys = { + get_s3_object_key( + s3_path="input", + prefix="", + start_time=start_time, + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}{suffix}", + ) + for suffix in ("first", "second", "third") + } + + assert len(keys) == 3 + + +def test_s3_object_key_bounding_matches_the_documented_layout(): + """The bounded key is `//_.json`.""" + import hashlib + + from litellm.integrations.s3 import get_s3_object_key + + file_name = f"time-06-18-41-948021_{_oversized_response_id()}" + + key = get_s3_object_key( + s3_path="input", + prefix="team/", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=file_name, + ) + + digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest() + assert key == f"input/team/2026-08-24/{file_name[:64]}_{digest}.json" + + +def test_s3_object_key_keeps_the_configured_prefix_when_only_the_id_overflows(): + """A 940 byte configured prefix survives whole when only the id overflows.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + prefix = "team-" + "b" * 934 + "/" + + key = get_s3_object_key( + s3_path="", + prefix=prefix, + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert key.startswith(prefix + "2026-08-24/") + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def test_s3_object_key_spends_the_whole_budget_when_the_prefix_must_be_trimmed(): + """A trimmed prefix keeps every byte the budget allows, not whole segments.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + s3_path = "p" * 400 + "/" + "q" * 600 + + key = get_s3_object_key( + s3_path=s3_path, + prefix="", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name="time-06-18-41-948021_abc", + ) + + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("p" * 400 + "/" + "q" * 500) + + +def test_s3_object_key_keeps_a_single_segment_path_as_far_as_it_fits(): + """A path with no separator is kept as far as it fits, never dropped to the bucket root.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + key = get_s3_object_key( + s3_path="a" * 1050, + prefix="", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name="time-06-18-41-948021_chatcmpl-xyz", + ) + + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("a" * 900) + + +def test_create_s3_batch_logging_element_bounds_key_and_keeps_full_response_id(): + """The batch element bounds the key and keeps the full response id in the payload.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + + logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) + response_id = _oversized_response_id() + payload = StandardLoggingPayload( + id=response_id, + metadata={"user_api_key_team_alias": "DefaultTeamProd", "user_api_key_alias": "prod-key"}, + messages=[], + ) + + result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) + + assert result is not None + assert len(result.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert result.s3_object_key.startswith("DefaultTeamProd/prod-key/2026-08-24/") + assert result.payload["id"] == response_id + + +def test_s3_object_download_filename_is_bounded_for_oversized_response_id(): + """The Content-Disposition filename is bounded too, or the PUT fails with MetadataTooLarge.""" + from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), _oversized_response_id()) + + assert len(file_name.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + assert file_name.startswith("time-2026-08-24T06-18-41-948021_resp_") + assert file_name.endswith(".json") + + +def test_s3_object_download_filenames_stay_distinct_when_shortened(): + """Shortened filenames stay distinct.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + file_names = { + get_s3_object_download_filename(start_time, _oversized_response_id() + suffix) + for suffix in ("first", "second", "third") + } + + assert len(file_names) == 3 + + +def test_s3_object_download_filename_short_id_is_unchanged(): + """An ordinary response id keeps the filename it had before.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), "resp_abc123") + + assert file_name == "time-2026-08-24T06-18-41-948021_resp_abc123.json" + + +def test_create_s3_batch_logging_element_bounds_the_download_filename(): + """The batch element carries a bounded Content-Disposition filename.""" + from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + logger = S3Logger() + payload = StandardLoggingPayload(id=_oversized_response_id(), metadata={}, messages=[]) + + result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) + + assert result is not None + assert len(result.s3_object_download_filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + +@pytest.mark.asyncio +async def test_audit_log_object_key_is_bounded_for_a_long_configured_path(): + """Audit log keys are bounded by the same builder.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + + logger = S3Logger() + logger.s3_path = "audit-archive/" + "z" * 1100 + + await logger.async_log_audit_log_event({"id": "1a4f7bd0-6f1e-4d0a-9b3c-9f2e1d5a7c88"}) + + assert len(logger.log_queue) == 1 + assert len(logger.log_queue[0].s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert logger.log_queue[0].s3_object_key.startswith("audit-archive/" + "z" * 900) + + +def test_s3_object_download_filename_drops_characters_that_break_the_header(): + """A quote or separator in the response id cannot escape the quoted header value.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), 'resp_a"b/c') + + assert file_name == "time-2026-08-24T06-18-41-948021_resp_a_b_c.json" + + # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # --------------------------------------------------------------
MetricValue
MetricValue
{row.metric}{row.value}
{row.metric}{row.value}