scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_stream_modify_response_chunks

Browse source
This commit is contained in:
mateo-berri 2026-09-01 14:49:06 -07:00
commit a38dfecd96
137 changed files with 6960 additions and 1364 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
.circleci/config.yml
View file

@ -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 legacy_resolver"
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_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 legacy_resolver"
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
installing_litellm_on_python_legacy_migration_resolver:
installing_litellm_on_python_v2_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 legacy migration resolver proxy smoke test
name: Run v2 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_legacy_resolver
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
helm_chart_testing:
machine:
@ -2879,8 +2879,7 @@ 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 || \
grep -q "Database migration cannot proceed" docker_output.log); then
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
echo "Expected error found. Test passed."
else
echo "Expected error not found. Test failed."
@ -3012,7 +3011,7 @@ workflows:
filters: *main_branches
- installing_litellm_on_python_3_13:
filters: *main_branches
- installing_litellm_on_python_legacy_migration_resolver:
- installing_litellm_on_python_v2_migration_resolver:
filters: *main_branches
- helm_chart_testing:
requires:

2
CLAUDE.md
View file

@ -37,7 +37,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 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 --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
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

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

@ -7,8 +7,7 @@ import subprocess
import tempfile
import time
from pathlib import Path
from types import MappingProxyType
from typing import Final, Optional
from typing import Optional
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
@ -51,17 +50,6 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
re.IGNORECASE,
)
_PRISMA_ATTEMPTS: Final = 4
_TRANSIENT_PRISMA_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",
}
)
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` "
@ -286,23 +274,6 @@ 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:
"""
@ -684,7 +655,7 @@ class ProxyExtrasDBManager:
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
v2 migration resolver (what the proxy CLI selects by default).
v2 migration resolver (opt-in via --use_v2_migration_resolver).
Runs `prisma migrate deploy` and handles standard recovery paths
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
@ -705,46 +676,20 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
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.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
)
# 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."
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),
check=True,
env=_get_prisma_env(),
)
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)
@ -754,7 +699,7 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
for attempt in range(_PRISMA_ATTEMPTS):
for attempt in range(4):
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
@ -869,36 +814,16 @@ class ProxyExtrasDBManager:
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
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 == _PRISMA_ATTEMPTS - 1:
raise RuntimeError(
f"Database migration failed after "
f"{_PRISMA_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 and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
raise RuntimeError(
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."
"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."
)
finally:
os.chdir(original_dir)
@ -946,11 +871,10 @@ class ProxyExtrasDBManager:
Args:
use_migrate: Whether to use prisma migrate instead of db push
use_v2_resolver: Run the v2 migration resolver (safer during
use_v2_resolver: Opt into the v2 migration resolver (safer during
rolling deploys; does not run the diff-and-force recovery
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.
that causes schema thrashing). Defaults to False for
backwards compatibility.
Returns:
bool: True if setup was successful, False otherwise
@ -968,7 +892,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")
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"

0
litellm-proxy-extras/tests/__init__.py Normal file
View file

242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py Normal file
View file

@ -0,0 +1,242 @@
"""Regression tests for ProxyExtrasDBManager 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.
"""
import subprocess
from unittest.mock import patch
import pytest
from litellm_proxy_extras.utils import (
ProxyExtrasDBManager,
_max_migration_timestamp,
_migration_timestamp,
)
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
def _run(*args, **kwargs):
raise subprocess.CalledProcessError(
returncode=returncode,
cmd=args[0],
stderr=stderr,
output="",
)
return _run
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a permission failure during migrate deploy raises RuntimeError."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3018\nMigration name: 20250326162113_baseline\n"
"Database error code: 42501\npermission denied for schema public"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="permission"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
'Reason: syntax error at or near "BRKN" LINE 42'
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_strip_prisma_query_params_removes_connection_limit():
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
assert "connection_limit" not in stripped
assert "pool_timeout" not in stripped
assert "sslmode=require" in stripped
def test_strip_prisma_query_params_passthrough_no_query():
"""URLs without query strings are returned unchanged."""
url = "postgresql://u:p@h:5432/db"
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
def test_migration_timestamp_extracts_leading_digits():
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
def test_migration_timestamp_returns_zero_on_malformed():
assert _migration_timestamp("0_init") == 0
assert _migration_timestamp("not_a_migration") == 0
def test_max_migration_timestamp():
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
assert _max_migration_timestamp(names) == 20260415000000
def test_max_migration_timestamp_empty_set():
assert _max_migration_timestamp(set()) == 0
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
This is the existing buggy behavior — we're not fixing it in v1, only
offering v2 as opt-in. This test pins the default so that a future
inadvertent default flip is caught.
"""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
# Stub `prisma migrate deploy` to claim success with pending migrations
# applied, which is the code path that triggers the legacy post-migration
# sanity check (a call to _resolve_all_migrations).
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
def fake_run(cmd, *args, **kwargs):
return FakeResult()
resolve_called = {"n": 0}
def fake_resolve(*args, **kwargs):
resolve_called["n"] += 1
monkeypatch.setattr("subprocess.run", fake_run)
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
assert ok is True
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = "db push error"
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="prisma db push failed"):
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
"""_warn_if_db_ahead_of_head must never raise — it's informational.
Non-connection DB errors (e.g. InsufficientPrivilege from a user
without SELECT on _prisma_migrations) must be caught, not propagated.
"""
import psycopg
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class _FakeConn:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def execute(self, *a, **kw):
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
raise psycopg.errors.InsufficientPrivilege("permission denied")
def _fake_connect(*a, **kw):
return _FakeConn()
monkeypatch.setattr("psycopg.connect", _fake_connect)
# Must not raise.
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
monkeypatch, tmp_path
):
"""If marking a migration as applied fails inside P3009 idempotent
recovery, the subprocess error must be re-raised as RuntimeError so
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
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(
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
)
# First call: migrate deploy -> P3009 idempotent error.
# Recovery path tries _resolve_specific_migration; that also raises.
def _failing_resolve(*a, **kw):
raise subprocess.CalledProcessError(
returncode=1,
cmd="prisma migrate resolve --applied",
stderr="resolve failed",
output="",
)
monkeypatch.setattr(
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
)
stderr = (
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
"relation already exists"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(
RuntimeError, match="Failed to mark migration .* as applied"
):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
resolve_called = {"n": 0}
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_all_migrations",
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
)
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"

7
litellm/constants.py
View file

@ -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))
@ -130,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.

20
litellm/experimental_mcp_client/client.py
View file

@ -7,6 +7,7 @@ import base64
import os
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from functools import partial
from importlib import metadata
from typing import Any, Final, TypeVar
@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.types.llms.custom_http import VerifyTypes
from litellm.types.mcp import (
@ -603,17 +605,19 @@ class MCPClient:
"""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
async def _list_tools_operation(session: ClientSession):
return await session.list_tools()
try:
result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
tool_count: Final = len(result.tools)
tool_names: Final = [tool.name for tool in result.tools]
# A per-server timeout above the global default extends the whole-walk deadline
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
tools: Final = await self.run_with_session(
partial(list_tools_with_pagination, listing_deadline=listing_deadline),
quiet_on_error=raise_on_error,
)
tool_count: Final = len(tools)
tool_names: Final = tuple(tool.name for tool in tools)
verbose_logger.info(
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
)
return result.tools
return tools
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
raise

74
litellm/experimental_mcp_client/tools.py
View file

@ -1,14 +1,22 @@
import json
from typing import Final, Literal
import anyio
from mcp import ClientSession
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import PaginatedRequestParams
from mcp.types import Tool as MCPTool
from openai.types.chat import ChatCompletionToolParam
from openai.types.responses.function_tool_param import FunctionToolParam
from openai.types.shared_params.function_definition import FunctionDefinition
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.types.llms.anthropic import AnthropicMessagesTool
from litellm.types.utils import ChatCompletionMessageToolCall
@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
)
async def list_tools_with_pagination(
session: ClientSession, listing_deadline: float | None = None
) -> list[MCPTool]: # mutable-ok: list return contract
"""Collect tools from every tools/list page by following nextCursor.
Stops and returns the tools collected so far when the upstream repeats a
cursor, the page cap is reached, or the whole-walk deadline expires, so a
buggy or slow upstream yields a partial catalog instead of an error.
listing_deadline overrides the default whole-walk deadline; callers with a
per-server timeout above the global default pass it through here.
"""
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
# The per-request session read timeout restarts on every page, so a multi-page
# walk needs its own overall deadline. max() keeps the pre-pagination guarantee
# that a single page slower than the listing timeout but within the client
# timeout still succeeds.
effective_deadline: Final = (
listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT)
)
with anyio.move_on_after(effective_deadline):
for _ in range(MCP_TOOL_LISTING_MAX_PAGES):
result = (
await session.list_tools()
if cursor is None
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
)
tools.extend(result.tools)
next_cursor = getattr(result, "nextCursor", None)
if not isinstance(next_cursor, str) or not next_cursor:
return tools
if next_cursor in seen_cursors:
verbose_logger.warning(
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
len(tools),
)
return tools
seen_cursors.add(next_cursor)
cursor = next_cursor
verbose_logger.warning(
"MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far",
MCP_TOOL_LISTING_MAX_PAGES,
len(tools),
)
return tools
verbose_logger.warning(
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
effective_deadline,
len(tools),
)
return tools
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> list[MCPTool] | list[ChatCompletionToolParam]:
@ -103,10 +169,12 @@ async def load_mcp_tools(
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
"""
tools: Final = await session.list_tools()
tools: Final = await list_tools_with_pagination(session)
if format == "openai":
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
return tools.tools
return [ # mutable-ok: public API returns a list
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
]
return tools
########################################################

52
litellm/integrations/prometheus.py
View file

@ -59,6 +59,8 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from prometheus_client.metrics import MetricWrapperBase
from litellm.router import Router
else:
AsyncIOScheduler = Any
@ -67,6 +69,8 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
_DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other"
_NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
(
"guardrail_name",
@ -154,6 +158,44 @@ def _get_budget_metrics_per_request_timeout() -> float:
return parsed
def _get_proxy_llm_router() -> Router | None:
try:
from litellm.proxy.proxy_server import llm_router
except Exception:
return None
return llm_router
def _bounded_requested_model_label(requested_model: str | None, router_originated: bool = False) -> str | None:
"""
Bound ``requested_model`` label cardinality: names the router recognizes
(model names, deployment ids, aliases, routing groups, team public model
names) or matches via a global or team wildcard/pattern route keep their
own label value; any other client-supplied string collapses into the
single ``other`` bucket. With no proxy router to vouch for the string,
client-supplied values collapse to ``other`` while ``router_originated``
values (emitted by an SDK ``Router``'s own deployment failure and
fallback events, where the proxy router never exists) pass through.
"""
if not requested_model:
return requested_model
llm_router: Final = _get_proxy_llm_router()
if llm_router is None:
return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL
if llm_router.is_recognized_model(requested_model):
return requested_model
if requested_model in llm_router.team_public_model_names:
return requested_model
if llm_router.pattern_router.route(requested_model) is not None:
return requested_model
if any(
team_pattern_router.route(requested_model) is not None
for team_pattern_router in llm_router.team_pattern_routers.values()
):
return requested_model
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -2407,7 +2449,7 @@ class PrometheusLogger(CustomLogger):
team_alias=user_api_key_dict.team_alias,
org_id=user_api_key_dict.org_id,
org_alias=user_api_key_dict.organization_alias,
requested_model=request_data.get("model", ""),
requested_model=_bounded_requested_model_label(request_data.get("model", "")),
status_code=str(status_code),
exception_status=str(status_code),
exception_class=self._get_exception_class_name(original_exception),
@ -2627,7 +2669,9 @@ class PrometheusLogger(CustomLogger):
label_model_id = ""
label_api_base = ""
label_api_provider = ""
label_requested_model = litellm_model_name or model_group or ""
label_requested_model = (
_bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or ""
)
enum_values: Final = UserAPIKeyLabelValues(
litellm_model_name=label_litellm_model_name,
@ -3186,7 +3230,7 @@ class PrometheusLogger(CustomLogger):
_tags: Final = cast(list[str], kwargs.get("tags") or [])
enum_values: Final = UserAPIKeyLabelValues(
requested_model=original_model_group,
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
fallback_model=_new_model,
hashed_api_key=standard_metadata["user_api_key_hash"],
api_key_alias=standard_metadata["user_api_key_alias"],
@ -3227,7 +3271,7 @@ class PrometheusLogger(CustomLogger):
)
enum_values: Final = UserAPIKeyLabelValues(
requested_model=original_model_group,
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
fallback_model=_new_model,
hashed_api_key=standard_metadata["user_api_key_hash"],
api_key_alias=standard_metadata["user_api_key_alias"],

81
litellm/integrations/s3.py
View file

@ -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"

20
litellm/integrations/s3_v2.py
View file

@ -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),

4
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider(
return model, custom_llm_provider
def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
"""The authenticating provider this pair already names, or None.
get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their
@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
adopt the declaration instead of resolving.
"""
declared: Final = custom_llm_provider or model.split("/", 1)[0]
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None

8
litellm/litellm_core_utils/litellm_logging.py
View file

@ -111,6 +111,7 @@ from litellm.types.mcp import MCPPostCallResponseObject
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
DEPLOYMENT_SCOPED_PRICING_FIELDS,
CachingDetails,
CallTypes,
CostBreakdown,
@ -255,6 +256,7 @@ _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggi
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
sentry_sdk_instance = None
capture_exception = None
@ -5033,7 +5035,9 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
"""
Check if the model uses custom pricing
Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info`
Returns True if any custom pricing field is present in `litellm_params`, or if
any custom pricing or deployment-scoped pricing field (such as
``off_peak_pricing``) is present in the metadata ``model_info``
"""
if litellm_params is None:
return False
@ -5051,7 +5055,7 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
model_info: dict = metadata.get("model_info", {}) or {}
if model_info:
matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys()
matching_keys = _MODEL_INFO_CUSTOM_PRICING_KEYS & model_info.keys()
for key in matching_keys:
if model_info.get(key) is not None:
return True

215
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -2,10 +2,12 @@
## 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, 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
@ -290,10 +292,187 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float,
)
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, 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()
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").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:
if start <= now < end:
return True
elif now >= start or now < end:
return True
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
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 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 isinstance(off_peak, Mapping) 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),
_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 _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,
service_tier: str | None = None,
current_time: datetime | None = None,
*,
threshold_is_inclusive: bool = False,
) -> tuple[float, float, float, float, float]:
@ -311,7 +490,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)
@ -345,12 +524,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:
return (
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
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+)
@ -451,12 +634,16 @@ def _get_token_base_cost(
except Exception:
continue
return (
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
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,
),
)

13
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -1281,6 +1281,19 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin
return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo
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
def _get_image_mime_type_from_url(url: str) -> str | None:
"""
Get mime type for common image URLs

52
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -1694,6 +1694,18 @@ def convert_function_to_anthropic_tool_invoke(
raise e
def _find_server_tool_result(
tool_id: str,
web_search_results: Sequence[object] | None,
tool_results: Sequence[object] | None,
) -> dict[str, object] | None:
candidates: Final = (*(web_search_results or ()), *(tool_results or ()))
return next(
(result for result in candidates if isinstance(result, dict) and result.get("tool_use_id") == tool_id),
None,
)
def convert_to_anthropic_tool_invoke(
tool_calls: list[ChatCompletionAssistantToolCall],
web_search_results: list[Any] | None = None,
@ -1758,32 +1770,22 @@ def convert_to_anthropic_tool_invoke(
context="Anthropic tool invoke",
)
# Check if this is a server-side tool (web_search, tool_search, etc.)
# Server tool IDs start with "srvtoolu_"
if tool_id.startswith("srvtoolu_"):
# Create server_tool_use block instead of tool_use
_anthropic_server_tool_use: dict[str, object] = {
"type": "server_tool_use",
"id": tool_id,
"name": tool_name,
"input": tool_input,
}
anthropic_tool_invoke.append(_anthropic_server_tool_use)
# Add corresponding tool result if available.
# Check both web_search_results (web_search_tool_result / web_fetch_tool_result)
# and tool_results (bash_code_execution_tool_result, etc.)
_all_tool_results: list[Any] = []
if web_search_results:
_all_tool_results.extend(web_search_results)
if tool_results:
_all_tool_results.extend(tool_results)
for result in _all_tool_results:
if result.get("tool_use_id") == tool_id:
anthropic_tool_invoke.append(result)
break
server_tool_result = (
_find_server_tool_result(tool_id, web_search_results, tool_results)
if tool_id.startswith("srvtoolu_")
else None
)
if server_tool_result is not None:
anthropic_tool_invoke.append(
{
"type": "server_tool_use",
"id": tool_id,
"name": tool_name,
"input": tool_input,
}
)
anthropic_tool_invoke.append(server_tool_result)
else:
# Regular tool_use
sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id)
_anthropic_tool_use_param = AnthropicMessagesToolUseParam(
type="tool_use",

2
litellm/llms/anthropic/chat/transformation.py
View file

@ -1516,7 +1516,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled)
if _tool is None:
continue
if not is_thinking_enabled:
if not is_thinking_enabled and not AnthropicModelInfo.forced_tool_use_unsupported(model):
_tool_choice = {
"name": RESPONSE_FORMAT_TOOL_NAME,
"type": "tool",

6
litellm/llms/anthropic/common_utils.py
View file

@ -325,13 +325,17 @@ class AnthropicModelInfo(BaseLLMModelInfo):
status_code=400,
)
@staticmethod
def forced_tool_use_unsupported(model: str) -> bool:
return AnthropicModelInfo._get_model_capability(model, "supports_forced_tool_use") is False
@staticmethod
def forced_tool_use_downgraded(model: str, drop_params: bool) -> bool:
"""True when the model map flags the model with
``supports_forced_tool_use: false`` (Fable 5.1 / Mythos 5.1 400 on
``any``/``tool``) and ``drop_params`` asks for the ``auto`` downgrade;
raises a clean client-side 400 for such models without ``drop_params``."""
if AnthropicModelInfo._get_model_capability(model, "supports_forced_tool_use") is not False:
if not AnthropicModelInfo.forced_tool_use_unsupported(model):
return False
if not (litellm.drop_params or drop_params):
raise litellm.utils.UnsupportedParamsError(

17
litellm/llms/azure/chat/gpt_transformation.py
View file

@ -1,3 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from httpx._models import Headers, Response
@ -6,6 +8,7 @@ import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
drop_tool_reference_parts_from_tool_messages,
hoist_images_from_tool_messages,
tool_with_flattened_parameters,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_azure_openai_messages,
@ -32,6 +35,19 @@ else:
LoggingClass = Any
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]:
tools: Final = optional_params.get("tools")
if not isinstance(tools, list):
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})
class AzureOpenAIConfig(BaseConfig):
"""
Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions
@ -261,6 +277,7 @@ class AzureOpenAIConfig(BaseConfig):
"model": model,
"messages": azure_messages,
**optional_params,
**flattened_tools_update(optional_params),
}
def transform_response(

7
litellm/llms/azure/chat/o_series_transformation.py
View file

@ -20,6 +20,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.utils import get_model_info, supports_reasoning
from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig
from .gpt_transformation import flattened_tools_update
class AzureOpenAIO1Config(OpenAIOSeriesConfig):
@ -108,4 +109,8 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig):
headers: dict,
) -> dict:
model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name
return super().transform_request(model, messages, optional_params, litellm_params, headers)
flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict
**optional_params,
**flattened_tools_update(optional_params),
}
return super().transform_request(model, messages, flattened_params, litellm_params, headers)

1
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -1073,6 +1073,7 @@ class AmazonConverseConfig(BaseConfig):
if (
litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider)
and not is_thinking_enabled
and not AnthropicModelInfo.forced_tool_use_unsupported(model)
):
optional_params["tool_choice"] = ToolChoiceValuesBlock(
tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME)

22
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_image_obj,
)
@ -11,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
convert_url_to_base64,
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
@ -74,10 +76,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
drop_params: bool,
) -> dict:
# Force tool-based structured outputs for Bedrock Invoke
# (similar to VertexAI fix in #19201)
# Bedrock Invoke doesn't support output_format parameter
# (similar to VertexAI fix in #19201) unless the model map advertises
# native structured output
from litellm.utils import supports_native_structured_output
original_model: Final = model
if "response_format" in non_default_params:
if "response_format" in non_default_params and not supports_native_structured_output(
model=model, custom_llm_provider="bedrock"
):
# Use a model name that forces tool-based approach
model = "claude-3-sonnet-20240229"
@ -101,6 +107,16 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
# Restore original model name
model = original_model
# The stub model hides the original model from the parent's forced-tool-use backstop
response_format_tool_choice: Final = optional_params.get("tool_choice")
if (
"response_format" in non_default_params
and isinstance(response_format_tool_choice, dict)
and response_format_tool_choice.get("name") == RESPONSE_FORMAT_TOOL_NAME
and AnthropicModelInfo.forced_tool_use_unsupported(original_model)
):
optional_params.pop("tool_choice")
return optional_params
@staticmethod

65
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -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
@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
drop_tool_reference_parts_from_tool_messages,
get_tool_call_names,
hoist_images_from_tool_messages,
tool_with_flattened_parameters,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
async_convert_url_to_base64,
@ -65,6 +67,9 @@ else:
LiteLLMLoggingObj = Any
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
"""
Reference: https://platform.openai.com/docs/api-reference/chat/create
@ -397,6 +402,21 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
)
return messages, tools
def _targets_openai_hosted_endpoint(
self,
custom_llm_provider: str | None,
api_base: str | None,
) -> bool:
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,
@ -408,15 +428,34 @@ 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.
"""
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,
@ -443,11 +482,14 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
optional_params["tools"] = tools
optional_params.pop("max_retries", None)
if not optional_params.get("tools") and not optional_params.get("functions"):
optional_params.pop("tool_choice", None)
return {
"model": model,
"messages": messages,
**optional_params,
**self._flattened_tools_update_for_openai(optional_params, litellm_params),
}
async def async_transform_request(
@ -473,10 +515,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
if tools is not None and len(tools) > 0:
optional_params["tools"] = tools
if self.__class__._is_base_class:
if not optional_params.get("tools") and not optional_params.get("functions"):
optional_params.pop("tool_choice", None)
return {
"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

30
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -104,6 +104,10 @@ class ResponsesStreamChunk(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
delta: ReadOnly[str]
item_id: ReadOnly[str]
output_index: ReadOnly[int]
content_index: ReadOnly[int]
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
@ -680,8 +684,32 @@ class OpenAIResponsesHandler(BaseTranslation):
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
"""
Get the string so far from the responses so far.
``response.output_text.done`` events carry the whole part in ``text``, while
``response.output_text.delta`` events carry fragments in ``delta``. A stream
that dies before its done event (``response.failed`` / ``response.incomplete``)
has text only in deltas, so per content part the done text wins when present
and the joined deltas fill in otherwise, never both.
"""
return "".join([response.get("text", "") for response in responses_so_far])
keyed_events: Final = tuple(
(
(event.get("item_id"), event.get("output_index"), event.get("content_index")),
event.get("text"),
event.get("delta"),
)
for event in responses_so_far
if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
)
def part_text(part_key: tuple[object, object, object]) -> str:
done_texts: Final = tuple(
text for key, text, _ in keyed_events if key == part_key and isinstance(text, str)
)
if done_texts:
return done_texts[-1]
return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str))
return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events))
def _has_text_content(self, response: "ResponsesAPIResponse") -> bool:
"""

27
litellm/llms/openai/responses/transformation.py
View file

@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
)
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import *
from litellm.types.responses.main import *
@ -36,6 +37,7 @@ else:
_NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
_MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4")
_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
class _DeleteResponseBody(TypedDict):
@ -210,8 +212,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
)
if sanitized_tools is not None:
response_api_optional_request_params["tools"] = sanitized_tools
replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input)
final_request_params: Final = dict(
ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params)
ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params)
)
return final_request_params
@ -248,6 +251,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
return input, tools
def _drop_foreign_tool_call_item_ids(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
if self.custom_llm_provider not in _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS or not isinstance(input, list):
return input
sanitized_items: Final = [self._without_foreign_tool_call_item_id(item) for item in input]
return cast("ResponseInputParam", sanitized_items) # cast-ok: items keep their shape, minus a rejected id
@staticmethod
def _without_foreign_tool_call_item_id(item: object) -> object:
if not isinstance(item, dict):
return item
item_type: Final = item.get("type")
item_id: Final = item.get("id")
genuine_prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type) if isinstance(item_type, str) else None
if genuine_prefix is None or not isinstance(item_id, str) or item_id.startswith(genuine_prefix):
return item
return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item
def _flatten_tool_schema_combinators_for_openai(
self,
model: str,
@ -773,7 +793,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
)
if sanitized_tools is not None:
response_api_optional_request_params["tools"] = sanitized_tools
data: Final = dict(ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params))
replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input)
data: Final = dict(
ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params)
)
return url, data

15
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py
View file

@ -153,14 +153,17 @@ class VertexAIAnthropicConfig(AnthropicConfig):
drop_params: bool,
) -> dict:
"""
Override parent method to ensure VertexAI always uses tool-based structured outputs.
VertexAI doesn't support the output_format parameter, so we force all models
to use the tool-based approach for structured outputs.
Override parent method so VertexAI uses tool-based structured outputs
unless the vertex map entry advertises native structured output
(``output_format``, which Vertex AI Claude forwards for those models).
"""
# Temporarily override model name to force tool-based approach
# This ensures Claude Sonnet 4.5 uses tools instead of output_format
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
original_model: Final = model
if "response_format" in non_default_params:
native_structured_output: Final = AnthropicModelInfo._get_provider_resolved_capability(
model, "supports_native_structured_output", "vertex_ai"
)
if "response_format" in non_default_params and native_structured_output is not True:
model = "claude-3-sonnet-20240229" # Use a model that will use tool-based approach
# Call parent method with potentially modified model name

11
litellm/model_prices_and_context_window_backup.json
View file

@ -1482,7 +1482,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",
@ -1557,7 +1557,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",
@ -1632,7 +1632,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",
@ -1707,7 +1707,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",
@ -3254,6 +3254,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@ -44132,6 +44133,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@ -44202,6 +44204,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,

32
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
@ -86,8 +88,6 @@ def _connection_error_message(exc: BaseException) -> str:
if MCP_AVAILABLE:
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
global_mcp_server_manager,
@ -1173,6 +1173,7 @@ if MCP_AVAILABLE:
transport=request.transport,
auth_type=request.auth_type,
mcp_info=request.mcp_info,
timeout=request.timeout,
command=request.command,
args=request.args,
env=request.env,
@ -1402,11 +1403,28 @@ if MCP_AVAILABLE:
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
async def _list_tools_operation(client):
async def _list_tools_session_operation(session):
return await session.list_tools()
list_tools_response: Final = await client.run_with_session(_list_tools_session_operation)
list_tools_result: Final[list[MCPTool]] = list_tools_response.tools
# Bound the whole pagination walk: without this the preview is limited only by the
# per-request timeout times the page cap. max() keeps the pre-pagination guarantee
# that a single slow page within the client timeout still succeeds, and a
# per-server timeout above the global default extends the deadline with it.
listing_deadline: Final = max(
getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT,
MCP_TOOL_LISTING_TIMEOUT,
)
list_tools_result = None # rebind-ok: set inside the timeout scope below
with anyio.move_on_after(listing_deadline):
list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above
if list_tools_result is None:
verbose_logger.warning(
"MCP tools/list preview timed out after %s seconds while paginating upstream tools",
listing_deadline,
)
return { # mutable-ok: error response payload
"status": "error",
"error": True,
"message": f"Timed out listing tools after {listing_deadline} seconds. "
"The MCP server may be responding slowly or paginating excessively.",
}
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
return {
"tools": model_dumped_tools,

19
litellm/proxy/_types.py
View file

@ -2436,9 +2436,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
database_socket_timeout: float | None = Field(
None,
description=(
"Prisma `socket_timeout` URL param (seconds). When set, an idle/slow "
"connection that has not produced data within this window is closed. "
"This is the main knob for capping idle DB connections from LiteLLM."
"Prisma `socket_timeout` URL param (seconds). When set, an in-flight "
"operation that has not produced data within this window is aborted. "
"For capping how long idle pooled connections are kept, see "
"`database_max_idle_connection_lifetime`."
),
)
database_max_idle_connection_lifetime: float | None = Field(
60,
description=(
"Prisma `max_idle_connection_lifetime` URL param (seconds). A pooled "
"connection idle longer than this is closed and replaced instead of "
"being handed to the next request. Defaults to 60 so connections are "
"recycled before common infra idle timeouts (AWS NLB / RDS Proxy "
"~350s, many LBs 60-350s) silently drop them and requests fail with "
"`Error { kind: Closed }`. A value pinned on the DATABASE_URL or set "
"via `database_extra_connection_params` takes precedence."
),
)
database_extra_connection_params: dict[str, Any] | None = Field(

20
litellm/proxy/db/db_url_settings.py
View file

@ -82,10 +82,30 @@ CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset(
"pool_timeout",
"connect_timeout",
"socket_timeout",
"max_idle_connection_lifetime",
"pgbouncer",
}
)
# Quaint never tests pooled connections on checkout and keeps them idle for
# 300s by default, past many infra idle timeouts, so dead sockets surface as
# `Error { kind: Closed }`. 60s recycles them first; explicit values win.
DEFAULT_MAX_IDLE_CONNECTION_LIFETIME: Final = 60
IDLE_LIFETIME_DEFAULT_PARAMS: Final[Mapping[str, int]] = MappingProxyType(
{"max_idle_connection_lifetime": DEFAULT_MAX_IDLE_CONNECTION_LIFETIME}
)
def idle_lifetime_params(configured: float | None) -> Mapping[str, str | int | float]:
"""The `max_idle_connection_lifetime` to add to URLs that do not pin one.
Applied via ``add_missing_query_params`` so a URL-pinned value always wins,
whether the operator configured `database_max_idle_connection_lifetime` or not.
"""
if configured is None:
return IDLE_LIFETIME_DEFAULT_PARAMS
return MappingProxyType({"max_idle_connection_lifetime": configured})
def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) -> str:
"""Return ``url`` with the ``params`` it does not already carry appended.

34
litellm/proxy/guardrails/guardrail_hooks/alice/__init__.py Normal file
View file

@ -0,0 +1,34 @@
from typing import TYPE_CHECKING, Final
from litellm.types.guardrails import SupportedGuardrailIntegrations
from .alice import AliceGuardrail
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
import litellm
_alice_guardrail_callback: Final = AliceGuardrail(
api_key=litellm_params.api_key,
api_base=litellm_params.api_base,
unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"),
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm.logging_callback_manager.add_litellm_callback(_alice_guardrail_callback)
return _alice_guardrail_callback
guardrail_initializer_registry: Final = { # mutable-ok: module-level registry, built once and never mutated
SupportedGuardrailIntegrations.ALICE.value: initialize_guardrail,
}
guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated
SupportedGuardrailIntegrations.ALICE.value: AliceGuardrail,
}

369
litellm/proxy/guardrails/guardrail_hooks/alice/alice.py Normal file
View file

@ -0,0 +1,369 @@
# +-------------------------------------------------------------+
#
# Use Alice for your LLM calls
# https://alice.io/
#
# +-------------------------------------------------------------+
import json
import os
from collections.abc import Mapping
from typing import (
TYPE_CHECKING,
Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml
Final,
Literal,
Optional,
)
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException, Timeout
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
GUARDRAIL_NAME: Final = "alice"
_DEFAULT_API_BASE: Final = "https://api.alice.io"
_EVALUATE_PATH: Final = "/v2/evaluate/litellm"
_VERDICT_ALLOW: Final = "ALLOW"
_VERDICT_BLOCK: Final = "BLOCK"
_VERDICT_MASK: Final = "MASK"
_VERDICT_DETECT: Final = "DETECT"
_KNOWN_VERDICTS: Final = frozenset({_VERDICT_ALLOW, _VERDICT_BLOCK, _VERDICT_MASK, _VERDICT_DETECT})
_DEFAULT_BLOCK_MESSAGE: Final = "Blocked by your organization's content policy."
# apply_guardrail selects nothing: it forwards whichever of these came populated and lets Alice
# decide what is worth evaluating. Only skip the call when every one of them is empty — there is
# then genuinely nothing to send.
_SELECTABLE_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls", "structured_messages")
# Caps on the outbound copy of request_data. A payload deeper or wider than this is malformed
# rather than large, and serializing it would cost more than the evaluation it feeds.
_MAX_DEPTH: Final = 12
_MAX_ITEMS: Final = 5000
# request_data carries the caller's raw credentials under these keys, at any nesting depth —
# a real captured payload puts inbound headers at request_data["proxy_server_request"]["headers"],
# again under ["metadata"]["headers"] / ["litellm_metadata"]["headers"], and again under
# ["metadata"]["requester_metadata"]["headers"], any of which can carry an Authorization or
# x-api-key value. LiteLLM's own spend-log sanitizer excludes `secret_fields` for the same reason
# (spend_tracking_utils._SENSITIVE_REQUEST_BODY_KEYS): `secret_fields.raw_headers` holds the
# caller's Authorization / x-api-key in the clear, and `api_key` can carry a forwarded provider
# credential. Stripping by key name rather than by path means a new nesting path can never
# reintroduce the leak. Posting any of these to a third-party guardrail endpoint would be worse
# than what the proxy already refuses to persist in its own audit trail — so none of them leave
# the process.
_CREDENTIAL_KEYS_TO_STRIP: Final = frozenset(
{"secret_fields", "api_key", "raw_headers", "headers", "provider_specific_header"}
)
class AliceReplacement(TypedDict):
"""A masked substitution, positional against the texts that were submitted."""
index: ReadOnly[NotRequired[int]]
text: ReadOnly[NotRequired[str]]
class AliceVerdict(TypedDict):
"""Body returned by Alice's LiteLLM evaluate endpoint."""
verdict: ReadOnly[NotRequired[str]]
categories: ReadOnly[NotRequired["tuple[str, ...]"]]
correlation_id: ReadOnly[NotRequired[str]]
message: ReadOnly[NotRequired[str]]
replacements: ReadOnly[NotRequired["tuple[AliceReplacement, ...]"]]
class AliceGuardrailMissingSecrets(Exception):
"""Raised when the Alice API key is not configured."""
class AliceGuardrail(CustomGuardrail):
"""
Alice — policy-based guardrails for prompts and model responses.
This forwards the hook's arguments as it received them and enforces the verdict that comes
back, with one deliberate exception: any key named `secret_fields`, `api_key`, `raw_headers`,
`headers`, or `provider_specific_header` is dropped from `request_data` at any nesting depth
before it is serialized, and never reaches Alice. Short of that, it selects nothing and
renames nothing: which parts of a conversation are worth evaluating, and how a verdict is
reached, are decided by Alice — so changing either is a change on their side rather than a
LiteLLM upgrade. A batch with nothing selectable at all (no `texts`, `images`, `tools`,
`tool_calls`, or `structured_messages`) still skips the call, since there would be nothing to
send.
Known limitation: the unified guardrail's `streaming_transform_mode` defaults to
`block_only`, whose streaming path discards any returned text rewrite. A MASK verdict is
therefore a no-op on a streamed response — the original, unmasked text still reaches the
caller — while BLOCK continues to function on both streamed and non-streamed responses.
This is `during_call`'s documented behavior generally, not specific to Alice; configure a
masking-aware `streaming_transform_mode` if that gap matters for your traffic.
Alice evaluates against policies configured per *application*, and one proxy typically fronts
several, so the application is named on the virtual key rather than in this config:
curl $PROXY/key/generate -H "Authorization: Bearer $LITELLM_MASTER_KEY" \\
-d '{"key_alias": "payments-bot",
"metadata": {"alice_app_id": "payments-bot"}}'
Alice reads that off the authenticated key. Because the proxy strips caller-supplied
`user_api_key_*` from the request before a guardrail sees it, a caller cannot point its own
traffic at an application with laxer policies than the one its key was issued for.
Configuration example (litellm config YAML):
guardrails:
- guardrail_name: alice
litellm_params:
guardrail: alice
mode: [pre_call, post_call]
api_key: os.environ/ALICE_API_KEY
api_base: https://api.alice.io # optional
unreachable_fallback: fail_closed # optional
"""
def __init__(
self,
api_key: str | None = None,
api_base: str | None = None,
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
**kwargs: Any, # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__, whose param list is wide and evolving
) -> None:
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
alice_api_key: Final = api_key or os.environ.get("ALICE_API_KEY")
if not alice_api_key:
raise AliceGuardrailMissingSecrets(
"Alice API key is required. Set the `ALICE_API_KEY` environment variable or "
"pass `api_key` in the guardrail config."
)
self.alice_api_key: str = alice_api_key
base: Final = (api_base or os.environ.get("ALICE_API_BASE") or _DEFAULT_API_BASE).rstrip("/")
self.api_base: str = f"{base}{_EVALUATE_PATH}"
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback
if "supported_event_hooks" not in kwargs:
kwargs["supported_event_hooks"] = [ # mutable-ok: CustomGuardrail.__init__ requires a list here
GuardrailEventHooks.pre_call,
GuardrailEventHooks.during_call,
GuardrailEventHooks.post_call,
]
super().__init__(**kwargs)
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object], # mutable-ok: overrides CustomGuardrail.apply_guardrail's plain-dict contract
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
if not any(inputs.get(field) for field in _SELECTABLE_INPUT_FIELDS):
return inputs
try:
verdict: AliceVerdict = await self._evaluate(
inputs=inputs, request_data=request_data, input_type=input_type
)
except Timeout as e:
return self._on_unreachable(e, inputs)
except httpx.HTTPStatusError as e:
status_code: Final = getattr(getattr(e, "response", None), "status_code", None)
# Any 5xx is an outage on Alice's side, not our misconfiguration — route the whole
# class through the configured policy. A 4xx (rejected credential, bad request) is
# ours to fix and must never fail open, so it is deliberately left to propagate.
if isinstance(status_code, int) and 500 <= status_code < 600:
return self._on_unreachable(e, inputs)
raise
except httpx.RequestError as e:
return self._on_unreachable(e, inputs)
except (json.JSONDecodeError, UnicodeDecodeError, TypeError) as e:
# A body that cannot be decoded, cannot be parsed as JSON, or parses to something
# other than an object, is as unreachable as a dropped connection: this deployment's
# policy decides, not a raw exception. UnicodeDecodeError is named explicitly because
# it is a sibling of JSONDecodeError under ValueError, not a subclass of it.
return self._on_unreachable(e, inputs)
return self._enforce(verdict, inputs)
async def _evaluate(
self,
inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, object],
input_type: str,
) -> AliceVerdict:
response: Final = await self.async_handler.post(
url=self.api_base,
json={ # mutable-ok: one-shot HTTP request body, never mutated after construction
"input_type": input_type,
"inputs": _json_safe(inputs),
"request_data": _json_safe(request_data, strip_keys=_CREDENTIAL_KEYS_TO_STRIP),
},
headers={ # mutable-ok: one-shot HTTP headers, never mutated after construction
"Content-Type": "application/json",
"af-api-key": self.alice_api_key,
},
)
response.raise_for_status()
body = response.json()
if not isinstance(body, dict):
raise TypeError("Alice returned a non-object body")
return body
def _enforce(self, verdict: AliceVerdict, inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs:
"""Act on the verdict. An answer we cannot read is treated as unavailable, never as a pass."""
name: Final = verdict.get("verdict")
if name not in _KNOWN_VERDICTS:
return self._on_unreachable(ValueError(f"unrecognized verdict: {name!r}"), inputs)
if name == _VERDICT_BLOCK:
raise GuardrailRaisedException(
guardrail_name=GUARDRAIL_NAME,
message=verdict.get("message") or _DEFAULT_BLOCK_MESSAGE,
should_wrap_with_default_message=False,
blocked_content=True,
)
if name == _VERDICT_DETECT:
# Recorded by Alice and allowed through. The correlation id is what ties this request
# to that record; the evaluated text itself is never logged.
verbose_proxy_logger.warning(
"Alice guardrail: detection recorded, request allowed (correlation_id=%s, categories=%s)",
verdict.get("correlation_id"),
verdict.get("categories"),
)
return inputs
if name == _VERDICT_MASK:
self._apply_replacements(verdict, inputs)
return inputs
def _apply_replacements(self, verdict: AliceVerdict, inputs: GenericGuardrailAPIInputs) -> None:
"""
Write each replacement onto the text it names.
Only `texts` is touched. The chat translation layer maps a returned `texts` list back onto
the request positionally, but takes a different branch entirely when `structured_messages`
comes back as a new object — which would drop these edits.
All-or-nothing: a single out-of-range or malformed replacement blocks the whole verdict
rather than being silently skipped, so content Alice meant to replace can never reach the
model unmasked alongside content that was replaced.
"""
texts: Final = inputs.get("texts") or [] # mutable-ok: empty-list fallback, replaced wholesale below
replacements: Final = verdict.get("replacements") or [] # mutable-ok: empty-list fallback for iteration only
if not replacements:
raise self._mask_rejected(verdict)
for replacement in replacements:
index = replacement.get("index")
text = replacement.get("text")
if not (isinstance(index, int) and isinstance(text, str) and 0 <= index < len(texts)):
raise self._mask_rejected(verdict)
texts[index] = text # mutable-ok: item assignment into the local working copy above
inputs["texts"] = texts
def _mask_rejected(self, verdict: AliceVerdict) -> GuardrailRaisedException:
"""A MASK verdict that cannot be applied in full is refused outright, never partially —
see `_apply_replacements`."""
return GuardrailRaisedException(
guardrail_name=GUARDRAIL_NAME,
message=verdict.get("message") or _DEFAULT_BLOCK_MESSAGE,
should_wrap_with_default_message=False,
blocked_content=True,
)
def _on_unreachable(self, error: Exception, inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs:
"""Apply the configured policy when Alice cannot be reached or cannot be understood."""
if self.unreachable_fallback == "fail_open":
verbose_proxy_logger.critical(
"Alice guardrail unreachable, allowing request per unreachable_fallback: %s",
error,
)
return inputs
raise GuardrailRaisedException(
guardrail_name=GUARDRAIL_NAME,
message="Alice guardrail is unavailable and this request cannot be checked",
should_wrap_with_default_message=False,
) from error
@staticmethod
def get_config_model() -> type | None:
from litellm.types.proxy.guardrails.guardrail_hooks.alice import (
AliceGuardrailConfigModel,
)
return AliceGuardrailConfigModel
def _json_safe(
value: object,
depth: int = 0,
seen: frozenset[int] = frozenset(),
strip_keys: frozenset[str] = frozenset(),
) -> object:
"""
Copy `value` into something `json.dumps` accepts, dropping only what cannot cross.
`request_data` carries live Python objects — an OpenTelemetry span among them — so it cannot
be serialized as it stands. What is dropped is decided by a mechanical rule rather than a
field list: a list drifts from what the far side needs, a rule cannot. Serializing naively
raises, and that error would be read as "guardrail unavailable" on every single request.
`strip_keys` drops a dict key by name at every depth it appears, not just the root — a caller
passes `_CREDENTIAL_KEYS_TO_STRIP` here so a credential nested under any path is caught the
same way a top-level one is, without maintaining a list of paths. The source object is never
mutated: every branch below builds a new container.
"""
if isinstance(value, (str, int, float, bool)) or value is None:
return value
if depth >= _MAX_DEPTH or id(value) in seen:
return None
nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately
if isinstance(value, dict):
out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is
for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view
if isinstance(key, str) and key not in strip_keys:
out[key] = _json_safe(item, depth + 1, nested, strip_keys)
return out
if isinstance(value, (list, tuple, set, frozenset)):
return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use
_json_safe(item, depth + 1, nested, strip_keys)
for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view
]
dump: Final = getattr(value, "model_dump", None)
if callable(dump):
try:
return _json_safe(dump(mode="json"), depth + 1, nested, strip_keys)
except Exception: # noqa: BLE001 # a model that will not dump is one we drop
return None
# Everything json.dumps handles natively — str, int, float, bool, None, dict, list — is
# caught above, and a dict/list subclass is caught by isinstance. So whatever reaches here
# (bytes, datetime, an OpenTelemetry span) cannot cross the wire.
return None

29
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -31,6 +31,7 @@ from litellm.caching import DualCache
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.litellm_core_utils.litellm_logging import (
_get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
@ -215,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
return redacted if isinstance(redacted, list) else assessments
_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
def _is_responses_api_route(request_route: str | None) -> bool:
if request_route is None:
return False
call_types: Final = get_call_types_for_route(request_route)
return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
@ -2709,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
yield streamed_chunk
return
# Responses-API events are neither chat-completions chunks nor raw
# Anthropic SSE, so the assembly below cannot scan them; the unified
# guardrail's translation layer can, with buffering semantics kept.
if _is_responses_api_route(user_api_key_dict.request_route):
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
guardrail_to_apply=self,
buffer_until_moderated_default=True,
):
yield translated_chunk
return
# Import here to avoid circular imports
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder

48
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -54,7 +54,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
coordination_redis_cache,
publish_config_change,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin
from litellm.proxy.management_endpoints.team_endpoints import (
@ -701,6 +704,12 @@ async def patch_model(
param="blocked",
)
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=patch_data.litellm_params,
user_api_key_dict=user_api_key_dict,
existing_litellm_params=db_model.litellm_params,
)
_raise_on_strategy_router_write_violation(
incoming_params=patch_data.litellm_params,
existing_params=db_model.litellm_params,
@ -1464,6 +1473,32 @@ class ModelManagementAuthChecks:
)
return True
@staticmethod
def can_user_attach_credential(
litellm_params: GenericLiteLLMParams | None,
user_api_key_dict: UserAPIKeyAuth,
existing_litellm_params: GenericLiteLLMParams | None = None,
) -> Literal[True]:
if litellm_params is None or litellm_params.litellm_credential_name is None:
return True
if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None:
existing_credential_name: Final = decrypt_value_helper(
value=existing_litellm_params.litellm_credential_name,
key="litellm_credential_name",
exception_type="debug",
return_original_value=True,
)
if litellm_params.litellm_credential_name == existing_credential_name:
return True
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
raise ProxyException(
message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.",
type=ProxyErrorTypes.auth_error.value,
code=status.HTTP_403_FORBIDDEN,
param="litellm_credential_name",
)
@staticmethod
async def allow_team_model_action(
model_params: Deployment | updateDeployment,
@ -1786,6 +1821,11 @@ async def add_new_model(
premium_user=premium_user,
)
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=model_params.litellm_params,
user_api_key_dict=user_api_key_dict,
)
_raise_on_strategy_router_write_violation(
incoming_params=model_params.litellm_params,
existing_params=None,
@ -1958,6 +1998,12 @@ async def update_model(
premium_user=premium_user,
)
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=model_params.litellm_params,
user_api_key_dict=user_api_key_dict,
existing_litellm_params=deployment.litellm_params,
)
_raise_on_strategy_router_write_violation(
incoming_params=model_params.litellm_params,
existing_params=deployment.litellm_params,

49
litellm/proxy/proxy_cli.py
View file

@ -913,14 +913,13 @@ class ProxyInitializationHelpers:
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
)
@click.option(
"--use_v2_migration_resolver/--use_legacy_migration_resolver",
default=True,
"--use_v2_migration_resolver",
is_flag=True,
default=False,
help=(
"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."
"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."
),
envvar="USE_V2_MIGRATION_RESOLVER",
)
@ -1226,6 +1225,7 @@ def run_server(
if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None:
from litellm.proxy.db.db_url_settings import (
add_missing_query_params,
idle_lifetime_params,
reader_shareable_params,
unsupported_db_scheme,
unsupported_db_scheme_message,
@ -1254,6 +1254,9 @@ def run_server(
disable_prepared_statements=db_disable_prepared_statements,
extra_params=db_extra_connection_params,
)
lifetime_params: Final = idle_lifetime_params(
general_settings.get("database_max_idle_connection_lifetime")
)
if os.getenv("DATABASE_URL", None) is not None:
database_url = get_secret("DATABASE_URL", default_value=None)
resolved_url: Final[str | None] = str(database_url) if database_url else None
@ -1271,11 +1274,11 @@ def run_server(
writer_url,
connection_url_params,
)
os.environ["DATABASE_URL"] = modified_url
os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, lifetime_params)
if os.getenv("DIRECT_URL", None) is not None:
database_url = os.getenv("DIRECT_URL")
modified_url = append_query_params(database_url, connection_url_params)
os.environ["DIRECT_URL"] = modified_url
os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, lifetime_params)
# The reader pool is a real pool against the same configured cap, so it
# gets the allowlisted pool params. Schema-affecting ones, including any
# the operator smuggled in through database_extra_connection_params, stay
@ -1289,10 +1292,13 @@ def run_server(
db_lock_timeout,
)
os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
_with_query_value(read_replica_url, "options", reader_options)
if reader_options
else read_replica_url,
reader_shareable_params(connection_url_params),
add_missing_query_params(
_with_query_value(read_replica_url, "options", reader_options)
if reader_options
else read_replica_url,
reader_shareable_params(connection_url_params),
),
lifetime_params,
)
subprocess.run(["prisma"], capture_output=True)
is_prisma_runnable = True
@ -1311,11 +1317,10 @@ def run_server(
else:
if not use_v2_migration_resolver:
print(
"\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"
"\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"
)
try:
setup_ok: Final = PrismaManager.setup_database(
@ -1323,10 +1328,10 @@ def run_server(
use_v2_resolver=use_v2_migration_resolver,
)
except RuntimeError as e:
# 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.
# 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.
print(
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m",
file=sys.stderr,

14
litellm/proxy/proxy_server.py
View file

@ -12481,11 +12481,21 @@ async def supported_openai_params(model: str):
--header 'Authorization: Bearer sk-1234'
```
"""
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
global llm_router
try:
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
resolved_models: Final = llm_router.resolved_litellm_models(model) if llm_router is not None else ()
target_model: Final = resolved_models[0] if resolved_models else model
declared_provider: Final = declared_authenticating_provider(target_model)
litellm_model, custom_llm_provider = (
(target_model.removeprefix(f"{declared_provider}/"), declared_provider)
if declared_provider is not None
else litellm.get_llm_provider(model=target_model)[:2]
)
return {
"supported_openai_params": litellm.get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
model=litellm_model, custom_llm_provider=custom_llm_provider
)
}
except Exception:

12
litellm/responses/litellm_completion_transformation/custom_tools.py
View file

@ -17,6 +17,7 @@ logic.
import json
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel, TypeAdapter, ValidationError
@ -28,6 +29,15 @@ from litellm.types.llms.openai import (
_MAX_ARGUMENTS_LEN: Final = 1_000_000
TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType({"function_call": "fc", "custom_tool_call": "ctc"})
def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str:
prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type)
if prefix is None or not tool_id or tool_id.startswith(prefix):
return tool_id
return f"{prefix}_{tool_id}"
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
"""Extract names of tools originally defined as ``type: "custom"``."""
@ -103,7 +113,7 @@ def build_tool_call_item_kwargs(
item_type: Final = "custom_tool_call" if custom else "function_call"
kwargs: Final[dict[str, str]] = {
"type": item_type,
"id": call_id,
"id": openai_shaped_tool_call_item_id(item_type, call_id),
"call_id": call_id,
"name": name,
"status": status,

10
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -114,6 +114,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = []
self._tool_output_index_by_call_id: dict[str, int] = {}
self._tool_args_by_call_id: dict[str, str] = {}
self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream
self._tool_call_id_by_index: dict[int, str] = {}
self._ambiguous_tool_call_indexes: set[int] = set()
self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item
@ -227,6 +228,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
event = OutputItemAddedEvent(
@ -248,7 +250,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=call_id,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
delta=delta_chunk,
)
@ -300,6 +302,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
event = OutputItemAddedEvent(
@ -325,7 +328,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
delta_event = FunctionCallArgumentsDeltaEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=call_id,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
delta=delta_chunk,
)
@ -335,7 +338,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
done_event = FunctionCallArgumentsDoneEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
item_id=call_id,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
arguments=final_args,
)
@ -345,6 +348,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names)
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
item_done_event = OutputItemDoneEvent(

8
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -93,6 +93,7 @@ from .custom_tools import (
convert_custom_tool_to_function_tool,
extract_custom_tool_names,
is_custom_tool_call,
openai_shaped_tool_call_item_id,
serialize_tool_call_arguments,
unwrap_custom_tool_arguments,
validated_allowed_callers,
@ -2034,7 +2035,7 @@ class LiteLLMCompletionResponsesConfig:
custom_item = CustomToolCallOutputItem(
type="custom_tool_call",
call_id=tool_id,
id=tool_id,
id=openai_shaped_tool_call_item_id("custom_tool_call", tool_id),
name=tool_name,
input=input_str,
status=function_definition.get("status") or "completed",
@ -2065,7 +2066,7 @@ class LiteLLMCompletionResponsesConfig:
name=tool_name,
arguments=tool_arguments,
call_id=tool_id,
id=tool_id,
id=openai_shaped_tool_call_item_id("function_call", tool_id),
type="function_call",
status=function_definition.get("status") or "completed",
)
@ -2502,8 +2503,7 @@ class LiteLLMCompletionResponsesConfig:
choice=choice,
)
message_output_items.extend(image_generation_items)
else:
# Regular message output
elif choice.message.content is not None:
message_output_items.append(
GenericResponseOutputItem(
type="message",

61
litellm/router.py
View file

@ -8165,6 +8165,52 @@ class Router:
if backend_value is not None:
model_info[field] = backend_value
@staticmethod
def _inherit_builtin_base_rates_for_off_peak(
model_info: dict, # mutable-ok: cost-map entry filled in place
backend_model: str,
custom_llm_provider: str | None,
) -> None:
"""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. 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
if any(
model_info.get(field) is not None
for field in ("input_cost_per_token", "input_cost_per_second", "tiered_pricing")
):
return
try:
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
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
model_info[field] = copy.deepcopy(backend_value)
@staticmethod
def _inherit_builtin_tiered_output_rate(
model_info: dict, backend_model: str, custom_llm_provider: str | None
@ -8253,6 +8299,11 @@ class Router:
if deployment.litellm_params.get(field) is not None:
_model_info[field] = deployment.litellm_params[field]
Router._inherit_builtin_base_rates_for_off_peak(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if _model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info,
@ -8994,6 +9045,11 @@ class Router:
if field_value is not None:
_model_info_dict[field] = field_value
Router._inherit_builtin_base_rates_for_off_peak(
model_info=_model_info_dict,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if _model_info_dict.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info_dict,
@ -9249,6 +9305,11 @@ class Router:
field_value = deployment.litellm_params.get(field)
if field_value is not None:
model_info[field] = field_value
Router._inherit_builtin_base_rates_for_off_peak(
model_info=model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=model_info,

27
litellm/router_utils/pattern_match_deployments.py
View file

@ -8,7 +8,7 @@ from re import Match
from typing import Final
from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
class PatternUtils:
@ -204,7 +204,7 @@ class PatternMatchRouter:
return litellm_deployment_litellm_model
def get_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict] | None:
def get_pattern(self, model: str | None, custom_llm_provider: str | None = None) -> list[dict] | None:
"""
Check if a pattern exists for the given model and custom llm provider
@ -215,18 +215,17 @@ class PatternMatchRouter:
Returns:
bool: True if pattern exists, False otherwise
"""
if custom_llm_provider is None:
try:
(
_,
custom_llm_provider,
_,
_,
) = get_llm_provider(model=model)
except Exception:
# get_llm_provider raises exception when provider is unknown
pass
return self.route(model) or self.route(f"{custom_llm_provider}/{model}")
provider: Final = (
custom_llm_provider or declared_authenticating_provider(model) or self._resolved_provider(model)
)
return self.route(model) or self.route(f"{provider}/{model}")
@staticmethod
def _resolved_provider(model: str | None) -> str | None:
try:
return get_llm_provider(model=model)[1] if model else None
except Exception: # noqa: BLE001 # get_llm_provider raises when the provider is unknown; the name then routes as-is
return None
def get_deployments_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict]:
"""

1
litellm/types/guardrails.py
View file

@ -136,6 +136,7 @@ class SupportedGuardrailIntegrations(Enum):
HEADROOM = "headroom"
COMPRESR = "compresr"
STRAIKER = "straiker"
ALICE = "alice"
class Role(Enum):

21
litellm/types/proxy/guardrails/guardrail_hooks/alice.py Normal file
View file

@ -0,0 +1,21 @@
from pydantic import Field
from .base import GuardrailConfigModel
class AliceGuardrailConfigModel(GuardrailConfigModel):
api_key: str | None = Field(
default=None,
description=("The API key for Alice. If not provided, the `ALICE_API_KEY` environment variable is checked."),
)
api_base: str | None = Field(
default=None,
description=(
"The API base URL for Alice. If not provided, the `ALICE_API_BASE` environment "
"variable is checked, then `https://api.alice.io`."
),
)
@staticmethod
def ui_friendly_name() -> str:
return "Alice"

50
litellm/types/utils.py
View file

@ -193,6 +193,38 @@ class AgenticLoopParams(TypedDict, total=False):
"""The LLM provider name (e.g., 'bedrock', 'anthropic')"""
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; 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]
class ModelInfoBase(ProviderSpecificModelInfo, total=False):
key: Required[str] # the key in litellm.model_cost which is returned
@ -225,6 +257,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: 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
@ -3486,17 +3519,22 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
return {k: v for k, v in model_info.items() if k not in cls.model_fields}
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = frozenset(
ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__
) - frozenset(CustomPricingLiteLLMParams.model_fields)
DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
- frozenset(CustomPricingLiteLLMParams.model_fields)
- DEPLOYMENT_SCOPED_PRICING_FIELDS
)
def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]:
"""Return only the fields safe to register under a shared ``{provider}/{model}``
key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus
per-deployment pricing overrides. Per-deployment metadata (``id``,
``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key;
it stays under the deployment's unique model id.
per-deployment pricing overrides and deployment-scoped pricing blocks such as
``off_peak_pricing``. Per-deployment metadata (``id``, ``access_via_team_ids``,
arbitrary custom keys) never belongs on the shared key; it stays under the
deployment's unique model id.
"""
return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}

6
litellm/utils.py
View file

@ -2869,10 +2869,9 @@ def _update_dictionary(existing_dict: dict, new_dict: dict) -> dict:
elif isinstance(v, dict):
existing_nested_dict = existing_dict.get(k)
if isinstance(existing_nested_dict, dict):
existing_nested_dict.update(v)
existing_dict[k] = existing_nested_dict
existing_dict[k] = {**existing_nested_dict, **v} # mutable-ok: copy-on-write merge
else:
existing_dict[k] = v
existing_dict[k] = dict(v) # mutable-ok: detached copy, never the caller's dict by reference
else:
existing_dict[k] = v
@ -5860,6 +5859,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),

11
model_prices_and_context_window.json
View file

@ -1482,7 +1482,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",
@ -1557,7 +1557,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",
@ -1632,7 +1632,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",
@ -1707,7 +1707,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",
@ -3254,6 +3254,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@ -44132,6 +44133,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@ -44202,6 +44204,7 @@
"supports_computer_use": true,
"supports_forced_tool_use": false,
"supports_function_calling": true,
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,

4
ruff-strict.toml
View file

@ -26,6 +26,10 @@ external = [
# caught a real mismatch, confirming Any is correct here, not a shortcut.
"litellm/litellm_core_utils/litellm_logging.py" = ["ANN401"]
"litellm/utils.py" = ["ANN401"]
# `**kwargs` forwards verbatim to CustomGuardrail.__init__, whose param list is wide and
# grows over time; typing it concretely (`object`) broke that forwarding call outright —
# basedpyright turned every named param into a reportArgumentType error. Any is correct here.
"litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"]
[lint.mccabe]
max-complexity = 15

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -63,6 +63,7 @@ IGNORE_FUNCTIONS = [
"json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks.
"_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible).
"_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible).
"_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input.
]

1
tests/code_coverage_tests/router_code_coverage.py
View file

@ -81,6 +81,7 @@ ignored_function_names = [
"_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name)
"_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name)
"has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call
"_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name)
]

20
tests/e2e/ui/helpers/premium.ts Normal file
View file

@ -0,0 +1,20 @@
import * as fs from "fs";
import { ADMIN_STORAGE_PATH } from "../constants";
/**
* Whether the proxy under test is licensed, read from the admin session JWT's `premium_user`
* claim. That is the same value the dashboard reads to enable premium-gated controls, so it
* describes the proxy Playwright is pointed at rather than the environment the runner happens
* to have, which are not the same machine when E2E_UI_BASE_URL points elsewhere.
*/
export function proxyIsPremium(): boolean {
const storage = JSON.parse(fs.readFileSync(ADMIN_STORAGE_PATH, "utf-8")) as {
cookies?: { name: string; value: string }[];
};
const token = storage.cookies?.find((cookie) => cookie.name === "token")?.value;
const payload = token?.split(".")[1];
if (!payload) {
return false;
}
return JSON.parse(Buffer.from(payload, "base64url").toString("utf-8")).premium_user === true;
}

73
tests/e2e/ui/helpers/traffic.ts
View file

@ -4,12 +4,16 @@ import { APIRequestContext, expect } from "@playwright/test";
export const CHAT_MODEL_A = "fake-openai-gpt-4";
export const CHAT_MODEL_B = "fake-anthropic-claude";
/** The deployment each of those models routes to, as spend logs and usage breakdowns name it. */
export const DEPLOYMENT_MODEL_A = "openai/fake-gpt-4";
export const DEPLOYMENT_MODEL_B = "openai/fake-claude";
/** The only completion text fixtures/mock_llm_server/server.py ever returns. */
export const MOCK_RESPONSE_TEXT = "This is a mock response.";
export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-1234";
const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? "";
export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? "";
interface ChatOptions {
model: string;
@ -114,13 +118,54 @@ export async function waitForSpendLogByPrompt(
const isoDay = (d: Date): string => d.toISOString().slice(0, 10);
interface DailyActivityKey {
metrics?: { api_requests?: number };
}
interface DailyActivityPage {
results?: { breakdown?: { api_keys?: Record<string, DailyActivityKey> } }[];
metadata?: { total_pages?: number };
}
const requestsOnPage = (body: DailyActivityPage, keyToken: string): number =>
(body.results ?? []).reduce((sum, day) => sum + (day.breakdown?.api_keys?.[keyToken]?.metrics?.api_requests ?? 0), 0);
/**
* The route paginates its per-key breakdown. Reading only the first page finds a key while the
* database is small and stops finding it once a run has generated more keys than one page holds,
* which reads as "the rollup is not running" when the rollup is fine.
*/
async function keyRequestsInDailyActivity(
request: APIRequestContext,
query: string,
keyToken: string,
page = 1,
seen = 0,
): Promise<number> {
const res = await request.get(`${rootPath()}/user/daily/activity?${query}&page=${page}`, {
headers: { Authorization: `Bearer ${masterKey()}` },
});
if (!res.ok()) {
return seen;
}
const body = (await res.json()) as DailyActivityPage;
const total = seen + requestsOnPage(body, keyToken);
return page >= (body.metadata?.total_pages ?? 1)
? total
: keyRequestsInDailyActivity(request, query, keyToken, page + 1, total);
}
/**
* The Usage page reads /user/daily/activity, a rollup written by a background job, and fetches it once
* on mount. Navigating before the rollup lands leaves a stale render that never refreshes.
*
* The rollup lands request by request, so waiting only for the key to appear leaves a caller that
* sent several requests reading a partial count. Pass `minRequests` to wait for all of them.
*/
export async function waitForKeyInDailyActivity(
request: APIRequestContext,
keyToken: string,
minRequests = 1,
timeoutMs = 120_000,
): Promise<void> {
const now = new Date();
@ -129,25 +174,17 @@ export async function waitForKeyInDailyActivity(
const query = `start_date=${isoDay(start)}&end_date=${isoDay(now)}`;
const deadline = Date.now() + timeoutMs;
let lastStatus = 0;
while (Date.now() < deadline) {
const res = await request.get(`${rootPath()}/user/daily/activity?${query}`, {
headers: { Authorization: `Bearer ${masterKey()}` },
});
lastStatus = res.status();
if (res.ok()) {
const body = await res.json();
const seen = (body?.results ?? []).some(
(day: { breakdown?: { api_keys?: Record<string, unknown> } }) => keyToken in (day.breakdown?.api_keys ?? {}),
for (;;) {
const seen = await keyRequestsInDailyActivity(request, query, keyToken);
if (seen >= minRequests) {
return;
}
if (Date.now() >= deadline) {
throw new Error(
`key ${keyToken} reached ${seen} of ${minRequests} requests in /user/daily/activity across every page; ` +
"the daily spend rollup may not be running",
);
if (seen) {
return;
}
}
await new Promise((r) => setTimeout(r, 3_000));
}
throw new Error(
`key ${keyToken} never appeared in /user/daily/activity (last status ${lastStatus}); ` +
"the daily spend rollup may not be running",
);
}

152
tests/e2e/ui/tests/logs/logsFilters.spec.ts Normal file
View file

@ -0,0 +1,152 @@
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
import {
CHAT_MODEL_A,
CHAT_MODEL_B,
createVirtualKey,
sendChatCompletion,
waitForSpendLog,
} from "../../helpers/traffic";
/**
* Every test mints its own key and asserts against request ids it generated, so a filter that
* quietly does nothing shows up as the other key's row still being on screen, and concurrent
* specs' traffic cannot decide the outcome.
*/
const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */
const requestLogsRows = (page: PlaywrightPage): Locator =>
page.locator("table").filter({ visible: true }).first().locator("tbody tr");
const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
async function openLogs(page: PlaywrightPage): Promise<void> {
await navigateToPage(page, Page.Logs);
await dismissFeedbackPopup(page);
await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
}
async function openFilterDrawer(page: PlaywrightPage): Promise<Locator> {
await visibleTestId(page, "datatable-filters-trigger").click();
const drawer = page.getByRole("dialog", { name: "Filters" });
await expect(drawer).toBeVisible({ timeout: 10_000 });
return drawer;
}
/** Picks a value in one of the drawer's searchable comboboxes and applies the filter. */
async function applyComboboxFilter(
page: PlaywrightPage,
drawer: Locator,
comboboxLabel: string,
value: string,
): Promise<void> {
await drawer.getByRole("combobox", { name: comboboxLabel }).click();
await page.keyboard.type(value);
await page.getByRole("option", { name: value, exact: true }).first().click();
await drawer.getByRole("button", { name: "Apply Filters" }).click();
await expect(drawer).not.toBeVisible({ timeout: 10_000 });
}
/** A request the key is not entitled to make, so the proxy refuses it and logs the refusal. */
async function sendDeniedCompletion(request: APIRequestContext, apiKey: string): Promise<void> {
const res = await request.post("/v1/chat/completions", {
headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
data: { model: CHAT_MODEL_B, messages: [{ role: "user", content: "denied" }] },
});
expect(res.status(), "a model outside the key's allow-list is refused").toBe(403);
}
test.describe("Logs page filters", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("the Key Alias filter narrows the table to that key's requests", async ({ page, request }) => {
const suffix = uniqueSuffix();
const mine = await createVirtualKey(request, { key_alias: `e2e-logs-mine-${suffix}` });
const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-theirs-${suffix}` });
const myRequestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `logs-filter-mine-${suffix}`,
apiKey: mine.key,
});
const theirRequestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `logs-filter-theirs-${suffix}`,
apiKey: theirs.key,
});
await waitForSpendLog(request, myRequestId);
await waitForSpendLog(request, theirRequestId);
await openLogs(page);
const drawer = await openFilterDrawer(page);
await applyComboboxFilter(page, drawer, "Search a key alias", mine.alias!);
await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 30_000 });
// The filter is only doing its job if the other key's request is gone, not merely if ours is present.
await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(0, { timeout: 10_000 });
});
test("the Status filter narrows the table to the refused request", async ({ page, request }) => {
const suffix = uniqueSuffix();
const alias = `e2e-logs-status-${suffix}`;
const scoped = await createVirtualKey(request, { key_alias: alias, models: [CHAT_MODEL_A] });
const servedRequestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `logs-filter-served-${suffix}`,
apiKey: scoped.key,
});
await sendDeniedCompletion(request, scoped.key);
await waitForSpendLog(request, servedRequestId);
await openLogs(page);
const drawer = await openFilterDrawer(page);
await drawer.getByRole("combobox", { name: "Search a key alias" }).click();
await page.keyboard.type(alias);
await page.getByRole("option", { name: alias, exact: true }).first().click();
// The Status field labels its group, not the trigger, so it is addressed by the value it shows.
await drawer.getByRole("combobox").filter({ hasText: "All Statuses" }).click();
await page.getByRole("option", { name: "Failure", exact: true }).click();
await drawer.getByRole("button", { name: "Apply Filters" }).click();
await expect(drawer).not.toBeVisible({ timeout: 10_000 });
// Both requests were made by this key, so a Status filter that does nothing leaves the served one on screen.
await expect(requestLogsRows(page)).toHaveCount(1, { timeout: 30_000 });
await expect(requestLogsRows(page)).toContainText("Failure");
await expect(requestLogsRows(page).filter({ hasText: servedRequestId })).toHaveCount(0);
});
test("Reset Filters brings back the rows a filter hid", async ({ page, request }) => {
const suffix = uniqueSuffix();
const mine = await createVirtualKey(request, { key_alias: `e2e-logs-reset-mine-${suffix}` });
const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-reset-theirs-${suffix}` });
const myRequestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `logs-reset-mine-${suffix}`,
apiKey: mine.key,
});
const theirRequestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `logs-reset-theirs-${suffix}`,
apiKey: theirs.key,
});
await waitForSpendLog(request, myRequestId);
await waitForSpendLog(request, theirRequestId);
await openLogs(page);
const drawer = await openFilterDrawer(page);
await applyComboboxFilter(page, drawer, "Search a key alias", mine.alias!);
await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(0, { timeout: 30_000 });
// A filter you cannot clear is a page that looks empty forever, which is how it reads to a user.
await page.getByRole("button", { name: "Reset Filters" }).filter({ visible: true }).click();
await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(1, { timeout: 30_000 });
await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 10_000 });
});
});

34
tests/e2e/ui/tests/modelsPage/addModel.spec.ts
View file

@ -5,6 +5,11 @@ import { navigateToPage } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
import { captureRequestBody, readBack } from "../../helpers/roundTrip";
import { sendChatCompletion } from "../../helpers/traffic";
import { proxyIsPremium } from "../../helpers/premium";
/** Four probes 13s apart span 39s, one PROXY_CONFIG_RELOAD_INTERVAL_SECONDS (30s) plus margin. */
const CREDENTIAL_PROBE_SUCCESSES = 4;
const CREDENTIAL_PROBE_SPACING_MS = 13_000;
/** The mock LLM as the proxy reaches it: same host locally, a sidecar in the deployed stack. */
const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`;
@ -35,7 +40,10 @@ async function selectProvider(page: PlaywrightPage, providerName: string) {
const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true });
await providerDropdown.click();
await providerDropdown.fill(providerName);
await page.getByRole("option").filter({ hasText: exactly(providerName) }).click();
await page
.getByRole("option")
.filter({ hasText: exactly(providerName) })
.click();
await expect(providerDropdown).toHaveValue(providerName);
}
@ -78,6 +86,9 @@ test.describe("Add Model", () => {
});
test("Edit team model TPM and RPM limits", async ({ page }) => {
// /model/new refuses a team-scoped deployment on an unlicensed proxy, so there this fails in
// setup on a product gate rather than on a regression in the edit it covers.
test.skip(!proxyIsPremium(), "proxy under test is unlicensed — team-scoped models are premium");
const masterKey = users[Role.ProxyAdmin].password;
const modelName = `e2e-team-model-${Date.now()}`;
@ -226,8 +237,11 @@ test.describe("Add Model", () => {
});
expect(createCred.ok(), `POST /credentials failed (${createCred.status()}): ${await createCred.text()}`).toBe(true);
// Multi-instance stacks propagate a new credential to the probe-serving instances on a periodic
// sync; consecutive successes guard against a load balancer alternating synced and stale replicas
// The proxy's periodic credential refresh prunes its in-memory list against a database snapshot
// it took before this credential landed, so a credential that resolves right after POST
// /credentials can stop resolving until the refresh after that. Successes spanning a whole
// PROXY_CONFIG_RELOAD_INTERVAL_SECONDS prove it survived a refresh, after which it stays.
// Resolution fails open onto the ambient key, so losing it reads as a confusing upstream 404.
let consecutiveProbeSuccesses = 0;
await expect
.poll(
@ -249,11 +263,12 @@ test.describe("Add Model", () => {
return consecutiveProbeSuccesses;
},
{
message: `stored credential ${credentialName} never became usable for a connection test`,
timeout: 60_000,
message: `stored credential ${credentialName} never stayed usable across a config reload`,
intervals: [0, CREDENTIAL_PROBE_SPACING_MS],
timeout: 110_000,
},
)
.toBeGreaterThanOrEqual(3);
.toBeGreaterThanOrEqual(CREDENTIAL_PROBE_SUCCESSES);
try {
await navigateToPage(page, Page.Models);
@ -458,7 +473,7 @@ test.describe("Add Model", () => {
await page.waitForLoadState("networkidle");
await page.getByPlaceholder("Search model names").fill("cohere");
// Clearer failure than timing out on a row assertion when the table is empty.
await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, {
timeout: 15_000,
@ -466,10 +481,7 @@ test.describe("Add Model", () => {
// Pin to one row carrying both the name and the team, so the sibling test's
// team-less cohere row can't satisfy it.
const teamCohereRow = page
.getByRole("row")
.filter({ hasText: "cohere/" })
.filter({ hasText: E2E_TEAM_CRUD_ID });
const teamCohereRow = page.getByRole("row").filter({ hasText: "cohere/" }).filter({ hasText: E2E_TEAM_CRUD_ID });
await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 });
} finally {
await deleteTeamScopedCohereModels();

32
tests/e2e/ui/tests/proxy-admin/keys.spec.ts
View file

@ -1,15 +1,17 @@
import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
import {
ADMIN_STORAGE_PATH,
E2E_DELETE_KEY_ALIAS,
E2E_REGENERATE_KEY_ALIAS,
E2E_UPDATE_LIMITS_KEY_ALIAS,
E2E_INTERNAL_USER_KEY_ALIAS,
E2E_TEAM_CRUD_ALIAS,
E2E_TEAM_CRUD_ID,
} from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
import { captureRequestBody, readBack } from "../../helpers/roundTrip";
import { masterKey } from "../../helpers/traffic";
import { proxyIsPremium } from "../../helpers/premium";
/**
* Looks a key up by alias, undefined when none carries it. `return_full_object=true` is what makes
@ -23,6 +25,17 @@ async function findKeyByAlias(page: PlaywrightPage, alias: string): Promise<Reco
return body.keys.find((row) => row.key_alias === alias);
}
/** A key this test owns, so deleting it costs the suite nothing on a retry or a second run. */
async function createDeletableKey(page: PlaywrightPage): Promise<string> {
const alias = `e2e-delete-key-${Date.now()}`;
const res = await page.request.post("/key/generate", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { key_alias: alias, team_id: E2E_TEAM_CRUD_ID },
});
expect(res.ok(), `POST /key/generate failed (${res.status()}): ${await res.text()}`).toBe(true);
return alias;
}
test.describe("Proxy Admin - Keys", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
@ -67,6 +80,9 @@ test.describe("Proxy Admin - Keys", () => {
});
test("Regenerate key", async ({ page }) => {
// The Regenerate Key button renders disabled when the proxy is unlicensed, so without one this
// fails on a product gate rather than on a regression.
test.skip(!proxyIsPremium(), "proxy under test is unlicensed — Regenerate Key is premium-gated");
await navigateToPage(page, Page.ApiKeys);
await dismissFeedbackPopup(page);
@ -143,12 +159,16 @@ test.describe("Proxy Admin - Keys", () => {
});
test("Delete key", async ({ page }) => {
// Deleting the seeded key leaves nothing for the next attempt, so the retries CI runs with are
// guaranteed to fail and the suite cannot run twice against one database. Bring our own.
const alias = await createDeletableKey(page);
await navigateToPage(page, Page.ApiKeys);
await dismissFeedbackPopup(page);
const keyRow = page.getByRole("row").filter({ hasText: E2E_DELETE_KEY_ALIAS });
const keyRow = page.getByRole("row").filter({ hasText: alias });
await expect(keyRow).toBeVisible({ timeout: 10_000 });
await keyRow.getByRole("button", { name: E2E_DELETE_KEY_ALIAS }).click();
await keyRow.getByRole("button", { name: alias }).click();
await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 });
@ -157,7 +177,7 @@ test.describe("Proxy Admin - Keys", () => {
const modal = page.getByRole("dialog", { name: "Delete Key" });
await expect(modal).toBeVisible({ timeout: 5_000 });
await modal.locator("input").fill(E2E_DELETE_KEY_ALIAS);
await modal.locator("input").fill(alias);
const deleteButton = modal.getByRole("button", { name: "Delete", exact: true });
await expect(deleteButton).toBeEnabled();
@ -167,8 +187,8 @@ test.describe("Proxy Admin - Keys", () => {
// The key is gone when the management API stops returning it, not when the toast says so.
await expect
.poll(async () => await findKeyByAlias(page, E2E_DELETE_KEY_ALIAS), {
message: `key ${E2E_DELETE_KEY_ALIAS} still readable from /key/list after delete`,
.poll(async () => await findKeyByAlias(page, alias), {
message: `key ${alias} still readable from /key/list after delete`,
timeout: 15_000,
})
.toBeUndefined();

162
tests/e2e/ui/tests/proxy-admin/teamSettings.spec.ts Normal file
View file

@ -0,0 +1,162 @@
import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
import { readBack } from "../../helpers/roundTrip";
import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic";
interface TeamInfo {
team_id: string;
team_alias: string;
models: string[];
max_budget: number | null;
tpm_limit: number | null;
rpm_limit: number | null;
metadata: Record<string, unknown> | null;
members_with_roles: { user_id?: string; role?: string }[];
}
/**
* Each test owns a team it created, rather than editing a seeded one, so a save that clobbers a
* field cannot take another spec's fixture down with it.
*/
async function createTeam(page: PlaywrightPage, alias: string, members: string[] = []): Promise<string> {
const res = await page.request.post("/team/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: {
team_alias: alias,
models: [CHAT_MODEL_A],
members_with_roles: members.map((user_id) => ({ user_id, role: "user" })),
},
});
expect(res.ok(), `POST /team/new failed (${res.status()}): ${await res.text()}`).toBe(true);
return (await res.json()).team_id as string;
}
/**
* A member of this test's own, not one of the seeded users. Putting a seeded user on an extra team
* changes what every spec that asserts on their memberships sees.
*/
async function createMember(page: PlaywrightPage, userId: string): Promise<string> {
const res = await page.request.post("/user/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { user_id: userId, user_role: "internal_user", auto_create_key: false },
});
expect(res.ok(), `POST /user/new failed (${res.status()}): ${await res.text()}`).toBe(true);
return userId;
}
async function teamInfo(page: PlaywrightPage, teamId: string): Promise<TeamInfo> {
const body = await readBack<{ team_info: TeamInfo }>(page, `/team/info?team_id=${encodeURIComponent(teamId)}`);
return body.team_info;
}
async function openTeamSettings(page: PlaywrightPage, teamId: string): Promise<void> {
await navigateToPage(page, Page.Teams);
await dismissFeedbackPopup(page);
await clickTeamId(page, teamId);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await expect(page.getByRole("button", { name: "Save Changes" })).toBeVisible({ timeout: 10_000 });
}
test.describe("Proxy Admin - Team settings", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("Setting a team's spend cap and rate limits leaves its models and members alone", async ({ page }) => {
const stamp = Date.now();
const alias = `e2e-team-limits-${stamp}`;
const member = await createMember(page, `e2e-team-limits-member-${stamp}`);
const teamId = await createTeam(page, alias, [member]);
const before = await teamInfo(page, teamId);
await openTeamSettings(page, teamId);
await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("42.5");
await page.getByRole("spinbutton", { name: "Tokens per minute Limit (TPM)" }).fill("7000");
await page.getByRole("spinbutton", { name: "Requests per minute Limit (RPM)" }).fill("70");
await page.getByRole("button", { name: "Save Changes" }).click();
await expect
.poll(
async () => {
const team = await teamInfo(page, teamId);
return [team.max_budget, team.tpm_limit, team.rpm_limit];
},
{ message: "team limits did not persist", timeout: 20_000 },
)
.toEqual([42.5, 7000, 70]);
// The Settings form posts the whole team. A field it fails to seed goes back as null, and
// the toast still says success, so pin the fields this edit had no business touching.
const after = await teamInfo(page, teamId);
expect(after.models, "model access untouched by a limits edit").toEqual(before.models);
expect(
after.members_with_roles.map((member) => member.user_id).sort(),
"membership untouched by a limits edit",
).toEqual(before.members_with_roles.map((member) => member.user_id).sort());
});
test("A model alias added on the Settings tab serves traffic under the alias name", async ({ page }) => {
const stamp = Date.now();
const alias = `e2e-team-alias-${stamp}`;
const modelAlias = `e2e-alias-${stamp}`;
const teamId = await createTeam(page, alias);
await openTeamSettings(page, teamId);
await page.getByRole("textbox", { name: "Alias Name" }).fill(modelAlias);
await page.getByRole("combobox", { name: "Select target model" }).click();
await page.getByRole("option", { name: CHAT_MODEL_A, exact: true }).first().click();
await page.getByRole("button", { name: "Add Alias" }).click();
await page.getByRole("button", { name: "Save Changes" }).click();
await expect
.poll(async () => (await teamInfo(page, teamId)).models, { message: "team lost its models", timeout: 20_000 })
.toEqual([CHAT_MODEL_A]);
const keyRes = await page.request.post("/key/generate", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { team_id: teamId, key_alias: `e2e-alias-key-${stamp}` },
});
expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()})`).toBe(true);
const teamKey = (await keyRes.json()).key as string;
// An alias the team can see but cannot call is the actual complaint; the readback alone
// would pass for an alias the router never resolves.
const served = await page.request.post("/v1/chat/completions", {
headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" },
data: { model: modelAlias, messages: [{ role: "user", content: "ping" }] },
});
expect(served.status(), `a team key calling ${modelAlias} is served`).toBe(200);
expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
});
test("Team metadata added as key-value pairs survives a reload", async ({ page }) => {
const stamp = Date.now();
const alias = `e2e-team-metadata-${stamp}`;
const metadataValue = `cost-center-${stamp}`;
const teamId = await createTeam(page, alias);
await openTeamSettings(page, teamId);
await page.getByRole("button", { name: "Add Key-Value Pair" }).click();
await page.getByPlaceholder("Key", { exact: true }).last().fill("owner");
await page.getByPlaceholder("Value", { exact: true }).last().fill(metadataValue);
await page.getByRole("button", { name: "Save Changes" }).click();
await expect
.poll(async () => (await teamInfo(page, teamId)).metadata?.owner, {
message: "team metadata did not persist",
timeout: 20_000,
})
.toBe(metadataValue);
// Reopening the form is the step that catches metadata the page writes but cannot read back.
await page.reload();
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await expect(page.getByPlaceholder("Key", { exact: true })).toHaveValue("owner", { timeout: 15_000 });
await expect(page.getByPlaceholder("Value", { exact: true })).toHaveValue(metadataValue);
});
});

32
tests/e2e/ui/tests/proxy-admin/teams.spec.ts
View file

@ -1,14 +1,9 @@
import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
import {
ADMIN_STORAGE_PATH,
E2E_TEAM_CRUD_ID,
E2E_TEAM_DELETE_ALIAS,
E2E_TEAM_NO_ADMIN_ID,
E2E_TEAM_ORG_ID,
} from "../../constants";
import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID, E2E_TEAM_NO_ADMIN_ID, E2E_TEAM_ORG_ID } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
import { readBack } from "../../helpers/roundTrip";
import { masterKey } from "../../helpers/traffic";
/** GET /team/list returns a bare array of teams, each carrying team_alias/team_id. */
async function findTeamByAlias(page: PlaywrightPage, alias: string): Promise<Record<string, any> | undefined> {
@ -25,6 +20,17 @@ async function teamMemberEmails(page: PlaywrightPage, teamId: string): Promise<s
return (info.team_info.members_with_roles ?? []).map((member) => member.user_email ?? "").filter(Boolean);
}
/** A team this test owns, so deleting it costs the suite nothing on a retry or a second run. */
async function createDeletableTeam(page: PlaywrightPage): Promise<string> {
const alias = `e2e-delete-team-${Date.now()}`;
const res = await page.request.post("/team/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { team_alias: alias, models: ["fake-openai-gpt-4"] },
});
expect(res.ok(), `POST /team/new failed (${res.status()}): ${await res.text()}`).toBe(true);
return alias;
}
test.describe("Proxy Admin - Teams", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
@ -121,10 +127,14 @@ test.describe("Proxy Admin - Teams", () => {
});
test("Delete a team", async ({ page }) => {
// Deleting the seeded team leaves nothing for the next attempt, so the retries CI runs with are
// guaranteed to fail and the suite cannot run twice against one database. Bring our own.
const alias = await createDeletableTeam(page);
await navigateToPage(page, Page.Teams);
await dismissFeedbackPopup(page);
const teamRow = page.locator("tr", { hasText: E2E_TEAM_DELETE_ALIAS }).first();
const teamRow = page.locator("tr", { hasText: alias }).first();
await expect(teamRow).toBeVisible({ timeout: 10_000 });
// Actions live in a kebab menu: open it, then click "Delete team".
await teamRow.locator('[data-testid^="team-actions-"]').click();
@ -132,15 +142,15 @@ test.describe("Proxy Admin - Teams", () => {
const modal = page.getByRole("dialog", { name: "Delete Team?" });
await expect(modal).toBeVisible({ timeout: 5_000 });
await modal.locator("input").fill(E2E_TEAM_DELETE_ALIAS);
await modal.locator("input").fill(alias);
await modal.getByRole("button", { name: /Force Delete|Delete/i }).click();
await expect(teamRow).not.toBeVisible({ timeout: 10_000 });
// A row vanishing is local state, which happens whether or not the delete landed.
await expect
.poll(async () => await findTeamByAlias(page, E2E_TEAM_DELETE_ALIAS), {
message: `team ${E2E_TEAM_DELETE_ALIAS} still readable from /team/list after delete`,
.poll(async () => await findTeamByAlias(page, alias), {
message: `team ${alias} still readable from /team/list after delete`,
timeout: 15_000,
})
.toBeUndefined();

51
tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts
View file

@ -35,7 +35,44 @@ async function findKeyByAlias(page: PlaywrightPage, alias: string): Promise<Reco
return body.keys.find((row) => row.key_alias === alias);
}
/** A member this test adds itself, so removing it costs the suite nothing on a retry or a re-run. */
async function addRemovableMember(page: PlaywrightPage, registerForCleanup: string[]): Promise<string> {
const userId = `e2e-removable-${Date.now()}`;
// Claimed before the call: /user/new can persist the user and still answer non-2xx, and the id is
// ours either way, so registering it up front is what no failure path can skip.
registerForCleanup.push(userId);
const created = await page.request.post("/user/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { user_id: userId, user_role: "internal_user", auto_create_key: false },
});
expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true);
const added = await page.request.post("/team/member_add", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { team_id: E2E_TEAM_CRUD_ID, member: { user_id: userId, role: "user" } },
});
expect(added.ok(), `POST /team/member_add failed (${added.status()}): ${await added.text()}`).toBe(true);
return userId;
}
test.describe("Team Admin", () => {
const createdMembers: string[] = [];
test.afterEach(async ({ page }) => {
// Runs on the failure path too, which a call at the end of the test body would not. Ids are
// claimed before the user is created, so the delete is attempted unconditionally and only its
// own 404 counts as never persisted; any other answer is a cleanup failure worth reporting
// rather than a reason to leave the user behind.
for (const userId of createdMembers.splice(0)) {
const deleted = await page.request.post("/user/delete", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { user_ids: [userId] },
});
const settled = deleted.ok() || deleted.status() === 404;
expect(settled, `POST /user/delete for ${userId} (${deleted.status()}): ${await deleted.text()}`).toBe(true);
}
});
test.use({ storageState: TEAM_ADMIN_STORAGE_PATH });
test("Team admin can see all team keys including internal user keys", async ({ page }) => {
@ -95,6 +132,10 @@ test.describe("Team Admin", () => {
});
test("Team admin can remove a member from their team", async ({ page }) => {
// Removing the seeded member leaves nothing for the next attempt, so the retries CI runs with
// are guaranteed to fail and the suite cannot run twice against one database. Bring our own.
const memberId = await addRemovableMember(page, createdMembers);
await navigateToPage(page, Page.Teams);
await dismissFeedbackPopup(page);
@ -102,9 +143,9 @@ test.describe("Team Admin", () => {
await page.getByRole("tab", { name: "Members" }).click();
// Seeded members appear in the roster by user_id (members_with_roles has no
// email), so match the row on the user_id rather than the email.
const row = page.locator("tr", { hasText: "e2e-removable-member" }).first();
// Members appear in the roster by user_id (members_with_roles has no email), so match
// the row on the user_id rather than the email.
const row = page.locator("tr", { hasText: memberId }).first();
await expect(row).toBeVisible({ timeout: 10_000 });
await row.getByTestId("delete-member").click();
@ -117,7 +158,7 @@ test.describe("Team Admin", () => {
// Removing the wrong member is exactly what a success toast hides, so pin both halves.
expect(remove.team_id, "delete targets the team being viewed").toBe(E2E_TEAM_CRUD_ID);
expect([remove.user_id, remove.user_email], "delete identifies the member whose row was clicked").toContain(
"e2e-removable-member",
memberId,
);
await expect(page.getByText("Team member removed successfully").first()).toBeVisible({ timeout: 10_000 });
@ -128,7 +169,7 @@ test.describe("Team Admin", () => {
message: "removed member is still on the team",
timeout: 15_000,
})
.not.toContain("e2e-removable-member");
.not.toContain(memberId);
});
test("Team admin sees all team models in the Playground model dropdown", async ({ page, request }) => {

135
tests/e2e/ui/tests/usage/usageActivityTabs.spec.ts Normal file
View file

@ -0,0 +1,135 @@
import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
import {
CHAT_MODEL_A,
CHAT_MODEL_B,
DEPLOYMENT_MODEL_A,
DEPLOYMENT_MODEL_B,
createVirtualKey,
masterKey,
rootPath,
sendChatCompletion,
waitForKeyInDailyActivity,
waitForSpendLog,
} from "../../helpers/traffic";
/**
* Covers the per-entity breakdowns on /ui/usage. The page-level totals move with every other spec's
* traffic, so each assertion is scoped to a key this test minted and to the requests it sent.
*/
/** Each breakdown renders one expandable card per entity, named "<entity> $x.xx N requests". */
const entityCard = (page: PlaywrightPage, tab: string, name: string): Locator =>
page.getByRole("tabpanel", { name: tab }).getByRole("button", { name: new RegExp(`^${name}\\s`) });
async function openUsageTab(page: PlaywrightPage, tab: string): Promise<Locator> {
await navigateToPage(page, Page.NewUsage);
await dismissFeedbackPopup(page);
await page.getByRole("tab", { name: tab }).click();
const panel = page.getByRole("tabpanel", { name: tab });
await expect(panel).toBeVisible({ timeout: 30_000 });
return panel;
}
/** Sends `count` completions on one model and waits for each to reach the spend log. */
async function sendTraffic(
request: Parameters<typeof sendChatCompletion>[0],
apiKey: string,
model: string,
count: number,
label: string,
): Promise<void> {
for (let i = 0; i < count; i++) {
const requestId = await sendChatCompletion(request, { model, prompt: `${label} ${i}`, apiKey });
await waitForSpendLog(request, requestId);
}
}
test.describe("Usage page activity tabs", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("Key Activity breaks a key's traffic down by model", async ({ page, request }) => {
const alias = `e2e-usage-keyact-${Date.now()}`;
const { key, token } = await createVirtualKey(request, { key_alias: alias });
// An uneven split, so a breakdown that lumps everything into one row or attributes to the
// wrong model cannot land on these numbers by accident.
await sendTraffic(request, key, CHAT_MODEL_A, 2, alias);
await sendTraffic(request, key, CHAT_MODEL_B, 1, alias);
await waitForKeyInDailyActivity(request, token, 3);
await openUsageTab(page, "Key Activity");
const card = entityCard(page, "Key Activity", alias);
await expect(card, `${alias} missing from Key Activity`).toBeVisible({ timeout: 30_000 });
await expect(card).toContainText("3 requests");
// Every key gets a card, and the page opens the first one. Scope to this key's own section,
// which the collapsible renders as the trigger's next sibling.
await card.click();
const details = card.locator("xpath=following-sibling::*[1]");
const successfulFor = (model: string) =>
details.getByRole("row").filter({ hasText: model }).getByRole("cell").nth(2); // Model | Spend | Successful | Failed | Tokens
await expect(successfulFor(DEPLOYMENT_MODEL_A)).toHaveText("2", { timeout: 20_000 });
await expect(successfulFor(DEPLOYMENT_MODEL_B)).toHaveText("1");
});
test("Model Activity can name its models by deployment instead of by public name", async ({ page, request }) => {
const alias = `e2e-usage-modelact-${Date.now()}`;
const { key, token } = await createVirtualKey(request, { key_alias: alias });
await sendTraffic(request, key, CHAT_MODEL_A, 1, alias);
await waitForKeyInDailyActivity(request, token);
const panel = await openUsageTab(page, "Model Activity");
await expect(entityCard(page, "Model Activity", CHAT_MODEL_A), `${CHAT_MODEL_A} missing`).toBeVisible({
timeout: 30_000,
});
// Nothing is published under the deployment's name, so its absence here is what makes the
// toggle below a real change of key rather than a relabelled button.
await expect(entityCard(page, "Model Activity", DEPLOYMENT_MODEL_A)).toHaveCount(0);
// Admins reconcile provider bills against the deployment, not the name their users call.
await panel.getByRole("button", { name: "Litellm Model Name" }).click();
await expect(entityCard(page, "Model Activity", DEPLOYMENT_MODEL_A)).toBeVisible({ timeout: 20_000 });
});
test("Filter by user narrows Key Activity to that user's keys", async ({ page, request }) => {
const stamp = Date.now();
const email = `e2e-usage-owner-${stamp}@test.local`;
const ownedAlias = `e2e-usage-owned-${stamp}`;
const otherAlias = `e2e-usage-other-${stamp}`;
const userRes = await request.post(`${rootPath()}/user/new`, {
headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" },
data: { user_email: email, user_role: "internal_user", auto_create_key: false },
});
expect(userRes.ok(), `POST /user/new failed (${userRes.status()})`).toBe(true);
const userId = (await userRes.json()).user_id as string;
const owned = await createVirtualKey(request, { key_alias: ownedAlias, user_id: userId });
const other = await createVirtualKey(request, { key_alias: otherAlias });
await sendTraffic(request, owned.key, CHAT_MODEL_A, 1, ownedAlias);
await sendTraffic(request, other.key, CHAT_MODEL_A, 1, otherAlias);
await waitForKeyInDailyActivity(request, owned.token);
await waitForKeyInDailyActivity(request, other.token);
await openUsageTab(page, "Key Activity");
await expect(entityCard(page, "Key Activity", otherAlias)).toBeVisible({ timeout: 30_000 });
await page.getByRole("combobox", { name: "Search users by email" }).click();
await page.keyboard.type(email);
await page
.getByRole("option", { name: new RegExp(email) })
.first()
.click();
// The filter earns its place only by dropping the other key; the owned key showing up
// proves nothing on a page that already listed every key.
await expect(entityCard(page, "Key Activity", otherAlias)).toHaveCount(0, { timeout: 30_000 });
await expect(entityCard(page, "Key Activity", ownedAlias)).toBeVisible({ timeout: 20_000 });
});
});

108
tests/e2e/ui/tests/usage/usagePage.spec.ts
View file

@ -1,10 +1,11 @@
import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
import {
CHAT_MODEL_A,
createVirtualKey,
masterKey,
rootPath,
sendChatCompletion,
waitForKeyInDailyActivity,
waitForSpendLog,
@ -27,9 +28,105 @@ async function openUsage(page: PlaywrightPage): Promise<Locator> {
return card;
}
/** The upstream fixtures/config.yml points its models at, so the mock server answers this too. */
const MOCK_DEPLOYMENT = "openai/fake-gpt-4";
/** A deployment whose traffic costs real money, so the key that used it outranks the $0 crowd. */
async function createPricedDeployment(
request: APIRequestContext,
label: string,
registerForCleanup: string[],
): Promise<{ modelName: string }> {
const modelName = `e2e-usage-priced-${label}`;
// Claimed before the call: /model/new can persist the deployment and still answer non-2xx, so a
// name recorded up front is the only registration no response shape can skip.
registerForCleanup.push(modelName);
const res = await request.post("/model/new", {
headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" },
data: {
model_name: modelName,
litellm_params: {
model: MOCK_DEPLOYMENT,
api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
api_key: "fake-key",
input_cost_per_token: 0.01,
output_cost_per_token: 0.01,
},
},
});
expect(res.ok(), `POST /model/new failed (${res.status()}): ${await res.text()}`).toBe(true);
// /model/new returns once the row is written, but the router only picks the deployment up on its
// next refresh, so sending traffic straight away can still get "no healthy deployments". A ping
// that fails writes no spend log, so retrying it costs the ranking this test asserts nothing.
await expect
.poll(
async () => {
const ping = await request.post(`${rootPath()}/v1/chat/completions`, {
headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" },
data: { model: modelName, messages: [{ role: "user", content: "readiness ping" }] },
});
return ping.ok();
},
{ message: `deployment ${modelName} never became routable`, timeout: 60_000 },
)
.toBe(true);
return { modelName };
}
test.describe("Usage page", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
const pricedDeployments: string[] = [];
test.afterEach(async ({ request }) => {
// A deployment left behind keeps its custom pricing, so it goes on changing what later runs
// route and what they cost. Runs on the failure path too, which the test body would not.
// Resolved by name rather than by a returned id, so a create that persisted without answering
// 2xx is still cleaned up. /model/info serves the router, and /model/new answers 2xx even when
// its in-request router reload failed, so the search-backed listing is what covers a deployment
// that reached the database only. Absent from both means it never persisted.
const names = pricedDeployments.splice(0);
if (names.length === 0) return;
const auth = { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" };
type Lookup =
| { readonly listed: true; readonly id: string | undefined }
| { readonly listed: false; readonly status: number };
const idIn = async (path: string, name: string): Promise<Lookup> => {
const listed = await request.get(path, { headers: auth });
if (!listed.ok()) return { listed: false, status: listed.status() };
const deployments = ((await listed.json()).data ?? []) as {
model_name?: string;
model_info?: { id?: string };
}[];
return { listed: true, id: deployments.find((d) => d.model_name === name)?.model_info?.id };
};
const remove = async (name: string, id: string) => {
const deleted = await request.post(`${rootPath()}/model/delete`, { headers: auth, data: { id } });
expect(deleted.ok(), `POST /model/delete for ${name} (${deleted.status()})`).toBe(true);
};
for (const name of names) {
const fromRouter = await idIn(`${rootPath()}/model/info`, name);
if (fromRouter.listed && fromRouter.id !== undefined) {
await remove(name, fromRouter.id);
continue;
}
const search = encodeURIComponent(name);
const fromDb = await idIn(`${rootPath()}/v2/model/info?search=${search}`, name);
expect(
fromDb.listed,
`GET /v2/model/info?search=${search} (${fromDb.listed ? 200 : fromDb.status}), so ${name} could not be checked`,
).toBe(true);
if (!fromDb.listed || fromDb.id === undefined) continue;
await remove(name, fromDb.id);
}
});
test("Top Virtual Keys lists a key that served traffic, toggles views, and opens key info", async ({
page,
request,
@ -39,8 +136,13 @@ test.describe("Usage page", () => {
key_alias: alias,
});
// Top Virtual Keys ranks by spend, and every mock deployment costs $0, so once a run has more
// keys than the list shows, whether this one makes the cut is down to how ties happen to sort.
// Give it a priced deployment of its own so it earns its place.
const { modelName } = await createPricedDeployment(request, alias, pricedDeployments);
const requestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
model: modelName,
prompt: `usage ping for ${alias}`,
apiKey: key,
});

26
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
View file

@ -40,6 +40,24 @@ def prometheus_logger() -> PrometheusLogger:
return PrometheusLogger()
@pytest.fixture
def known_model_router():
router = litellm.Router(
model_list=[
{
"model_name": "gpt-5-mini",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"},
},
{
"model_name": "us/azure/openai/gpt-5-mini",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"},
},
]
)
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
yield router
def create_standard_logging_payload() -> StandardLoggingPayload:
return StandardLoggingPayload(
id="test_id",
@ -741,7 +759,7 @@ async def test_async_log_failure_event(prometheus_logger):
@pytest.mark.asyncio
async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger):
async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger, known_model_router):
"""LiteLLM-side reject (no deployment picked) routes the requested model
into `requested_model` and skips the partial-outage flag."""
standard_logging_object = create_standard_logging_payload()
@ -786,7 +804,7 @@ async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger
@pytest.mark.asyncio
async def test_async_post_call_failure_hook(prometheus_logger):
async def test_async_post_call_failure_hook(prometheus_logger, known_model_router):
"""
Test for the async_post_call_failure_hook method
@ -1069,7 +1087,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
@pytest.mark.asyncio
async def test_log_success_fallback_event(prometheus_logger):
async def test_log_success_fallback_event(prometheus_logger, known_model_router):
prometheus_logger.litellm_deployment_successful_fallbacks = MagicMock()
original_model_group = "gpt-5-mini"
@ -1107,7 +1125,7 @@ async def test_log_success_fallback_event(prometheus_logger):
@pytest.mark.asyncio
async def test_log_failure_fallback_event(prometheus_logger):
async def test_log_failure_fallback_event(prometheus_logger, known_model_router):
prometheus_logger.litellm_deployment_failed_fallbacks = MagicMock()
original_model_group = "gpt-5-mini"

571
tests/litellm-proxy-extras/test_setup_database_fail_fast.py
View file

@ -1,571 +0,0 @@
"""Regression tests for ProxyExtrasDBManager's v2 migration resolver.
v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver`
kwarg, which still defaults to False for direct callers.
"""
import subprocess
from unittest.mock import patch
import pytest
from litellm_proxy_extras.utils import (
_PRISMA_ATTEMPTS,
ProxyExtrasDBManager,
_max_migration_timestamp,
_migration_timestamp,
)
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
def _run(*args, **kwargs):
raise subprocess.CalledProcessError(
returncode=returncode,
cmd=args[0],
stderr=stderr,
output="",
)
return _run
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a permission failure during migrate deploy raises RuntimeError."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3018\nMigration name: 20250326162113_baseline\n"
"Database error code: 42501\npermission denied for schema public"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="permission"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
'Reason: syntax error at or near "BRKN" LINE 42'
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_strip_prisma_query_params_removes_connection_limit():
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
assert "connection_limit" not in stripped
assert "pool_timeout" not in stripped
assert "sslmode=require" in stripped
def test_strip_prisma_query_params_passthrough_no_query():
"""URLs without query strings are returned unchanged."""
url = "postgresql://u:p@h:5432/db"
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
def test_migration_timestamp_extracts_leading_digits():
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
def test_migration_timestamp_returns_zero_on_malformed():
assert _migration_timestamp("0_init") == 0
assert _migration_timestamp("not_a_migration") == 0
def test_max_migration_timestamp():
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
assert _max_migration_timestamp(names) == 20260415000000
def test_max_migration_timestamp_empty_set():
assert _max_migration_timestamp(set()) == 0
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
This is the existing buggy behavior — we're not fixing it in v1, only
offering v2 as opt-in. This test pins the default so that a future
inadvertent default flip is caught.
"""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
# Stub `prisma migrate deploy` to claim success with pending migrations
# applied, which is the code path that triggers the legacy post-migration
# sanity check (a call to _resolve_all_migrations).
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
def fake_run(cmd, *args, **kwargs):
return FakeResult()
resolve_called = {"n": 0}
def fake_resolve(*args, **kwargs):
resolve_called["n"] += 1
monkeypatch.setattr("subprocess.run", fake_run)
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
assert ok is True
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = "db push error"
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="prisma db push failed"):
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
"""_warn_if_db_ahead_of_head must never raise — it's informational.
Non-connection DB errors (e.g. InsufficientPrivilege from a user
without SELECT on _prisma_migrations) must be caught, not propagated.
"""
import psycopg
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class _FakeConn:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def execute(self, *a, **kw):
# 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)
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(
monkeypatch, tmp_path
):
"""If marking a migration as applied fails inside P3009 idempotent
recovery, the subprocess error must be re-raised as RuntimeError so
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
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(
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
)
# First call: migrate deploy -> P3009 idempotent error.
# Recovery path tries _resolve_specific_migration; that also raises.
def _failing_resolve(*a, **kw):
raise subprocess.CalledProcessError(
returncode=1,
cmd="prisma migrate resolve --applied",
stderr="resolve failed",
output="",
)
monkeypatch.setattr(
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
)
stderr = (
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
"relation already exists"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(
RuntimeError, match=r"Failed to mark migration .* as applied"
):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
resolve_called = {"n": 0}
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_all_migrations",
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
)
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`.
Scoped by argv so the Prisma toolchain check cannot consume the mock first.
"""
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):
"""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):
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):
"""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):
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):
"""v2: a database not accepting connections yet is retried, not fatal."""
_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):
"""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):
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):
"""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`"
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_db_push_retries_transient_failures(monkeypatch, tmp_path):
"""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}
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_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 _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)
def _side_effect(n, cmd):
raise subprocess.CalledProcessError(
returncode=1,
cmd=cmd,
stderr="Error: relation \"LiteLLM_SpendLogs\" does not exist",
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

6
tests/llm_translation/conftest.py
View file

@ -44,7 +44,11 @@ _VCR_AUTO_MARKER_SKIP_FILES = frozenset(
{"test_vcr_redis_persister.py", "test_ws_vcr.py"}
)
_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ()
_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = (
"test_nvidia_nim.py::test_embedding_nvidia_nim",
"test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[False]",
"test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[True]",
)
_verbose_state = VerboseReporterState()

14
tests/llm_translation/test_prompt_factory.py
View file

@ -1446,9 +1446,17 @@ def test_convert_to_anthropic_tool_invoke_sanitizes_invalid_ids():
def test_convert_to_anthropic_tool_invoke_server_tool():
"""
Test that server_tool_use (srvtoolu_) is reconstructed as server_tool_use.
Test that a server tool call (srvtoolu_) with no stored result is replayed
as a regular tool_use block.
Fixes: https://github.com/BerriAI/litellm/issues/17737
A server_tool_use block is only valid when paired with its result block, so
an unpaired one must degrade to tool_use for Anthropic to accept the replay.
A paired call still becomes server_tool_use, covered by
test_convert_to_anthropic_tool_invoke_with_web_search_results.
Context: https://github.com/BerriAI/litellm/issues/17737 (original
server_tool_use reconstruction) and LIT-6622 / PR #39144 (unpaired calls
degrade instead of 400ing at Anthropic).
"""
tool_calls = [
{
@ -1464,7 +1472,7 @@ def test_convert_to_anthropic_tool_invoke_server_tool():
result = convert_to_anthropic_tool_invoke(tool_calls)
assert len(result) == 1
assert result[0]["type"] == "server_tool_use" # NOT tool_use
assert result[0]["type"] == "tool_use"
assert result[0]["id"] == "srvtoolu_01ABC123"
assert result[0]["name"] == "web_search"
assert result[0]["input"] == {"query": "elephant weight"}

1
tests/local_testing/conftest.py
View file

@ -90,6 +90,7 @@ _VCR_INCOMPATIBLE_FILES = frozenset(
# carry no real provider cost.
_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = (
"test_router.py::test_router_text_completion_client",
"test_embedding.py::test_encoding_format_omitted_by_default_for_openai_sdk",
)

14
tests/local_testing/test_basic_python_version.py
View file

@ -305,16 +305,14 @@ def _run_proxy_server_smoke_test(extra_proxy_args=None):
def test_litellm_proxy_server_config_no_general_settings():
"""Exercises the default (v2) migration resolver."""
"""Exercises the default (v1) migration resolver."""
_run_proxy_server_smoke_test()
def test_litellm_proxy_server_config_no_general_settings_legacy_resolver():
"""Exercises the legacy (v1) migration resolver against a real database.
def test_litellm_proxy_server_config_no_general_settings_v2_resolver():
"""Exercises the opt-in v2 migration resolver.
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.
Runs in a separate CI job against a local Postgres to avoid collisions
with the v1 variant when they share a database.
"""
_run_proxy_server_smoke_test(extra_proxy_args=["--use_legacy_migration_resolver"])
_run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"])

5
tests/local_testing/test_get_llm_provider.py
View file

@ -155,6 +155,11 @@ def test_default_api_base():
continue
elif provider == "github" and other_provider.value == "azure":
continue
elif (
provider in ("qwencloud", "qwen_ai_platform")
and other_provider.value == "dashscope"
):
continue
assert other_provider.value not in api_base.replace("/openai", "")

78
tests/mcp_tests/test_mcp_client_unit.py
View file

@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY
import litellm.experimental_mcp_client.client as mcp_client_module
from litellm.experimental_mcp_client.client import MCPClient
from litellm.types.mcp import MCPAuth, MCPTransport
from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import ListToolsResult, PaginatedRequestParams
from mcp.types import Tool as MCPTool
def test_mcp_client_uses_configurable_default_timeout():
@ -185,6 +187,80 @@ class TestMCPClientUnitTests:
mock_session_instance.initialize.assert_called_once()
mock_session_instance.list_tools.assert_called_once()
@pytest.mark.asyncio
@patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
@patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
async def test_list_tools_follows_next_cursor_until_exhausted(
self,
mock_session_class,
mock_transport,
):
"""Test listing tools follows MCP pagination cursors until exhausted."""
mock_transport_ctx = AsyncMock()
mock_transport.return_value = mock_transport_ctx
mock_transport_instance = MagicMock()
mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance)
mock_session_ctx = AsyncMock()
mock_session_class.return_value = mock_session_ctx
mock_session_instance = AsyncMock()
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
first_page_tools = [
MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100)
]
second_page_tool = MCPTool(
name="tool_100",
description="Tool 100",
inputSchema={},
)
mock_session_instance.list_tools.side_effect = [
ListToolsResult(tools=first_page_tools, nextCursor="page-2"),
ListToolsResult(tools=[second_page_tool]),
]
client = MCPClient("http://example.com")
result = await client.list_tools()
assert result == [*first_page_tools, second_page_tool]
assert mock_session_instance.list_tools.call_count == 2
second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"]
assert isinstance(second_call_params, PaginatedRequestParams)
assert second_call_params.cursor == "page-2"
@pytest.mark.asyncio
@patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
@patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
async def test_list_tools_swallows_mid_walk_error_without_raise_on_error(
self,
mock_session_class,
mock_transport,
):
"""Test a mid-walk failure returns [] when raise_on_error is False."""
mock_transport_ctx = AsyncMock()
mock_transport.return_value = mock_transport_ctx
mock_transport_instance = MagicMock()
mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance)
mock_session_ctx = AsyncMock()
mock_session_class.return_value = mock_session_ctx
mock_session_instance = AsyncMock()
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
mock_session_instance.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})],
nextCursor="page-2",
),
RuntimeError("transient upstream failure"),
]
client = MCPClient("http://example.com")
result = await client.list_tools()
assert result == []
assert mock_session_instance.list_tools.call_count == 2
@pytest.mark.asyncio
@patch.object(mcp_client_module, "streamable_http_client")
@patch.object(mcp_client_module, "ClientSession")

130
tests/test_litellm/experimental_mcp_client/test_tools.py
View file

@ -8,11 +8,13 @@ from mcp.types import (
CallToolRequestParams,
CallToolResult,
ListToolsResult,
PaginatedRequestParams,
TextContent,
)
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.tools import (
list_tools_with_pagination,
transform_mcp_tool_to_anthropic_tool,
_get_function_arguments,
_normalize_mcp_input_schema,
@ -106,6 +108,134 @@ async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result
mock_session.list_tools.assert_called_once()
@pytest.mark.asyncio()
async def test_load_mcp_tools_follows_pagination(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[
MCPTool(name="tool_a", description="a", inputSchema={}),
MCPTool(name="tool_b", description="b", inputSchema={}),
],
nextCursor="page-2",
),
ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]),
]
result = await load_mcp_tools(mock_session, format="mcp")
assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"]
assert mock_session.list_tools.call_count == 2
second_call_params = mock_session.list_tools.call_args_list[1].kwargs["params"]
assert isinstance(second_call_params, PaginatedRequestParams)
assert second_call_params.cursor == "page-2"
@pytest.mark.asyncio()
async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch):
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2)
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="page-2",
),
ListToolsResult(
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
nextCursor="page-3",
),
ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]),
]
result = await list_tools_with_pagination(mock_session)
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
assert mock_session.list_tools.call_count == 2
@pytest.mark.asyncio()
async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="same-cursor",
),
ListToolsResult(
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
nextCursor="same-cursor",
),
]
result = await list_tools_with_pagination(mock_session)
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
assert mock_session.list_tools.call_count == 2
@pytest.mark.asyncio()
async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="",
),
]
result = await list_tools_with_pagination(mock_session)
assert [tool.name for tool in result] == ["tool_0"]
mock_session.list_tools.assert_called_once()
@pytest.mark.asyncio()
async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch):
import anyio
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2)
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.2)
async def slow_page(params=None):
await anyio.sleep(0.15)
idx = int(params.cursor) if params is not None else 0
return ListToolsResult(
tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})],
nextCursor=str(idx + 1),
)
mock_session.list_tools = slow_page
result = await list_tools_with_pagination(mock_session)
assert [tool.name for tool in result] == ["tool_0"]
@pytest.mark.asyncio()
async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_session, monkeypatch):
import anyio
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1)
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.1)
async def slow_page(params=None):
await anyio.sleep(0.15)
idx = int(params.cursor) if params is not None else 0
tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})]
if idx == 0:
return ListToolsResult(tools=tools, nextCursor="1")
return ListToolsResult(tools=tools)
mock_session.list_tools = slow_page
result = await list_tools_with_pagination(mock_session, listing_deadline=2.0)
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
@pytest.mark.asyncio()
async def test_load_mcp_tools_openai_format_spans_pages(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_a", description="a", inputSchema={})],
nextCursor="page-2",
),
ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]),
]
result = await load_mcp_tools(mock_session, format="openai")
assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"]
def test_get_function_arguments():
# Test with string arguments
function = {"arguments": '{"test": "value"}'}

279
tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py Normal file
View file

@ -0,0 +1,279 @@
"""
LIT-6611: every unique client-supplied model name that fails routing used to
mint permanent Prometheus series carrying ``requested_model="<junk>"`` on the
proxy request metrics and the deployment metrics, with no eviction. The fix
collapses any requested model the router does not recognize (and no wildcard
pattern matches) into the single ``other`` label bucket, while recognized
names, aliases, and wildcard-matched names keep their own label values.
"""
import sys
import types
from unittest.mock import patch
import pytest
from prometheus_client import REGISTRY
import litellm
from litellm.integrations.prometheus import (
UNRECOGNIZED_REQUESTED_MODEL_LABEL,
PrometheusLogger,
)
from litellm.proxy._types import UserAPIKeyAuth
class _ClientSideError(Exception):
status_code = 400
@pytest.fixture(autouse=True)
def cleanup_prometheus_registry():
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
yield
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
@pytest.fixture
def router():
return litellm.Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"},
},
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": "fake-key"},
},
],
model_group_alias={"gpt4o-alias": "gpt-4o-mini"},
)
@pytest.fixture
def team_router():
return litellm.Router(
model_list=[
{
"model_name": "team-internal-gpt",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"},
"model_info": {"team_id": "team-1", "team_public_model_name": "team-alias-gpt"},
},
{
"model_name": "team-internal-bedrock",
"litellm_params": {"model": "openai/*", "api_key": "fake-key"},
"model_info": {"team_id": "team-1", "team_public_model_name": "team-models/*"},
},
]
)
def _requested_model_values(metric) -> set[str]:
index = metric._labelnames.index("requested_model")
return {sample_key[index] for sample_key in metric._metrics}
def _series_count(metric) -> int:
return len(metric._metrics)
def _total_value(metric) -> float:
return sum(child._value.get() for child in metric._metrics.values())
async def _fire_proxy_failure(logger: PrometheusLogger, model: str) -> None:
await logger.async_post_call_failure_hook(
request_data={"model": model, "metadata": {}, "proxy_server_request": {}},
original_exception=_ClientSideError(f"model {model} does not exist"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key-1"),
)
@pytest.mark.asyncio
async def test_unknown_models_collapse_to_one_series_on_proxy_request_metrics(router):
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
for index in range(25):
await _fire_proxy_failure(logger, f"agent-typo-{index}")
for metric in (
logger.litellm_proxy_failed_requests_metric,
logger.litellm_proxy_total_requests_metric,
):
assert _requested_model_values(metric) == {UNRECOGNIZED_REQUESTED_MODEL_LABEL}
assert _series_count(metric) == 1
assert _total_value(metric) == 25
@pytest.mark.asyncio
async def test_known_alias_and_wildcard_models_keep_their_own_labels(router):
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
await _fire_proxy_failure(logger, "gpt-4o-mini")
await _fire_proxy_failure(logger, "gpt4o-alias")
await _fire_proxy_failure(logger, "openai/gpt-4o-audio-preview")
await _fire_proxy_failure(logger, "agent-typo-hallucinated")
for metric in (
logger.litellm_proxy_failed_requests_metric,
logger.litellm_proxy_total_requests_metric,
):
assert _requested_model_values(metric) == {
"gpt-4o-mini",
"gpt4o-alias",
"openai/gpt-4o-audio-preview",
UNRECOGNIZED_REQUESTED_MODEL_LABEL,
}
@pytest.mark.asyncio
async def test_team_alias_and_team_wildcard_models_keep_their_own_labels(team_router):
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", team_router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
await _fire_proxy_failure(logger, "team-alias-gpt")
await _fire_proxy_failure(logger, "team-models/gpt-4o-audio-preview")
await _fire_proxy_failure(logger, "agent-typo-hallucinated")
for metric in (
logger.litellm_proxy_failed_requests_metric,
logger.litellm_proxy_total_requests_metric,
):
assert _requested_model_values(metric) == {
"team-alias-gpt",
"team-models/gpt-4o-audio-preview",
UNRECOGNIZED_REQUESTED_MODEL_LABEL,
}
@pytest.mark.asyncio
async def test_unknown_models_collapse_to_other_when_router_is_unavailable():
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", None, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
await _fire_proxy_failure(logger, "agent-typo-no-router")
await _fire_proxy_failure(logger, "gpt-4o-mini")
assert _requested_model_values(logger.litellm_proxy_failed_requests_metric) == {
UNRECOGNIZED_REQUESTED_MODEL_LABEL
}
@pytest.mark.asyncio
async def test_sdk_router_originated_metrics_keep_labels_without_proxy_router():
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", None, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
logger.set_llm_deployment_failure_metrics(
request_kwargs={
"model": "sdk-deployment-group",
"litellm_params": {"metadata": {}},
"standard_logging_object": {},
"exception": _ClientSideError("model does not exist"),
}
)
await logger.log_failure_fallback_event(
original_model_group="sdk-fallback-group",
kwargs={"model": "sdk-fallback-group", "metadata": {}},
original_exception=_ClientSideError("upstream unavailable"),
)
assert _requested_model_values(logger.litellm_deployment_failure_responses) == {"sdk-deployment-group"}
assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == {"sdk-fallback-group"}
@pytest.mark.asyncio
async def test_sdk_fallback_labels_survive_non_import_errors_from_proxy_module(monkeypatch):
logger = PrometheusLogger()
broken_proxy_module = types.ModuleType("litellm.proxy.proxy_server")
def _raise_value_error(_name: str):
raise ValueError("bad proxy env var")
broken_proxy_module.__getattr__ = _raise_value_error # test-quality-ok: reproduces a proxy_server import raising non-ImportError, no injection seam
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", broken_proxy_module) # test-quality-ok: reproduces a proxy_server import raising non-ImportError, no injection seam
await logger.log_failure_fallback_event(
original_model_group="sdk-fallback-group",
kwargs={"model": "sdk-fallback-group", "metadata": {}},
original_exception=_ClientSideError("upstream unavailable"),
)
assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == {"sdk-fallback-group"}
def test_unknown_models_collapse_to_one_series_on_deployment_metrics(router):
logger = PrometheusLogger()
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
for index in range(25):
logger.set_llm_deployment_failure_metrics(
request_kwargs={
"model": f"agent-typo-{index}",
"litellm_params": {"metadata": {}},
"standard_logging_object": {},
"exception": _ClientSideError("model does not exist"),
}
)
logger.set_llm_deployment_failure_metrics(
request_kwargs={
"model": "gpt-4o-mini",
"litellm_params": {"metadata": {}},
"standard_logging_object": {},
"exception": _ClientSideError("all deployments cooling down"),
}
)
for metric in (
logger.litellm_deployment_failure_responses,
logger.litellm_deployment_total_requests,
):
assert _requested_model_values(metric) == {
UNRECOGNIZED_REQUESTED_MODEL_LABEL,
"gpt-4o-mini",
}
assert _series_count(metric) == 2
assert _total_value(metric) == 26
@pytest.mark.asyncio
async def test_fallback_event_requested_model_is_bounded(router):
logger = PrometheusLogger()
kwargs = {"model": "gpt-4o-mini", "metadata": {}}
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
await logger.log_failure_fallback_event(
original_model_group="agent-typo-hallucinated",
kwargs=kwargs,
original_exception=_ClientSideError("model does not exist"),
)
await logger.log_success_fallback_event(
original_model_group="agent-typo-hallucinated",
kwargs=kwargs,
original_exception=_ClientSideError("model does not exist"),
)
await logger.log_failure_fallback_event(
original_model_group="gpt-4o-mini",
kwargs=kwargs,
original_exception=_ClientSideError("upstream unavailable"),
)
assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == {
UNRECOGNIZED_REQUESTED_MODEL_LABEL,
"gpt-4o-mini",
}
assert _requested_model_values(logger.litellm_deployment_successful_fallbacks) == {
UNRECOGNIZED_REQUESTED_MODEL_LABEL
}

42
tests/test_litellm/integrations/test_s3.py
View file

@ -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

288
tests/test_litellm/integrations/test_s3_v2.py
View file

@ -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 `<prefix>/<date>/<head>_<sha256>.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)
# --------------------------------------------------------------

373
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -32,6 +32,8 @@ 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,
get_token_type_cost_breakdown,
@ -409,6 +411,377 @@ 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_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
# 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_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 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_non_mapping_off_peak_block_bills_standard_rates():
"""A truthy non-mapping off_peak_pricing value (a bare string or a list in
YAML) must bill standard rates rather than raising, matching how every
other malformed piece of the block behaves.
"""
from datetime import datetime, timezone
from typing import cast
from litellm.types.utils import ModelInfo
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
when = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)
for malformed_block in ("16:00-19:00", ["16:00-19:00"], 5e-7, True):
model_info = cast(
ModelInfo,
{
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
"off_peak_pricing": malformed_block,
},
)
result = _get_token_base_cost(model_info, usage, current_time=when)
assert result[0] == 1e-6
assert result[1] == 2e-6
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_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"

74
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
View file

@ -1435,6 +1435,80 @@ class TestFlattenTopLevelSchemaCombinators:
assert schema == snapshot
class TestToolWithFlattenedParameters:
def _anyof_tool(self):
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 test_flattens_anyof_parameters_into_new_tool(self):
from litellm.litellm_core_utils.prompt_templates.common_utils import (
tool_with_flattened_parameters,
)
tool = self._anyof_tool()
result = tool_with_flattened_parameters(tool)
assert result is not tool
parameters = result["function"]["parameters"]
assert "anyOf" not in parameters
assert parameters["type"] == "object"
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
assert parameters["required"] == ["id"]
assert result["function"]["name"] == "automation_update"
assert tool == self._anyof_tool()
def test_clean_parameters_return_the_same_tool_object(self):
from litellm.litellm_core_utils.prompt_templates.common_utils import (
tool_with_flattened_parameters,
)
tool = {
"type": "function",
"function": {
"name": "lookup",
"parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]},
},
}
assert tool_with_flattened_parameters(tool) is tool
@pytest.mark.parametrize(
"tool",
[
{"type": "function"},
{"type": "function", "function": "not-a-dict"},
{"type": "function", "function": {"name": "no_params"}},
{"type": "function", "function": {"name": "bad_params", "parameters": "not-a-dict"}},
],
)
def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool):
from litellm.litellm_core_utils.prompt_templates.common_utils import (
tool_with_flattened_parameters,
)
assert tool_with_flattened_parameters(tool) is tool
class TestRequestContainsImageContent:
"""One detector for every dialect that reaches pre-routing hooks untranslated."""

64
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -3627,3 +3627,67 @@ def test_convert_gemini_tool_call_result_answers_tool_reference_only_result():
)
assert result == {"function_response": {"name": "ToolSearch", "response": {"content": ""}}}
def test_convert_to_anthropic_tool_invoke_degrades_unpaired_server_tool_use():
"""A replayed srvtoolu_ call whose server tool result is not available
(e.g. the Responses bridge replays items without provider_specific_fields)
must become a plain client tool_use so the client's tool_result can pair
with it. A dangling server_tool_use makes Anthropic 400 the request with
"unexpected `tool_use_id` found in `tool_result` blocks"."""
from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_tool_invoke
result = convert_to_anthropic_tool_invoke(
tool_calls=[
{
"id": "srvtoolu_01Unpaired",
"type": "function",
"function": {"name": "web_search", "arguments": '{"query": "zig version"}'},
}
],
web_search_results=None,
tool_results=None,
)
assert result == [
{
"type": "tool_use",
"id": "srvtoolu_01Unpaired",
"name": "web_search",
"input": {"query": "zig version"},
}
]
def test_convert_to_anthropic_tool_invoke_keeps_paired_server_tool_use():
"""When the paired server tool result is available, the srvtoolu_ call is
still reconstructed as server_tool_use followed by its result block."""
from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_tool_invoke
server_result = {
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01Paired",
"content": [{"type": "web_search_result", "url": "https://ziglang.org", "title": "Zig"}],
}
result = convert_to_anthropic_tool_invoke(
tool_calls=[
{
"id": "srvtoolu_01Paired",
"type": "function",
"function": {"name": "web_search", "arguments": '{"query": "zig version"}'},
}
],
web_search_results=[server_result],
tool_results=None,
)
assert result == [
{
"type": "server_tool_use",
"id": "srvtoolu_01Paired",
"name": "web_search",
"input": {"query": "zig version"},
},
server_result,
]

2
tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py
View file

@ -188,6 +188,8 @@ class TestDeclaredAuthenticatingProvider:
("gpt-4o", "github_copilot", "github_copilot"),
("openai/gpt-4o", None, None),
("gpt-4o", "openai", None),
("github_copilot", None, None),
("chatgpt", None, None),
],
)
def test_names_only_the_providers_whose_resolution_authenticates(self, model, provider, expected):

30
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -6354,3 +6354,33 @@ def test_anthropic_drop_params_reduces_mixed_output_config_to_format(monkeypatch
)
assert result.get("output_config") == {"format": schema_format}
def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(local_model_cost_map, monkeypatch):
"""Backstop: on the tool-based structured-output path, a model flagged
``supports_forced_tool_use: false`` must not get the forced response-format
tool_choice the provider would 400 on."""
monkeypatch.setitem(
litellm.model_cost,
"claude-test-no-forced-tools",
{"litellm_provider": "anthropic", "mode": "chat", "supports_forced_tool_use": False},
)
config = AnthropicConfig()
result = config.map_openai_params(
non_default_params={
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "test_schema",
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
},
}
},
optional_params={},
model="claude-test-no-forced-tools",
drop_params=False,
)
assert "tools" in result
assert "tool_choice" not in result

89
tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py
View file

@ -11,6 +11,7 @@ sys.path.insert(
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY
from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
from litellm.utils import get_optional_params
@ -195,3 +196,91 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None:
assert "presence_penalty" not in mapped
assert "logit_bias" not in mapped
assert "reasoning_effort" in supported
class TestAzureToolSchemaCombinatorFlattening:
"""
Regression tests for LIT-6510: Azure's chat completions validator rejects
tool parameters carrying a top-level anyOf/oneOf/allOf for every model
family, so AzureOpenAIConfig.transform_request must flatten them.
"""
@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, tools):
return config.transform_request(
model=model,
messages=[{"role": "user", "content": "hi"}],
optional_params={"tools": tools},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
def test_transform_request_flattens_top_level_anyof(self):
request = self._transform(AzureOpenAIConfig(), "gpt-4o", [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_config_flattens_via_shared_transform(self):
request = self._transform(AzureOpenAIGPT5Config(), "gpt-5.4-mini", [self._anyof_tool()])
parameters = request["tools"][0]["function"]["parameters"]
assert "anyOf" not in parameters
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
def test_caller_tool_dict_is_not_mutated(self):
tool = self._anyof_tool()
self._transform(AzureOpenAIConfig(), "gpt-4o", [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(AzureOpenAIConfig(), "gpt-4o", [tool])
assert request["tools"][0] is tool
def test_non_dict_tool_entries_pass_through_unchanged(self):
request = self._transform(AzureOpenAIConfig(), "gpt-4o", ["not-a-tool"])
assert request["tools"] == ["not-a-tool"]
def test_request_without_tools_is_unchanged(self):
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"temperature": 0.2},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert "tools" not in request
assert request["temperature"] == 0.2

45
tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py
View file

@ -23,3 +23,48 @@ async def test_azure_chat_o_series_transformation():
)
print(response)
assert response["model"] == "web-interface-o1-mini"
def test_azure_o_series_transform_request_flattens_top_level_anyof():
"""Regression test for LIT-6510: the o-series super() chain ends in
OpenAIGPTConfig, whose flatten gate skips provider 'azure', so
AzureOpenAIO1Config must flatten tool schema combinators itself."""
tool = {
"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"],
},
},
}
optional_params = {"tools": [tool]}
request = AzureOpenAIO1Config().transform_request(
model="o3-mini",
messages=[{"role": "user", "content": "hi"}],
optional_params=optional_params,
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
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 "anyOf" in tool["function"]["parameters"]
assert optional_params["tools"][0] is tool

28
tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
View file

@ -643,3 +643,31 @@ def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_mode
assert "output_config" not in result
last_content = result["messages"][-1]["content"]
assert json.loads(last_content[-1]["text"]) == schema
@pytest.mark.parametrize(
"model",
["us.anthropic.claude-fable-5-1", "anthropic.claude-fable-5-1"],
)
def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(local_model_cost_map, model):
"""Regression: Bedrock rejects both native ``output_config.format`` and forced
tool_choice for Fable 5.1, so invoke must use the tool-based path without a
forced ``tool_choice``."""
result = AmazonAnthropicClaudeConfig().map_openai_params(
non_default_params={
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "test_schema",
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
},
}
},
optional_params={},
model=model,
drop_params=False,
)
assert "output_format" not in result
assert "tools" in result
assert "tool_choice" not in result

33
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -6497,6 +6497,39 @@ def test_unforced_tool_choice_unaffected_on_fable_5_1_converse(local_model_cost_
assert result == ({"auto": {}} if tool_choice == "auto" else None)
@pytest.mark.parametrize(
"model",
["anthropic.claude-fable-5-1", "us.anthropic.claude-fable-5-1"],
)
def test_response_format_avoids_native_and_forced_tool_choice_on_fable_5_1_converse(
local_model_cost_map, model
):
"""Regression: Bedrock rejects both ``outputConfig`` structured output and forced
tool_choice for Fable 5.1, so response_format must map to a tool without a forced
tool_choice."""
config = AmazonConverseConfig()
result = config.map_openai_params(
non_default_params={
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "test_schema",
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
},
}
},
optional_params={},
model=model,
drop_params=False,
)
assert "outputConfig" not in result
assert "tools" in result
assert "tool_choice" not in result
assert result.get("json_mode") is True
def test_forced_tool_choice_forwarded_on_converse_models_that_support_it(
local_model_cost_map, monkeypatch
):

243
tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py
View file

@ -871,6 +871,134 @@ class TestCacheControlPreservationForCustomEndpoint:
assert all("cache_control" not in m for m in body["messages"])
class TestToolChoiceWithoutToolsDropped:
def setup_method(self):
self.config = OpenAIGPTConfig()
@staticmethod
def _pi_compact_summarization_messages():
return [
{
"role": "system",
"content": "You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.",
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "<conversation>\n[User]: Reply with exactly: ok-1\n\n[Assistant]: ok-1\n</conversation>\n\nThe messages above are a conversation to summarize.",
}
],
},
]
def _transform(self, optional_params, config=None, model="gpt-5.6-sol"):
return (config or self.config).transform_request(
model=model,
messages=self._pi_compact_summarization_messages(),
optional_params=optional_params,
litellm_params={"custom_llm_provider": "openai", "api_base": None},
headers={},
)
def test_pi_compact_shape_drops_tool_choice_none_without_tools(self):
body = self._transform(
{
"stream": True,
"stream_options": {"include_usage": True},
"store": False,
"max_completion_tokens": 13107,
"tool_choice": "none",
}
)
assert "tool_choice" not in body
assert "tools" not in body
assert body["model"] == "gpt-5.6-sol"
assert body["stream"] is True
assert body["stream_options"] == {"include_usage": True}
assert body["store"] is False
assert body["max_completion_tokens"] == 13107
def test_drops_tool_choice_auto_without_tools(self):
body = self._transform({"tool_choice": "auto"})
assert "tool_choice" not in body
def test_drops_named_function_tool_choice_without_tools(self):
body = self._transform(
{"tool_choice": {"type": "function", "function": {"name": "get_weather"}}}
)
assert "tool_choice" not in body
def test_drops_tool_choice_but_keeps_empty_tools_array(self):
body = self._transform({"tools": [], "tool_choice": "none"})
assert "tool_choice" not in body
assert body["tools"] == []
def test_gpt5_config_drops_tool_choice_without_tools(self):
body = self._transform({"tool_choice": "none"}, config=OpenAIGPT5Config())
assert "tool_choice" not in body
@pytest.mark.parametrize(
"tool_choice",
[
"none",
"auto",
"required",
{"type": "function", "function": {"name": "get_weather"}},
],
)
def test_preserves_tool_choice_when_tools_present(self, tool_choice):
tools = [
{
"type": "function",
"function": {"name": "get_weather", "parameters": {}},
}
]
body = self._transform({"tools": tools, "tool_choice": tool_choice})
assert body["tool_choice"] == tool_choice
assert body["tools"] == tools
def test_preserves_tool_choice_with_legacy_functions(self):
functions = [{"name": "get_weather", "parameters": {}}]
body = self._transform({"functions": functions, "tool_choice": "auto"})
assert body["tool_choice"] == "auto"
assert body["functions"] == functions
def test_preserves_function_call_without_functions(self):
body = self._transform({"function_call": "none"})
assert body["function_call"] == "none"
@pytest.mark.asyncio
async def test_async_transform_drops_tool_choice_without_tools(self):
body = await self.config.async_transform_request(
model="gpt-5.6-sol",
messages=self._pi_compact_summarization_messages(),
optional_params={"stream": True, "tool_choice": "none"},
litellm_params={"custom_llm_provider": "openai", "api_base": None},
headers={},
)
assert "tool_choice" not in body
@pytest.mark.asyncio
async def test_async_transform_preserves_tool_choice_when_tools_present(self):
tools = [
{
"type": "function",
"function": {"name": "get_weather", "parameters": {}},
}
]
body = await self.config.async_transform_request(
model="gpt-5.6-sol",
messages=self._pi_compact_summarization_messages(),
optional_params={"tools": tools, "tool_choice": "auto"},
litellm_params={"custom_llm_provider": "openai", "api_base": None},
headers={},
)
assert body["tool_choice"] == "auto"
assert body["tools"] == tools
class TestToolMessageImageHoisting:
"""transform_request moves tool-message images into a following user message
(OpenAI-compatible APIs only accept text in role:"tool" messages)."""
@ -1038,3 +1166,118 @@ 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.
"""
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"}

92
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -828,6 +828,24 @@ class MockPassThroughGuardrail(CustomGuardrail):
return inputs
class MockRecordingGuardrail(MockPassThroughGuardrail):
"""Pass-through guardrail that records every apply_guardrail inputs payload"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.seen_inputs: List[GenericGuardrailAPIInputs] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.seen_inputs.append(inputs)
return inputs
class TestOpenAIResponsesHandlerStreamingOutputProcessing:
"""Test streaming output processing functionality"""
@ -1104,6 +1122,80 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
assert output_text == original_text
@pytest.mark.asyncio
async def test_failed_stream_scans_delta_text(self):
"""A stream ending in response.failed has text only in delta events; the
fallback scan must assemble and scan it instead of skipping on an empty string."""
handler = OpenAIResponsesHandler()
guardrail = MockRecordingGuardrail(guardrail_name="test")
responses_so_far = [
{"type": "response.created", "response": {"id": "resp_123"}},
{"type": "response.output_item.added", "item": {"type": "message", "id": "msg_123"}},
{
"type": "response.output_text.delta",
"item_id": "msg_123",
"output_index": 0,
"content_index": 0,
"delta": "Hello",
},
{
"type": "response.output_text.delta",
"item_id": "msg_123",
"output_index": 0,
"content_index": 0,
"delta": " world",
},
{"type": "response.failed", "response": {"id": "resp_123", "status": "failed"}},
]
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert result == responses_so_far
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["Hello world"]]
def test_get_streaming_string_so_far_prefers_done_text_over_deltas(self):
"""The done event repeats the whole part, so deltas must not be double counted;
a part with no done event yet still contributes its joined deltas."""
handler = OpenAIResponsesHandler()
events = [
{
"type": "response.output_text.delta",
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": "Hello",
},
{
"type": "response.output_text.delta",
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": " world",
},
{
"type": "response.output_text.done",
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"text": "Hello world",
},
{
"type": "response.output_text.delta",
"item_id": "msg_2",
"output_index": 1,
"content_index": 0,
"delta": "; unfinished",
},
]
assert handler.get_streaming_string_so_far(events) == "Hello world; unfinished"
class TestGetStructuredMessages:
"""Test the get_structured_messages method for Responses API handler."""

105
tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
View file

@ -220,6 +220,111 @@ class TestOpenAIResponsesAPIConfig:
assert result["input"] == input_clean
def test_transform_drops_foreign_tool_call_item_ids(self):
"""Replayed tool call items whose ids are not OpenAI-shaped (e.g.
Anthropic toolu_/srvtoolu_ ids after a router fallback) must be sent
without an id: OpenAI 400s foreign ids ("Expected an ID that begins
with 'fc'") but accepts the items with no id at all. Genuine fc_/ctc_
ids and non-tool-call items pass through untouched."""
replayed_input = [
{"role": "user", "content": [{"type": "input_text", "text": "hi"}]},
{
"type": "function_call",
"id": "toolu_01Foreign",
"call_id": "toolu_01Foreign",
"name": "get_weather",
"arguments": '{"city": "SF"}',
},
{"type": "function_call_output", "call_id": "toolu_01Foreign", "output": "sunny"},
{
"type": "custom_tool_call",
"id": "srvtoolu_01Foreign",
"call_id": "srvtoolu_01Foreign",
"name": "apply_patch",
"input": "patch",
},
{
"type": "function_call",
"id": "fc_genuine",
"call_id": "call_genuine",
"name": "get_weather",
"arguments": "{}",
},
{"type": "message", "id": "msg_1", "role": "assistant", "content": []},
]
result = self.config.transform_responses_api_request(
model=self.model,
input=replayed_input,
response_api_optional_request_params={},
litellm_params={},
headers={},
)
assert "id" not in result["input"][1]
assert result["input"][1]["call_id"] == "toolu_01Foreign"
assert "id" not in result["input"][3]
assert result["input"][3]["call_id"] == "srvtoolu_01Foreign"
assert result["input"][4]["id"] == "fc_genuine"
assert result["input"][5]["id"] == "msg_1"
assert replayed_input[1]["id"] == "toolu_01Foreign"
assert replayed_input[3]["id"] == "srvtoolu_01Foreign"
def test_transform_keeps_foreign_tool_call_item_ids_for_other_providers(self):
"""Providers reusing this config that do not enforce OpenAI's id
shapes must keep replayed ids untouched."""
from litellm.types.utils import LlmProviders
class _OpenRouterLikeConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.OPENROUTER
replayed_input = [
{
"type": "function_call",
"id": "toolu_01Foreign",
"call_id": "toolu_01Foreign",
"name": "get_weather",
"arguments": "{}",
}
]
result = _OpenRouterLikeConfig().transform_responses_api_request(
model="openrouter/some-model",
input=replayed_input,
response_api_optional_request_params={},
litellm_params={},
headers={},
)
assert result["input"][0]["id"] == "toolu_01Foreign"
def test_transform_compact_drops_foreign_tool_call_item_ids(self):
"""The compact request path replays input the same way, so it must
apply the same id drop."""
replayed_input = [
{
"type": "function_call",
"id": "toolu_01Foreign",
"call_id": "toolu_01Foreign",
"name": "get_weather",
"arguments": "{}",
}
]
_url, data = self.config.transform_compact_response_api_request(
model=self.model,
input=replayed_input,
response_api_optional_request_params={},
api_base="https://api.openai.com/v1/responses",
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert "id" not in data["input"][0]
assert data["input"][0]["call_id"] == "toolu_01Foreign"
def test_transform_streaming_response(self):
"""Test streaming response transformation"""
# Test with a text delta event

25
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py
View file

@ -727,3 +727,28 @@ def test_sanitize_strips_effort_for_haiku_45():
data = {"output_config": {"effort": "high"}}
sanitize_vertex_anthropic_output_params(data, "vertex_ai/claude-opus-4-6")
assert data["output_config"] == {"effort": "high"}
def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_model_cost_map):
"""Regression: Fable 5.1 rejects forced tool use, so the vertex map entry
advertises native structured output and ``response_format`` must map to
``output_format`` instead of the tool-based path's forced tool_choice."""
config = VertexAIAnthropicConfig()
response_format = {
"type": "json_schema",
"json_schema": {
"name": "test_schema",
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
},
}
result_params = config.map_openai_params(
non_default_params={"response_format": response_format},
optional_params={},
model="claude-fable-5-1",
drop_params=False,
)
assert "output_format" in result_params
assert "tool_choice" not in result_params
assert "tools" not in result_params

294
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -214,6 +214,46 @@ class TestExecuteWithMcpClient:
assert server.scopes == ["read", "write"]
assert server.has_client_credentials is True
async def test_preview_forwards_per_server_timeout_to_client_factory(self, monkeypatch):
"""The request's per-server timeout must reach the temporary MCPServer model:
the client factory reads ``server.timeout`` for both the per-request timeout
and the preview's whole-walk listing deadline."""
captured: dict = {}
def fake_build_stdio_env(server, raw_headers):
return None
async def fake_create_client(*args, **kwargs):
captured["server"] = kwargs.get("server")
return object()
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"_build_stdio_env",
fake_build_stdio_env,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"_create_mcp_client",
fake_create_client,
raising=False,
)
async def ok_operation(client):
return {"status": "ok"}
payload = NewMCPServerRequest(
server_name="slow-catalog-server",
url="https://example.com",
timeout=120.5,
)
result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation)
assert result["status"] == "ok"
assert captured["server"].timeout == 120.5
@pytest.mark.asyncio
async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch):
"""For M2M OAuth servers the incoming Authorization header (which carries
@ -524,6 +564,131 @@ class TestTestToolsList:
assert captured["oauth2_headers"] is None
assert oauth_call_counter["count"] == 0
async def test_preview_tools_list_times_out_on_slow_pagination(self, monkeypatch):
"""A preview whose upstream paginates past the listing deadline returns a
timeout error instead of holding the request open."""
monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False)
monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False)
class SlowClient:
async def list_tools(self, raise_on_error=False):
await asyncio.sleep(1)
return []
async def fake_execute(
request,
operation,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=None,
):
return await operation(SlowClient())
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
from litellm.proxy._types import LitellmUserRoles
request = _build_request()
payload = NewMCPServerRequest(
server_name="example",
url="https://example.com",
auth_type=MCPAuth.api_key,
credentials={"auth_value": "secret-key"},
)
result = await rest_endpoints.test_tools_list(
request,
payload,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result["status"] == "error"
assert result["error"] is True
assert "Timed out listing tools" in result["message"]
async def test_preview_tools_list_succeeds_within_deadline(self, monkeypatch):
"""The preview timeout scope passes a fast listing through untouched."""
from mcp.types import Tool as MCPTool
class QuickClient:
async def list_tools(self, raise_on_error=False):
return [MCPTool(name="quick_tool", description="q", inputSchema={})]
async def fake_execute(
request,
operation,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=None,
):
return await operation(QuickClient())
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
from litellm.proxy._types import LitellmUserRoles
request = _build_request()
payload = NewMCPServerRequest(
server_name="example",
url="https://example.com",
auth_type=MCPAuth.api_key,
credentials={"auth_value": "secret-key"},
)
result = await rest_endpoints.test_tools_list(
request,
payload,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result["error"] is None
assert result["message"] == "Successfully retrieved tools"
assert [tool["name"] for tool in result["tools"]] == ["quick_tool"]
async def test_preview_tools_list_honors_per_server_timeout(self, monkeypatch):
"""A per-server timeout above the global default extends the preview deadline."""
monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False)
monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False)
from mcp.types import Tool as MCPTool
class SlowConfiguredClient:
timeout = 1.0
async def list_tools(self, raise_on_error=False):
await asyncio.sleep(0.2)
return [MCPTool(name="slow_tool", description="s", inputSchema={})]
async def fake_execute(
request,
operation,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=None,
):
return await operation(SlowConfiguredClient())
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
from litellm.proxy._types import LitellmUserRoles
request = _build_request()
payload = NewMCPServerRequest(
server_name="example",
url="https://example.com",
auth_type=MCPAuth.api_key,
credentials={"auth_value": "secret-key"},
)
result = await rest_endpoints.test_tools_list(
request,
payload,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result["error"] is None
assert [tool["name"] for tool in result["tools"]] == ["slow_tool"]
async def test_extracts_oauth2_headers(self, monkeypatch):
"""Ensure oauth2 auth type pulls oauth headers and omits MCP auth header."""
@ -786,9 +951,7 @@ class TestListToolsRestAPI:
they do for a gateway session, never to the bare session key."""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
session_auth = UserAPIKeyAuth(
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
)
session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user")
admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org")
async def fake_reload(user_id):
@ -868,9 +1031,7 @@ class TestListToolsRestAPI:
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
session_auth = UserAPIKeyAuth(
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
)
session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user")
scoped_auth = UserAPIKeyAuth(
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="toolset-scope",
@ -952,6 +1113,123 @@ class TestListToolsRestAPI:
assert scope_inputs == [session_auth]
assert reload_calls == []
async def test_single_server_response_includes_paginated_upstream_tools(
self,
monkeypatch,
):
"""The REST tools/list path should include tools beyond the upstream first page."""
import litellm.experimental_mcp_client.client as mcp_client_module
from mcp.types import ListToolsResult, PaginatedRequestParams
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.server import MCPServer
from litellm.types.mcp import MCPTransport
async def fake_contexts(user_api_key_auth):
return [user_api_key_auth]
async def fake_get_allowed_mcp_servers(*args, **kwargs):
return ["server-1"]
stub_server = MCPServer(
server_id="server-1",
name="stub",
server_name="stub",
alias="stub",
url="https://example.com/mcp",
transport=MCPTransport.http,
mcp_info={"server_name": "stub"},
)
stub_server.available_on_public_internet = True
mock_transport_ctx = AsyncMock()
mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
mock_transport_ctx.__aexit__ = AsyncMock(return_value=None)
monkeypatch.setattr(
mcp_client_module,
"streamable_http_client",
MagicMock(return_value=mock_transport_ctx),
raising=False,
)
mock_session_ctx = AsyncMock()
mock_session_instance = AsyncMock()
mock_session_instance.initialize = AsyncMock(return_value=None)
mock_session_instance.list_tools.side_effect = [
ListToolsResult(
tools=[
MCPTool(
name="first_page_tool",
description="First page tool",
inputSchema={},
)
],
nextCursor="page-2",
),
ListToolsResult(
tools=[
MCPTool(
name="second_page_tool",
description="Second page tool",
inputSchema={},
)
]
),
]
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
monkeypatch.setattr(
mcp_client_module,
"ClientSession",
MagicMock(return_value=mock_session_ctx),
raising=False,
)
monkeypatch.setattr(
rest_endpoints,
"build_effective_auth_contexts",
fake_contexts,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_allowed_mcp_servers",
fake_get_allowed_mcp_servers,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"filter_server_ids_by_ip_with_info",
lambda server_ids, client_ip: (server_ids, 0),
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: stub_server if server_id == "server-1" else None,
raising=False,
)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id="server-1",
user_api_key_dict=UserAPIKeyAuth(),
)
assert set(result.keys()) == {"tools", "error", "message"}
assert [tool.name for tool in result["tools"]] == [
"first_page_tool",
"second_page_tool",
]
assert result["error"] is None
assert result["message"] == "Successfully retrieved tools"
assert mock_session_instance.list_tools.call_count == 2
second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"]
assert isinstance(second_call_params, PaginatedRequestParams)
assert second_call_params.cursor == "page-2"
async def test_include_disabled_tools_is_admin_only(self, monkeypatch):
"""include_disabled_tools skips the allowlist filter only for PROXY_ADMIN;
a non-admin passing it stays filtered so the REST endpoint can't be used
@ -3021,9 +3299,7 @@ class TestRestListToolsetFiltering:
mock_manager = MagicMock()
mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {})
mock_manager.resolve_toolset_tool_permissions = AsyncMock(
return_value={"server-a": ["lookup_status"]}
)
mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value={"server-a": ["lookup_status"]})
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,

614
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_alice.py Normal file
View file

@ -0,0 +1,614 @@
import json
import os
from copy import deepcopy
from unittest.mock import AsyncMock
import httpx
import pytest
from httpx import Request, Response
import litellm
from litellm.exceptions import GuardrailRaisedException
from litellm.proxy.guardrails.guardrail_hooks.alice.alice import (
GUARDRAIL_NAME,
AliceGuardrail,
AliceGuardrailMissingSecrets,
_json_safe,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
def _guardrail(**overrides: object) -> AliceGuardrail:
params: dict[str, object] = {"api_key": "test-key", "guardrail_name": "alice", "event_hook": "pre_call"}
params.update(overrides)
return AliceGuardrail(**params)
def _verdict(payload: dict[str, object], status_code: int = 200) -> Response:
return Response(
status_code=status_code,
json=payload,
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
def test_alice_guardrail_config(monkeypatch: pytest.MonkeyPatch):
"""Should register through init_guardrails_v2 like any other provider."""
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setenv("ALICE_API_KEY", "test-key")
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "alice",
"litellm_params": {"guardrail": "alice", "mode": "pre_call", "default_on": True},
}
],
config_file_path="",
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, AliceGuardrail)]
assert len(registered) == 1
assert registered[0].guardrail_name == "alice"
class TestAliceGuardrailInitialization:
def setup_method(self):
for key in ("ALICE_API_KEY", "ALICE_API_BASE"):
os.environ.pop(key, None)
def test_missing_api_key_raises(self):
with pytest.raises(AliceGuardrailMissingSecrets, match="API key"):
AliceGuardrail(guardrail_name="alice", event_hook="pre_call")
def test_reads_credentials_from_environment(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("ALICE_API_KEY", "env-key")
monkeypatch.setenv("ALICE_API_BASE", "https://env.alice.test")
guardrail = AliceGuardrail(guardrail_name="alice", event_hook="pre_call")
assert guardrail.alice_api_key == "env-key"
assert guardrail.api_base == "https://env.alice.test/v2/evaluate/litellm"
def test_defaults_the_api_base(self):
assert _guardrail().api_base == "https://api.alice.io/v2/evaluate/litellm"
def test_trailing_slash_does_not_double_up(self):
assert _guardrail(api_base="https://api.alice.io/").api_base == ("https://api.alice.io/v2/evaluate/litellm")
class TestAliceForwarding:
"""The hook's arguments cross the wire as they were received — nothing selected, nothing
renamed — except the caller's raw credentials, which are stripped before request_data is
serialized (see TestAliceCredentialStripping)."""
@pytest.mark.asyncio
async def test_forwards_the_hook_arguments_verbatim(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
inputs = {"texts": ["hello"], "structured_messages": [{"role": "user", "content": "hello"}]}
request_data = {"model": "gpt-4o", "metadata": {"user_api_key_alias": "payments-bot"}}
# Snapshot before the call: @log_guardrail_information writes its own entry into
# request_data["metadata"] afterwards, so the original is no longer what was sent.
sent_inputs = deepcopy(inputs)
sent_request_data = deepcopy(request_data)
await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request")
body = guardrail.async_handler.post.call_args.kwargs["json"]
assert body["input_type"] == "request"
assert body["inputs"] == sent_inputs
assert body["request_data"] == sent_request_data
@pytest.mark.asyncio
async def test_sends_the_credential(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request")
assert guardrail.async_handler.post.call_args.kwargs["headers"]["af-api-key"] == "test-key"
@pytest.mark.asyncio
async def test_marks_a_completion_as_a_response(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
await guardrail.apply_guardrail(inputs={"texts": ["answer"]}, request_data={}, input_type="response")
assert guardrail.async_handler.post.call_args.kwargs["json"]["input_type"] == "response"
@pytest.mark.asyncio
async def test_nothing_selectable_reaches_no_evaluation(self):
"""No texts, images, tools, tool_calls, or structured_messages: genuinely nothing to send."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock()
result = await guardrail.apply_guardrail(inputs={"texts": []}, request_data={}, input_type="request")
assert result == {"texts": []}
guardrail.async_handler.post.assert_not_called()
@pytest.mark.asyncio
async def test_tool_calls_only_still_reaches_alice(self):
"""A batch with empty texts but populated tool_calls is still a selection decision Alice
should make, not the plugin — see the class docstring."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
inputs = {"texts": [], "tool_calls": [{"id": "call_1", "function": {"name": "get_weather"}}]}
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
guardrail.async_handler.post.assert_called_once()
assert guardrail.async_handler.post.call_args.kwargs["json"]["inputs"]["tool_calls"] == inputs["tool_calls"]
@pytest.mark.asyncio
async def test_images_only_still_reaches_alice(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
await guardrail.apply_guardrail(
inputs={"texts": [], "images": ["data:image/png;base64,abc"]}, request_data={}, input_type="request"
)
guardrail.async_handler.post.assert_called_once()
@pytest.mark.asyncio
async def test_structured_messages_only_still_reaches_alice(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
await guardrail.apply_guardrail(
inputs={"texts": [], "structured_messages": [{"role": "user", "content": []}]},
request_data={},
input_type="request",
)
guardrail.async_handler.post.assert_called_once()
@pytest.mark.asyncio
async def test_makes_exactly_one_attempt(self):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(side_effect=httpx.ConnectError("refused"))
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request")
assert guardrail.async_handler.post.call_count == 1
class TestAliceCredentialStripping:
"""request_data's raw-credential keys never leave the process."""
@pytest.mark.asyncio
async def test_secret_fields_and_api_key_are_stripped(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
request_data = {
"model": "gpt-4o",
"api_key": "sk-forwarded-provider-secret",
"secret_fields": {"raw_headers": {"authorization": "Bearer caller-virtual-key"}},
"metadata": {"user_api_key_alias": "payments-bot"},
}
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request")
sent_request_data = guardrail.async_handler.post.call_args.kwargs["json"]["request_data"]
assert "secret_fields" not in sent_request_data
assert "api_key" not in sent_request_data
assert sent_request_data == {"model": "gpt-4o", "metadata": {"user_api_key_alias": "payments-bot"}}
@pytest.mark.asyncio
async def test_nested_credentials_are_stripped_at_every_depth(self):
"""Shaped after a real captured Claude Code payload: the caller's Authorization/x-api-key
lives under several independent nesting paths, none of which are the root."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
request_data = {
"model": "claude-3-5-sonnet",
"secret_fields": {"raw_headers": {"authorization": "Bearer caller-virtual-key"}},
"provider_specific_header": {"extra_headers": {"authorization": "sk-ant-oat01-nested-oauth"}},
"proxy_server_request": {
"url": "/v1/messages",
"headers": {"authorization": "Bearer inbound-caller-secret", "x-request-id": "req-1"},
"body": {
"model": "claude-3-5-sonnet",
"metadata": {"headers": {"authorization": "Bearer body-metadata-secret"}},
},
},
"metadata": {
"user_api_key_alias": "payments-bot",
"headers": {"authorization": "Bearer metadata-secret"},
"requester_metadata": {"headers": {"authorization": "Bearer requester-metadata-secret"}},
},
"litellm_metadata": {"headers": {"authorization": "Bearer litellm-metadata-secret"}},
}
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request")
posted_body = guardrail.async_handler.post.call_args.kwargs["json"]
serialized = json.dumps(posted_body)
assert "authorization" not in serialized.lower()
assert "caller-virtual-key" not in serialized
assert "nested-oauth" not in serialized
assert "inbound-caller-secret" not in serialized
assert "body-metadata-secret" not in serialized
assert "metadata-secret" not in serialized
assert "requester-metadata-secret" not in serialized
assert "litellm-metadata-secret" not in serialized
sent_request_data = posted_body["request_data"]
assert sent_request_data["model"] == "claude-3-5-sonnet"
assert sent_request_data["proxy_server_request"]["url"] == "/v1/messages"
assert "headers" not in sent_request_data["proxy_server_request"]
assert sent_request_data["proxy_server_request"]["body"]["model"] == "claude-3-5-sonnet"
assert "headers" not in sent_request_data["proxy_server_request"]["body"]["metadata"]
assert sent_request_data["metadata"]["user_api_key_alias"] == "payments-bot"
assert "headers" not in sent_request_data["metadata"]
assert "requester_metadata" in sent_request_data["metadata"]
assert "headers" not in sent_request_data["metadata"]["requester_metadata"]
assert "headers" not in sent_request_data["litellm_metadata"]
assert "secret_fields" not in sent_request_data
assert "provider_specific_header" not in sent_request_data
@pytest.mark.asyncio
async def test_the_original_request_data_is_not_mutated(self):
"""Stripping must only affect the outbound copy — api_key still has to reach the
provider, and secret_fields still has to reach the rest of the request pipeline."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
request_data = {"api_key": "sk-forwarded-provider-secret", "secret_fields": {"raw_headers": {}}}
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request")
assert request_data["api_key"] == "sk-forwarded-provider-secret"
assert request_data["secret_fields"] == {"raw_headers": {}}
class TestAliceVerdicts:
@pytest.mark.asyncio
async def test_allow_leaves_the_inputs_untouched(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []}))
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
async def test_block_surfaces_the_policy_message(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict(
{
"verdict": "BLOCK",
"categories": ["self_harm"],
"correlation_id": "c1",
"message": "Blocked by your organization's policy",
}
)
)
with pytest.raises(GuardrailRaisedException) as error:
await guardrail.apply_guardrail(inputs={"texts": ["bad"]}, request_data={}, input_type="request")
assert "Blocked by your organization's policy" in str(error.value)
@pytest.mark.asyncio
async def test_block_without_a_message_still_blocks(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "BLOCK", "categories": []}))
with pytest.raises(GuardrailRaisedException):
await guardrail.apply_guardrail(inputs={"texts": ["bad"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_mask_substitutes_by_position(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict(
{
"verdict": "MASK",
"categories": ["pii"],
"replacements": [{"index": 1, "text": "my ssn is ***"}],
}
)
)
result = await guardrail.apply_guardrail(
inputs={"texts": ["untouched", "my ssn is 123-45-6789"]},
request_data={},
input_type="request",
)
assert result["texts"] == ["untouched", "my ssn is ***"]
@pytest.mark.asyncio
async def test_mask_that_lands_nowhere_blocks(self):
"""A mask that wrote nothing would let the text through under a verdict that said not to."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict({"verdict": "MASK", "categories": [], "replacements": [{"index": 9, "text": "***"}]})
)
with pytest.raises(GuardrailRaisedException):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_mask_with_no_replacements_blocks(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "MASK", "categories": []}))
with pytest.raises(GuardrailRaisedException):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_mask_with_one_invalid_replacement_blocks_entirely(self):
"""A mixed valid/invalid replacement list must not let the valid half through: that
would leave the content named by the invalid entry unmasked while looking like success."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict(
{
"verdict": "MASK",
"categories": ["pii"],
"replacements": [{"index": 0, "text": "***"}, {"index": 9, "text": "***"}],
}
)
)
with pytest.raises(GuardrailRaisedException):
await guardrail.apply_guardrail(
inputs={"texts": ["my ssn is 123-45-6789"]}, request_data={}, input_type="request"
)
@pytest.mark.asyncio
async def test_mask_leaves_structured_messages_identical(self):
"""A new structured_messages object makes the translation layer skip the texts write-back."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict({"verdict": "MASK", "categories": [], "replacements": [{"index": 0, "text": "***"}]})
)
messages = [{"role": "user", "content": "secret"}]
result = await guardrail.apply_guardrail(
inputs={"texts": ["secret"], "structured_messages": messages},
request_data={},
input_type="request",
)
assert result["structured_messages"] is messages
@pytest.mark.asyncio
async def test_detect_allows_and_leaves_the_text_alone(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=_verdict({"verdict": "DETECT", "categories": ["profanity"], "correlation_id": "c1"})
)
result = await guardrail.apply_guardrail(inputs={"texts": ["mild"]}, request_data={}, input_type="request")
assert result["texts"] == ["mild"]
class TestAliceUnreachable:
@pytest.mark.parametrize(
"failure",
[
pytest.param({"side_effect": httpx.ConnectError("refused")}, id="connect-error"),
pytest.param({"return_value": _verdict({"verdict": "MAYBE"})}, id="unrecognized-verdict"),
pytest.param({"return_value": _verdict({})}, id="no-verdict"),
],
)
@pytest.mark.asyncio
async def test_fails_closed_by_default(self, failure: dict):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(**failure)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_fails_open_when_configured(self):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(side_effect=httpx.ConnectError("refused"))
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
class TestAliceTransportFailures:
"""Every path out of the HTTP call, since each decides whether traffic flows unscreened."""
@pytest.mark.asyncio
async def test_a_timeout_is_unreachable(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
side_effect=litellm.exceptions.Timeout(message="slow", model="gpt-4o", llm_provider="openai")
)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.parametrize("status", [500, 502, 503, 504])
@pytest.mark.asyncio
async def test_upstream_5xx_is_unreachable(self, status: int):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(
side_effect=httpx.HTTPStatusError(
"server error",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
response=_verdict({}, status_code=status),
)
)
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
async def test_a_500_fails_closed_by_default(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
side_effect=httpx.HTTPStatusError(
"server error",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
response=_verdict({}, status_code=500),
)
)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_a_4xx_is_not_treated_as_unreachable(self):
"""A rejected credential is our misconfiguration, not an outage — it must not fail open."""
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(
side_effect=httpx.HTTPStatusError(
"unauthorized",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
response=_verdict({}, status_code=401),
)
)
with pytest.raises(httpx.HTTPStatusError):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_a_non_object_body_fails_closed_by_default(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
json=["not", "an", "object"],
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_a_non_object_body_fails_open_when_configured(self):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
json=["not", "an", "object"],
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
async def test_malformed_json_fails_closed_by_default(self):
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
content=b"not json",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_malformed_json_fails_open_when_configured(self):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
content=b"not json",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
async def test_an_undecodable_body_fails_closed_by_default(self):
"""UnicodeDecodeError is a sibling of JSONDecodeError under ValueError, not a subclass."""
guardrail = _guardrail()
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
content=b"\xff\xfe not utf-8",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
with pytest.raises(GuardrailRaisedException, match="unavailable"):
await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
@pytest.mark.asyncio
async def test_an_undecodable_body_fails_open_when_configured(self):
guardrail = _guardrail(unreachable_fallback="fail_open")
guardrail.async_handler.post = AsyncMock(
return_value=Response(
status_code=200,
content=b"\xff\xfe not utf-8",
request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"),
)
)
result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
class TestAliceSerialization:
"""`request_data` carries live objects, so it cannot be posted as it stands."""
def test_drops_what_cannot_serialize_and_keeps_the_rest(self):
class Span:
pass
result = _json_safe({"model": "x", "metadata": {"span": Span(), "user": "u1"}, "n": 1})
assert result == {"model": "x", "metadata": {"span": None, "user": "u1"}, "n": 1}
def test_survives_a_cycle(self):
data: dict = {"a": 1}
data["self"] = data
assert _json_safe(data) == {"a": 1, "self": None}
def test_drops_a_model_that_will_not_dump(self):
class Stubborn:
def model_dump(self, mode: str = "python") -> dict:
raise RuntimeError("cannot serialise")
assert _json_safe({"m": Stubborn()}) == {"m": None}
def test_drops_a_bare_unserialisable_value(self):
class Span:
pass
assert _json_safe(Span()) is None
def test_dumps_pydantic_models(self):
from pydantic import BaseModel
class Model(BaseModel):
name: str
assert _json_safe({"m": Model(name="x")}) == {"m": {"name": "x"}}
def test_config_model_is_exposed_for_the_ui():
config_model = AliceGuardrail.get_config_model()
assert config_model is not None
assert config_model.ui_friendly_name() == "Alice"
def test_guardrail_name_constant():
assert GUARDRAIL_NAME == "alice"

150
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -5595,6 +5595,156 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail"
def _responses_stream_events() -> list:
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
deltas = [
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_lit6457",
output_index=0,
content_index=0,
delta=part,
)
for part in ("Hello", " world")
]
completed = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_lit6457",
created_at=1234567890,
model="gpt-4o",
object="response",
status="completed",
output=[
{
"type": "message",
"id": "msg_lit6457",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "Hello world"}],
}
],
),
)
return [*deltas, completed]
@pytest.mark.asyncio
async def test_responses_api_stream_scans_output_and_replays_buffered_events():
"""Streamed /v1/responses events must be scanned via the unified translation
layer, not fed to stream_chunk_builder (which raises APIError on them)."""
guardrail = BedrockGuardrail(
guardrail_name="bedrock-responses-stream",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
stream_events = _responses_stream_events()
order = []
yielded = []
async def record_scan(*args, **kwargs):
order.append("scan")
return {"action": "NONE", "assessments": [], "outputs": []}
async def mock_stream():
for event in stream_events:
yield event
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
response=mock_stream(),
request_data={"model": "gpt-4o", "input": "hi"},
):
order.append("chunk")
yielded.append(chunk)
assert order == ["scan", "chunk", "chunk", "chunk"]
assert len(yielded) == len(stream_events)
assert all(emitted is original for emitted, original in zip(yielded, stream_events))
def _responses_failed_stream_events() -> list:
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseFailedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
deltas = [
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_lit6457_failed",
output_index=0,
content_index=0,
delta=part,
)
for part in ("Hello", " world")
]
failed = ResponseFailedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_FAILED,
response=ResponsesAPIResponse(
id="resp_lit6457_failed",
created_at=1234567890,
model="gpt-4o",
object="response",
status="failed",
output=[],
),
)
return [*deltas, failed]
@pytest.mark.asyncio
async def test_responses_api_failed_stream_scans_delta_text_before_replay():
"""A responses stream that dies mid-generation carries its text only in delta
events; the end-of-stream scan must still see that text instead of skipping
on an empty assembled string and replaying the buffer unmoderated."""
guardrail = BedrockGuardrail(
guardrail_name="bedrock-responses-failed-stream",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
stream_events = _responses_failed_stream_events()
order = []
scan_payloads = []
yielded = []
async def record_scan(*args, **kwargs):
order.append("scan")
scan_payloads.append(str(args) + str(kwargs))
return {"action": "NONE", "assessments": [], "outputs": []}
async def mock_stream():
for event in stream_events:
yield event
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
response=mock_stream(),
request_data={"model": "gpt-4o", "input": "hi"},
):
order.append("chunk")
yielded.append(chunk)
assert order == ["scan", "chunk", "chunk", "chunk"]
assert "Hello world" in scan_payloads[0]
assert len(yielded) == len(stream_events)
assert all(emitted is original for emitted, original in zip(yielded, stream_events))
@pytest.mark.asyncio
async def test_apply_guardrail_debug_log_masks_signed_request_headers():
import logging

126
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -18,6 +18,7 @@ from litellm.proxy._types import (
ReconcileOutcome,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
_get_team_deployments,
@ -263,6 +264,131 @@ class TestModelManagementAuthChecks:
)
assert "403" in str(exc_info.value)
def test_can_user_attach_credential_admin_success(self):
result = ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
user_api_key_dict=self.admin_user,
)
assert result is True
def test_can_user_attach_credential_without_credential_allows_any_role(self):
result = ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model"),
user_api_key_dict=self.team_admin_user,
)
assert result is True
def test_can_user_attach_credential_team_admin_fails(self):
with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info:
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
user_api_key_dict=self.team_admin_user,
)
assert exc_info.value.code == "403"
def test_can_user_attach_credential_unchanged_existing_allows_any_role(self):
result = ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
user_api_key_dict=self.team_admin_user,
existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
)
assert result is True
def test_can_user_attach_credential_unchanged_encrypted_existing_allows_any_role(self, monkeypatch):
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
encrypted_name = encrypt_value_helper(value="shared-credential")
assert encrypted_name != "shared-credential"
result = ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
user_api_key_dict=self.team_admin_user,
existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name=encrypted_name),
)
assert result is True
@pytest.mark.asyncio
async def test_add_new_model_rejects_credential_attach_for_non_admin(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
add_new_model,
)
mock_prisma = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # test-quality-ok: prior auth check needs a live DB; only the credential check is under test
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
):
with pytest.raises(ProxyException) as exc_info:
await add_new_model(
model_params=Deployment(
model_name="credential-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o", litellm_credential_name="shared-credential"
),
model_info={"id": "credential-create-test"},
),
user_api_key_dict=self.team_admin_user,
)
assert exc_info.value.code == "403"
mock_prisma.db.litellm_proxymodeltable.create.assert_not_called()
@pytest.mark.asyncio
async def test_patch_model_rejects_credential_attach_for_non_admin(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
from litellm.types.router import updateLiteLLMParams
model_id = "credential-patch-test"
db_model = Deployment(
model_name="credential-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info={"id": model_id},
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # test-quality-ok: stubs the DB row fetch; only the credential check is under test
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
new=AsyncMock(return_value=db_model),
),
patch( # test-quality-ok: prior auth check needs a live DB; only the credential check is under test
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
patch( # test-quality-ok: asserts the DB write is never reached on rejection
"litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db",
new=AsyncMock(),
) as mock_update,
):
with pytest.raises(ProxyException) as exc_info:
await patch_model(
model_id=model_id,
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(
model="openai/gpt-4o", litellm_credential_name="shared-credential"
)
),
user_api_key_dict=self.team_admin_user,
)
assert exc_info.value.code == "403"
mock_update.assert_not_awaited()
def test_can_user_attach_credential_internal_user_fails(self):
with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info:
ModelManagementAuthChecks.can_user_attach_credential(
litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"),
user_api_key_dict=self.normal_user,
)
assert exc_info.value.code == "403"
class MockModelTable:
def __init__(self, model_aliases: Dict[str, str], include: Optional[dict] = None):

96
tests/test_litellm/proxy/proxy_server/test_routes_utils.py
View file

@ -9,12 +9,13 @@ Pins (PR2):
from __future__ import annotations
import asyncio
from unittest.mock import AsyncMock, MagicMock
import json
import pytest
import litellm
from litellm.proxy import proxy_server
from litellm.router_utils import pattern_match_deployments
from .conftest import normalize # type: ignore[import-not-found]
@ -99,6 +100,7 @@ def test_token_counter_missing_input_returns_400(
@pytest.fixture
def patched_supported_params(monkeypatch):
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(
litellm,
"get_llm_provider",
@ -124,12 +126,104 @@ def test_supported_openai_params_happy_path(client, auth_as, patched_supported_p
}
def test_supported_openai_params_resolves_router_alias(client, auth_as, monkeypatch):
"""A router alias absent from the cost map resolves through the deployment's underlying model."""
router = litellm.Router(
model_list=[
{
"model_name": "claude-opus-4-6-cached",
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
with auth_as():
response = client.get("/utils/supported_openai_params", params={"model": "claude-opus-4-6-cached"})
assert response.status_code == 200
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
assert response.json() == {"supported_openai_params": expected}
assert "max_tokens" in response.json()["supported_openai_params"]
def test_supported_openai_params_declared_prefix_alias_resolves_through_router(client, auth_as, monkeypatch):
"""Regression: an alias whose name starts with an authenticating provider's prefix skipped
router resolution and answered with that provider's params instead of the deployment's."""
router = litellm.Router(
model_list=[
{
"model_name": "github_copilot/gpt-4o",
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
with auth_as():
response = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
assert response.status_code == 200
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
assert response.json() == {"supported_openai_params": expected}
def test_supported_openai_params_never_runs_oauth_for_authenticating_providers(client, auth_as, monkeypatch, tmp_path):
"""Regression: github_copilot/chatgpt names answer from their declaration; resolving them
through ``get_llm_provider`` would run the provider's OAuth device flow and block the event loop."""
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
(tmp_path / "access-token").write_text("fake-access-token")
(tmp_path / "api-key.json").write_text(
json.dumps(
{
"token": "fake-api-key",
"expires_at": 4102444800,
"endpoints": {"api": "https://api.githubcopilot.com"},
}
)
)
router = litellm.Router(
model_list=[
{
"model_name": "copilot-alias",
"litellm_params": {"model": "github_copilot/gpt-4o"},
},
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*"},
},
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
resolution_attempts: list[str] = []
def _oauth_tripwire(model, *args, **kwargs):
resolution_attempts.append(model)
raise AssertionError("get_llm_provider would run the OAuth device flow")
monkeypatch.setattr(litellm, "get_llm_provider", _oauth_tripwire)
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire)
expected = litellm.get_supported_openai_params(model="gpt-4o", custom_llm_provider="github_copilot")
with auth_as():
via_alias = client.get("/utils/supported_openai_params", params={"model": "copilot-alias"})
via_direct_name = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
assert via_alias.status_code == 200
assert via_alias.json() == {"supported_openai_params": expected}
assert via_direct_name.status_code == 200
assert via_direct_name.json() == {"supported_openai_params": expected}
assert resolution_attempts == []
def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch):
"""Pins ``GET /utils/supported_openai_params`` (error: unknown model)."""
def _raise(model):
raise Exception("unknown")
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "get_llm_provider", _raise)
with auth_as():
response = client.get("/utils/supported_openai_params", params={"model": "??"})

129
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -1737,8 +1737,7 @@ class TestRunServerDbSetup:
mock_atexit_register,
mock_subprocess_run,
):
"""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."""
"""Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter"""
from litellm.proxy.proxy_cli import run_server
# Mock subprocess.run to simulate prisma being available
@ -1788,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=True
use_migrate=True, use_v2_resolver=False
)
# Reset mocks
@ -1803,38 +1802,9 @@ class TestRunServerDbSetup:
standalone_mode=False,
)
mock_setup_database.assert_called_with(
use_migrate=False, use_v2_resolver=True
use_migrate=False, use_v2_resolver=False
)
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")
@ -1899,7 +1869,7 @@ class TestRunServerDbSetup:
)
assert exc_info.value.code == 1
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=True
use_migrate=True, use_v2_resolver=False
)
@patch("subprocess.run")
@ -2011,6 +1981,7 @@ class TestRunServerDbSetup:
use_migrate=True, use_v2_resolver=True
)
# --- Module-level helpers for worker startup hook tests ---
_dummy_hook_called = False
@ -2481,6 +2452,96 @@ class TestReadReplicaConnectionParams:
assert "DATABASE_URL_READ_REPLICA" not in captured
class TestMaxIdleConnectionLifetimeDefault:
"""The proxy defaults `max_idle_connection_lifetime` below common infra idle
timeouts so stale pooled connections are recycled instead of failing requests."""
def _config(self, tmp_path, general_settings):
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings}))
return str(config_path)
def test_default_applied_to_database_and_direct_url(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
direct_url="postgresql://t:t@localhost:5432/t",
)
for env_var in ("DATABASE_URL", "DIRECT_URL"):
query = urlparse.parse_qs(urlparse.urlparse(captured[env_var]).query)
assert query["max_idle_connection_lifetime"] == ["60"], env_var
def test_url_pinned_value_wins_over_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["300"]
def test_url_pinned_value_wins_over_config_key(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["300"]
def test_config_key_overrides_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["45"]
def test_extra_connection_params_override_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(
tmp_path,
{"database_extra_connection_params": {"max_idle_connection_lifetime": 120}},
),
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["120"]
def test_read_replica_gets_the_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
read_replica_url="postgresql://t:t@reader:5432/t",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["60"]
def test_replica_pinned_value_wins(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
read_replica_url="postgresql://t:t@reader:5432/t?max_idle_connection_lifetime=200",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["200"]
def test_config_key_reaches_the_read_replica(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
read_replica_url="postgresql://t:t@reader:5432/t",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["45"]
def test_idle_lifetime_params_prefers_configured_value(self):
from litellm.proxy.db.db_url_settings import idle_lifetime_params
assert dict(idle_lifetime_params(45)) == {"max_idle_connection_lifetime": 45}
assert dict(idle_lifetime_params(None)) == {"max_idle_connection_lifetime": 60}
class TestTokenAuthCliFlags:
"""`--azure_postgresql_auth` has to reach the URL assembly the same way the env var does."""

67
tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py
View file

@ -648,6 +648,72 @@ class TestLiteLLMCompletionResponsesConfig:
assert responses_api_response.status == "incomplete"
def test_tool_call_only_response_emits_no_null_text_message_item(self):
"""A tool-calls-only turn (message content None, e.g. from Anthropic)
must not emit a message output item whose output_text has text null.
OpenAI rejects such an item on replay with
"Invalid type for 'input[..].content[..].text': expected a string, but
got null instead." Native OpenAI tool-only turns carry no message item."""
chat_completion_response = ModelResponse(
id="test-response-id",
created=1234567890,
model="claude-sonnet-4-5",
object="chat.completion",
choices=[
Choices(
finish_reason="tool_calls",
index=0,
message=Message(
content=None,
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
id="toolu_01OnlyToolCall",
type="function",
function=Function(name="get_weather", arguments='{"city": "SF"}'),
)
],
),
)
],
)
responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
request_input="what's the weather in SF?",
responses_api_request={},
chat_completion_response=chat_completion_response,
)
output_types = [item.type for item in responses_api_response.output]
assert "message" not in output_types
assert "function_call" in output_types
def test_content_bearing_response_still_emits_message_item(self):
"""Turns with real text content must keep their message output item."""
chat_completion_response = ModelResponse(
id="test-response-id",
created=1234567890,
model="claude-sonnet-4-5",
object="chat.completion",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(content="It is sunny.", role="assistant"),
)
],
)
responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
request_input="what's the weather in SF?",
responses_api_request={},
chat_completion_response=chat_completion_response,
)
message_items = [item for item in responses_api_response.output if item.type == "message"]
assert len(message_items) == 1
assert message_items[0].content[0].text == "It is sunny."
def test_transform_chat_completion_response_preserves_hidden_params(self):
"""Test that _hidden_params from chat completion response are preserved in responses API response"""
# Setup
@ -3343,6 +3409,7 @@ class TestEnsureOutputItemContentPartAdded:
iterator._pending_tool_events = []
iterator._tool_output_index_by_call_id = {}
iterator._tool_args_by_call_id = {}
iterator._tool_item_id_by_call_id = {}
iterator._tool_call_id_by_index = {}
iterator._ambiguous_tool_call_indexes = set()
iterator._next_tool_output_index = 1

76
tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -132,7 +132,7 @@ def test_tool_call_delta_is_emitted_as_responses_events():
evt2 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk)
assert evt2 is not None
assert evt2.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA
assert evt2.item_id == "call_1"
assert evt2.item_id == "fc_call_1"
assert evt2.output_index == 1
# The delta will be a chunk of the arguments, not the full arguments
assert len(evt2.delta) <= 10 # Chunks are max 10 characters
@ -197,7 +197,7 @@ def test_tool_calls_present_only_in_final_response_are_emitted_before_completed(
# The last event should be FUNCTION_CALL_ARGUMENTS_DONE
assert evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE
assert evt.item_id == "call_2"
assert evt.item_id == "fc_call_2"
assert evt.output_index == 1
assert evt.arguments == '{"y":2}'
@ -291,7 +291,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior():
# Verify each delta is at most 10 characters
for evt in delta_events:
assert len(evt.delta) <= 10
assert evt.item_id == "call_test"
assert evt.item_id == "fc_call_test"
assert evt.output_index == 1
assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__
@ -349,7 +349,8 @@ def test_tool_call_delta_without_id_uses_index_mapping():
if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
]
assert len(output_item_added_events) == 1
assert output_item_added_events[0].item.id == "call_abc123"
assert output_item_added_events[0].item.id == "fc_call_abc123"
assert output_item_added_events[0].item.call_id == "call_abc123"
def test_parallel_tool_calls_without_ids_use_index_mapping():
@ -404,8 +405,8 @@ def test_parallel_tool_calls_without_ids_use_index_mapping():
arguments_by_call_id.setdefault(evt.item_id, "")
arguments_by_call_id[evt.item_id] += evt.delta
assert arguments_by_call_id["call_a"] == '{"x":1}'
assert arguments_by_call_id["call_b"] == '{"y":2}'
assert arguments_by_call_id["fc_call_a"] == '{"x":1}'
assert arguments_by_call_id["fc_call_b"] == '{"y":2}'
def test_reused_index_with_new_call_id_marks_fallback_ambiguous():
@ -461,10 +462,10 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous():
arguments_by_call_id.setdefault(evt.item_id, "")
arguments_by_call_id[evt.item_id] += evt.delta
assert arguments_by_call_id["call_a"] == '{"a":'
assert arguments_by_call_id["call_b"] == '{"b":'
assert arguments_by_call_id["call_a"] != '{"a":1}'
assert arguments_by_call_id["call_b"] != '{"b":1}'
assert arguments_by_call_id["fc_call_a"] == '{"a":'
assert arguments_by_call_id["fc_call_b"] == '{"b":'
assert arguments_by_call_id["fc_call_a"] != '{"a":1}'
assert arguments_by_call_id["fc_call_b"] != '{"b":1}'
@pytest.mark.asyncio
@ -557,3 +558,58 @@ def test_object_tool_call_arguments_stream_as_valid_json():
)
assert json.loads(streamed_arguments) == {"command": "ls", "flags": ["-l"]}
def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id():
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=AsyncMock(),
request_input="Test input",
responses_api_request={},
)
response = ModelResponse(
id="resp-anthropic",
created=123,
model="test-model",
object="chat.completion",
choices=[
{
"index": 0,
"finish_reason": "tool_calls",
"message": {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "toolu_01AbCdEf",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
"index": 0,
}
],
},
}
],
)
iterator.litellm_model_response = response
events = []
while True:
evt = iterator.common_done_event_logic(sync_mode=True)
events.append(evt)
if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE:
break
added = [e for e in events if e.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
deltas = [e for e in events if e.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA]
dones = [e for e in events if e.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE]
item_dones = [e for e in events if e.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
assert len(added) == 1 and len(dones) == 1 and len(item_dones) == 1 and deltas
assert added[0].item.id == "fc_toolu_01AbCdEf"
assert added[0].item.call_id == "toolu_01AbCdEf"
assert item_dones[0].item.id == "fc_toolu_01AbCdEf"
assert item_dones[0].item.call_id == "toolu_01AbCdEf"
for evt in deltas + dones:
assert evt.item_id == added[0].item.id

82
tests/test_litellm/responses/test_custom_tool_call.py
View file

@ -20,6 +20,7 @@ from litellm.responses.litellm_completion_transformation.transformation import (
from litellm.responses.litellm_completion_transformation.custom_tools import (
extract_custom_tool_names,
is_custom_tool_call,
openai_shaped_tool_call_item_id,
unwrap_custom_tool_arguments,
build_tool_call_item_kwargs,
convert_custom_tool_to_function_tool,
@ -129,6 +130,41 @@ class TestCustomToolUtilities:
assert kwargs["arguments"] == raw
assert "input" not in kwargs
def test_openai_shaped_tool_call_item_id_prefixes_foreign_ids(self):
"""Anthropic-style tool ids must be normalized to OpenAI's item id
shapes (fc/ctc prefixes) so replaying the item to OpenAI does not 400
with "Expected an ID that begins with 'fc'"."""
assert openai_shaped_tool_call_item_id("function_call", "toolu_01Abc") == "fc_toolu_01Abc"
assert openai_shaped_tool_call_item_id("function_call", "srvtoolu_01Xyz") == "fc_srvtoolu_01Xyz"
assert openai_shaped_tool_call_item_id("custom_tool_call", "toolu_01Abc") == "ctc_toolu_01Abc"
assert openai_shaped_tool_call_item_id("function_call", "fc_already") == "fc_already"
assert openai_shaped_tool_call_item_id("custom_tool_call", "ctc_already") == "ctc_already"
assert openai_shaped_tool_call_item_id("function_call", "") == ""
assert openai_shaped_tool_call_item_id("message", "toolu_01Abc") == "toolu_01Abc"
def test_build_tool_call_item_kwargs_normalizes_item_id_keeps_call_id(self):
"""The streaming item id gets the OpenAI shape while call_id stays raw
so tool_result pairing (which keys off call_id) keeps working."""
function_kwargs = build_tool_call_item_kwargs(
call_id="toolu_01Abc",
name="get_weather",
arguments_or_input="{}",
status="completed",
custom_tool_names=set(),
)
assert function_kwargs["id"] == "fc_toolu_01Abc"
assert function_kwargs["call_id"] == "toolu_01Abc"
custom_kwargs = build_tool_call_item_kwargs(
call_id="toolu_01Def",
name="apply_patch",
arguments_or_input=json.dumps({"content": "patch"}),
status="completed",
custom_tool_names={"apply_patch"},
)
assert custom_kwargs["id"] == "ctc_toolu_01Def"
assert custom_kwargs["call_id"] == "toolu_01Def"
def test_unwrap_custom_tool_arguments_oversized_returns_raw(self):
"""Arguments larger than the safety cap are returned unchanged to avoid
OOM on JSON parsing a pathologically large string."""
@ -293,6 +329,52 @@ class TestTransformationCustomTools:
assert item.name == "regular_tool"
assert item.arguments == json.dumps({"param": "value"})
def test_transform_anthropic_tool_call_ids_get_openai_item_id_shape(self):
"""Anthropic tool ids (toolu_/srvtoolu_) surfacing through the bridge
must be emitted with fc/ctc-prefixed item ids so a Responses client can
replay them to OpenAI verbatim, while call_id stays raw for pairing."""
from litellm.types.utils import ModelResponse, Choices, Message, ChatCompletionMessageToolCall, Function
client_call = ChatCompletionMessageToolCall(
id="toolu_01ClientCall",
type="function",
function=Function(name="get_weather", arguments=json.dumps({"city": "SF"})),
)
server_call = ChatCompletionMessageToolCall(
id="srvtoolu_01ServerCall",
type="function",
function=Function(name="web_search", arguments=json.dumps({"query": "zig"})),
)
custom_call = ChatCompletionMessageToolCall(
id="toolu_01CustomCall",
type="function",
function=Function(name="apply_patch", arguments=json.dumps({"content": "patch content"})),
)
message = Message(role="assistant", content=None, tool_calls=[client_call, server_call, custom_call])
choices = [Choices(index=0, message=message, finish_reason="tool_calls")]
response = ModelResponse(
id="test_response", choices=choices, created=1234567890, model="claude-sonnet-4-5", object="chat.completion"
)
responses_api_request = {
"tools": [{"type": "custom", "name": "apply_patch"}, {"type": "function", "name": "get_weather"}]
}
result = LiteLLMCompletionResponsesConfig.transform_chat_completion_tools_to_responses_tools(
response, responses_api_request=responses_api_request
)
assert [item.id for item in result] == [
"fc_toolu_01ClientCall",
"fc_srvtoolu_01ServerCall",
"ctc_toolu_01CustomCall",
]
assert [item.call_id for item in result] == [
"toolu_01ClientCall",
"srvtoolu_01ServerCall",
"toolu_01CustomCall",
]
def test_transform_mixed_tool_calls(self):
"""Test transformation with both custom and regular tool calls."""
from litellm.types.utils import ModelResponse, Choices, Message, ChatCompletionMessageToolCall, Function

78
tests/test_litellm/router_utils/test_pattern_match_deployments.py Normal file
View file

@ -0,0 +1,78 @@
"""Behavior pins for ``litellm/router_utils/pattern_match_deployments.py``."""
from __future__ import annotations
from litellm.router_utils import pattern_match_deployments
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
def _wildcard_deployment(model_name: str) -> dict:
return {"model_name": model_name, "litellm_params": {"model": model_name}}
def _matched_models(matches: list[dict] | None) -> list[str]:
return [deployment["litellm_params"]["model"] for deployment in matches or []]
def test_get_pattern_never_resolves_declared_authenticating_providers(monkeypatch):
"""Regression: resolving a github_copilot/chatgpt name through ``get_llm_provider`` runs the
provider's OAuth device flow; the auth layer walks every wildcard router on every request, so
a single metadata lookup for an unserved name would block the proxy's event loop."""
resolution_attempts: list[str] = []
def _oauth_tripwire(model, *args, **kwargs):
resolution_attempts.append(model)
raise AssertionError("get_llm_provider would run the OAuth device flow")
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire)
unmatched_router = PatternMatchRouter()
unmatched_router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*"))
assert unmatched_router.get_pattern("github_copilot/gpt-4o") is None
matched_router = PatternMatchRouter()
matched_router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*"))
assert _matched_models(matched_router.get_pattern("github_copilot/gpt-4o")) == ["github_copilot/gpt-4o"]
assert _matched_models(matched_router.get_pattern("gpt-4o", custom_llm_provider="github_copilot")) == [
"github_copilot/gpt-4o"
]
assert resolution_attempts == []
def test_get_pattern_bare_provider_name_never_matches_that_providers_wildcard(monkeypatch):
"""Regression: a bare ``github_copilot`` adopted itself as its provider and retried as
``github_copilot/github_copilot``, false-matching the wildcard for a name no deployment serves."""
def _unknown_provider(model, *args, **kwargs):
raise ValueError(f"unknown provider for {model}")
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider)
router = PatternMatchRouter()
router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*"))
assert router.get_pattern("github_copilot") is None
def test_get_pattern_missing_model_returns_none(monkeypatch):
"""Regression: a request without a model reaches the auth layer's pattern walk as ``None``; the
declared-provider guard raised ``TypeError`` where the old inline resolve swallowed every
resolver error, so the proxy's missing-model 400 became a crash."""
def _unknown_provider(model, *args, **kwargs):
raise ValueError(f"unknown provider for {model}")
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider)
router = PatternMatchRouter()
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
assert router.get_pattern(None) is None
def test_get_pattern_still_resolves_unqualified_names(monkeypatch):
monkeypatch.setattr(
pattern_match_deployments,
"get_llm_provider",
lambda model, **kwargs: (model, "openai", None, None),
)
router = PatternMatchRouter()
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"]

200
tests/test_litellm/test_register_model_custom_pricing.py
View file

@ -793,3 +793,203 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key():
finally:
litellm.model_cost.pop(model_key, None)
_invalidate_model_cost_lowercase_map()
def test_update_dictionary_merges_nested_dicts_without_aliasing():
"""A nested dict must be merged copy-on-write: the pre-existing nested dict
object stays untouched, and the caller's incoming nested dict is never
inserted by reference into the merged result.
"""
from litellm.utils import _update_dictionary
existing_nested = {"hours_utc": "01:00-02:00"}
existing = {"off_peak_pricing": existing_nested}
incoming_nested = {"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}]}
incoming = {"off_peak_pricing": incoming_nested}
merged = _update_dictionary(existing, incoming)
assert merged["off_peak_pricing"] == {
"hours_utc": "01:00-02:00",
"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}],
}
assert existing_nested == {"hours_utc": "01:00-02:00"}
assert merged["off_peak_pricing"] is not incoming_nested
fresh = _update_dictionary({}, incoming)
assert fresh["off_peak_pricing"] == incoming_nested
assert fresh["off_peak_pricing"] is not incoming_nested
def test_router_deployments_sharing_backend_keep_their_own_off_peak_pricing():
"""Two deployments of the same backend model with different
``off_peak_pricing`` blocks must each keep their own schedule under their
unique model id, and neither block may leak onto the shared backend keys.
Before the fix, ``register_model`` inserted the first deployment's block by
reference into the built-in ``gpt-4o-mini`` entry, and the second
deployment's registration merged its keys into that same object, corrupting
the first deployment's schedule and polluting the built-in entry.
"""
from litellm import Router
active_block = {
"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}],
"input_cost_per_token": 5e-07,
"output_cost_per_token": 1e-06,
}
inactive_block = {
"hours_utc": "05:00-06:00",
"input_cost_per_token": 5e-07,
"output_cost_per_token": 1e-06,
}
shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"]
deployment_ids = ["offpeak-alias-dep-1", "offpeak-alias-dep-2"]
original_entries = _snapshot_model_cost_entries(shared_keys)
router = Router(
model_list=[
{
"model_name": "offpeak-active-weekday",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {
"id": deployment_ids[0],
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
"off_peak_pricing": dict(active_block),
},
},
{
"model_name": "offpeak-inactive-hours",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {
"id": deployment_ids[1],
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
"off_peak_pricing": dict(inactive_block),
},
},
]
)
try:
registered_first = litellm.model_cost[deployment_ids[0]]["off_peak_pricing"]
registered_second = litellm.model_cost[deployment_ids[1]]["off_peak_pricing"]
assert registered_first == active_block
assert registered_second == inactive_block
for shared_key in shared_keys:
shared_entry = litellm.model_cost.get(shared_key) or {}
assert not shared_entry.get("off_peak_pricing")
finally:
for deployment_id in deployment_ids:
litellm.model_cost.pop(deployment_id, None)
_restore_model_cost_entries(original_entries)
del router
def test_router_off_peak_only_deployment_inherits_builtin_base_rates():
"""A deployment that sets only ``off_peak_pricing`` on its model_info must
still be costed from its deployment-scoped entry: the base token rates are
inherited from the backend model's built-in cost map entry, since the
shared backend key deliberately never carries the off-peak block.
"""
from litellm import Router
block = {
"hours_utc": "00:00-00:00",
"input_cost_per_token": 5e-05,
"output_cost_per_token": 1e-04,
}
shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"]
deployment_id = "offpeak-only-dep-1"
original_entries = _snapshot_model_cost_entries(shared_keys + [deployment_id])
builtin_info = litellm.get_model_info(model="openai/gpt-4o-mini")
router = Router(
model_list=[
{
"model_name": "offpeak-only",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {"id": deployment_id, "off_peak_pricing": dict(block)},
}
]
)
try:
entry = litellm.model_cost[deployment_id]
assert entry["off_peak_pricing"] == block
assert entry["input_cost_per_token"] is not None
assert entry["input_cost_per_token"] == builtin_info["input_cost_per_token"]
assert entry["output_cost_per_token"] == builtin_info["output_cost_per_token"]
for shared_key in shared_keys:
shared_entry = litellm.model_cost.get(shared_key) or {}
assert not shared_entry.get("off_peak_pricing")
finally:
_restore_model_cost_entries(original_entries)
del router
def test_use_custom_pricing_for_model_sees_off_peak_only_model_info():
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
block = {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-05}
assert use_custom_pricing_for_model({"metadata": {"model_info": {"off_peak_pricing": block}}}) is True
assert use_custom_pricing_for_model({"metadata": {"model_info": {"off_peak_pricing": None}}}) is False
assert use_custom_pricing_for_model({"metadata": {"model_info": {"id": "some-id"}}}) is False
def test_completion_cost_applies_off_peak_only_deployment_pricing():
"""End to end through the cost calculator: with ``custom_pricing`` set and
a ``router_model_id`` whose entry carries only an always-on off-peak block,
the request bills at the block's rates rather than the shared backend rate.
"""
from litellm import Router
from litellm.types.utils import ModelResponse, Usage
block = {
"hours_utc": "00:00-00:00",
"input_cost_per_token": 5e-05,
"output_cost_per_token": 1e-04,
}
shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"]
deployment_id = "offpeak-only-dep-2"
original_entries = _snapshot_model_cost_entries(shared_keys + [deployment_id])
router = Router(
model_list=[
{
"model_name": "offpeak-only",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {"id": deployment_id, "off_peak_pricing": dict(block)},
}
]
)
try:
response = ModelResponse(
model="gpt-4o-mini",
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
)
cost = litellm.completion_cost(
completion_response=response,
model="openai/gpt-4o-mini",
custom_llm_provider="openai",
custom_pricing=True,
router_model_id=deployment_id,
)
assert cost == pytest.approx(100 * 5e-05 + 50 * 1e-04)
finally:
_restore_model_cost_entries(original_entries)
del router

151
tests/test_litellm/test_router_model_cost_isolation.py
View file

@ -538,6 +538,157 @@ def test_inherit_builtin_cache_pricing_noop_for_unknown_backend():
assert model_info == {"input_cost_per_token": 0.000003}
def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates():
"""Direct unit test of the helper: an entry carrying only an
off_peak_pricing block inherits the backend model's built-in base token
rates, so cost lookup via the deployment id can bill standard rates
outside the windows.
"""
backend_model = "gpt-4o-mini"
builtin_info = litellm.get_model_info(model=backend_model, custom_llm_provider="openai")
off_peak_block = {
"hours_utc": "00:00-00:00",
"input_cost_per_token": 5e-07,
"output_cost_per_token": 1e-06,
}
model_info = {"off_peak_pricing": off_peak_block}
Router._inherit_builtin_base_rates_for_off_peak(
model_info=model_info,
backend_model=backend_model,
custom_llm_provider="openai",
)
assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"]
assert model_info["output_cost_per_token"] == builtin_info["output_cost_per_token"]
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_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
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.
"""
model_info = {
"off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07},
"input_cost_per_token": 3e-06,
}
Router._inherit_builtin_base_rates_for_off_peak(
model_info=model_info,
backend_model="gpt-4o-mini",
custom_llm_provider="openai",
)
assert model_info["input_cost_per_token"] == 3e-06
assert "output_cost_per_token" not in model_info
def test_inherit_builtin_base_rates_for_off_peak_noop_without_block_or_backend():
"""Nothing happens without an off_peak_pricing block, and an unmapped
backend model leaves the entry unchanged rather than raising.
"""
plain_info = {"id": "dep-1"}
Router._inherit_builtin_base_rates_for_off_peak(
model_info=plain_info,
backend_model="gpt-4o-mini",
custom_llm_provider="openai",
)
assert plain_info == {"id": "dep-1"}
off_peak_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=off_peak_info,
backend_model="this-backend-model-does-not-exist-x9y8z7",
custom_llm_provider=None,
)
assert "input_cost_per_token" not in off_peak_info
def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields():
"""The shared-backend-key stripping in Router relies on
CustomPricingLiteLLMParams enumerating every per-deployment pricing field.

4
ui/litellm-dashboard/public/assets/logos/alice.svg Normal file
View file

@ -0,0 +1,4 @@
<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
<rect width="24" height="24" rx="8" fill="#E686B4"/>
<path d="M12.0001 4C16.165 4.00002 19 7.1202 19 10.7146V19.1637C19 19.6256 18.6256 20 18.1637 20H16.9259C16.4641 20 16.0896 19.6256 16.0896 19.1637V11.3386C16.0896 9.31676 14.9356 6.74578 12.0001 6.74575C10.4696 6.74575 7.91038 7.76917 7.91038 11.3386V14.0344H13.801C14.2629 14.0344 14.6373 14.4088 14.6373 14.8707V15.9188C14.6373 16.3807 14.2629 16.7551 13.801 16.7551H7.91038V19.1637C7.91038 19.6256 7.53595 20 7.07406 20H5.83632C5.37443 20 5 19.6256 5 19.1637V10.7146C5 7.12018 7.83522 4 12.0001 4Z" fill="white"/>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 674 B

36
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx
View file

@ -74,6 +74,42 @@ describe("AgentsTable", () => {
expect(onDeleteClick).toHaveBeenCalledWith("agent-9", "Doomed Agent");
});
it("filters agents by name or by agent card description", async () => {
const user = userEvent.setup();
render(
<AgentsTable
agents={[
makeAgent({ agent_id: "a1", agent_name: "Billing Router" }),
makeAgent({
agent_id: "a2",
agent_name: "Second Agent",
agent_card_params: { description: "handles support tickets" },
}),
]}
{...baseProps}
/>,
);
const search = screen.getByPlaceholderText("Search agent names or descriptions...");
await user.type(search, "billing");
expect(screen.getByText("Billing Router")).toBeInTheDocument();
expect(screen.queryByText("Second Agent")).not.toBeInTheDocument();
await user.clear(search);
await user.type(search, "support tickets");
expect(screen.getByText("Second Agent")).toBeInTheDocument();
expect(screen.queryByText("Billing Router")).not.toBeInTheDocument();
});
it("shows the no-match empty state when the search matches nothing", async () => {
const user = userEvent.setup();
render(<AgentsTable agents={[makeAgent()]} {...baseProps} />);
await user.type(screen.getByPlaceholderText("Search agent names or descriptions..."), "zzzz");
expect(screen.queryByText("Test Agent")).not.toBeInTheDocument();
expect(screen.getByText("No matching agents")).toBeInTheDocument();
});
it("hides the actions column entirely for non-admins", () => {
const agent = makeAgent({ agent_id: "agent-2" });
render(<AgentsTable agents={[agent]} {...baseProps} isAdmin={false} />);

42
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx
View file

@ -1,13 +1,15 @@
"use client";
import { SortingState } from "@tanstack/react-table";
import { Bot, CircleCheck } from "lucide-react";
import { Bot, CircleCheck, Search as SearchIcon, X } from "lucide-react";
import React, { useMemo, useState } from "react";
import { Agent } from "@/components/agents/types";
import { DataTable } from "@/components/shared/DataTable";
import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group";
import { Switch } from "@/components/ui/switch";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { filterBySearchTerm } from "@/utils/searchUtils";
import { getAgentsTableColumns } from "./AgentsTableColumns";
@ -24,14 +26,18 @@ interface AgentsTableProps {
const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }];
function EmptyState() {
function EmptyState({ isFiltered }: { isFiltered: boolean }) {
return (
<div className="flex flex-col items-center gap-1 py-6">
<div className="mb-1 flex size-10 items-center justify-center rounded-lg bg-muted">
<Bot className="size-5 text-muted-foreground" />
</div>
<div className="text-sm font-medium text-foreground">No agents yet</div>
<div className="text-sm text-muted-foreground">Add an agent to make it available in your organization.</div>
<div className="text-sm font-medium text-foreground">{isFiltered ? "No matching agents" : "No agents yet"}</div>
<div className="text-sm text-muted-foreground">
{isFiltered
? "Adjust the search to see more agents."
: "Add an agent to make it available in your organization."}
</div>
</div>
);
}
@ -47,6 +53,11 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
onDeleteClick,
}) => {
const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
const [searchTerm, setSearchTerm] = useState("");
const filteredAgents = useMemo(
() => filterBySearchTerm(agents, searchTerm, (agent) => [agent.agent_name, agent.agent_card_params?.description]),
[agents, searchTerm],
);
const columns = useMemo(
() => getAgentsTableColumns({ isAdmin, onAgentClick, onDeleteClick }),
@ -55,7 +66,7 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
return (
<DataTable
data={agents}
data={filteredAgents}
columns={columns}
getRowId={(agent, index) => agent.agent_id || String(index)}
sortingMode="client"
@ -63,10 +74,27 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
onSortingChange={setSorting}
isLoading={isLoading}
loadingMessage="Loading agents…"
noDataMessage={<EmptyState />}
noDataMessage={<EmptyState isFiltered={agents.length > 0} />}
size="compact"
toolbar={() => (
<div className="flex items-center justify-end">
<div className="flex flex-wrap items-center justify-between gap-2">
<InputGroup className="max-w-sm">
<InputGroupAddon>
<SearchIcon className="size-4 text-muted-foreground" />
</InputGroupAddon>
<InputGroupInput
placeholder="Search agent names or descriptions..."
value={searchTerm}
onChange={(e) => setSearchTerm(e.target.value)}
/>
{searchTerm && (
<InputGroupAddon align="inline-end">
<InputGroupButton size="icon-xs" aria-label="Clear search" onClick={() => setSearchTerm("")}>
<X />
</InputGroupButton>
</InputGroupAddon>
)}
</InputGroup>
<TooltipProvider delay={300}>
<Tooltip>
<TooltipTrigger

4
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx
View file

@ -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/ }));

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